diff --git a/snowflake/id.go b/snowflake/id.go index 528dfd849aa..951b95c0669 100644 --- a/snowflake/id.go +++ b/snowflake/id.go @@ -2,30 +2,90 @@ package snowflake import ( "encoding/binary" + "errors" "math/rand" + "sync" "time" "github.com/influxdata/influxdb/pkg/snowflake" "github.com/influxdata/platform" ) -// TODO: rename to id.go - func init() { rand.Seed(time.Now().UnixNano()) + SetGlobalMachineID(rand.Intn(1023)) +} + +var globalmachineID struct { + id int + set bool + sync.RWMutex +} + +// ErrGlobalIDBadVal means that the global machine id value wasn't properly set. +var ErrGlobalIDBadVal = errors.New("globalID must be a number between (inclusive) 0 and 1023") + +// SetGlobalMachineID returns the global machine id. This number is limited to a number between 0 and 1023 inclusive. +func SetGlobalMachineID(id int) error { + if id > 1023 || id < 0 { + return ErrGlobalIDBadVal + } + globalmachineID.Lock() + globalmachineID.id = id + globalmachineID.set = true + globalmachineID.Unlock() + return nil +} + +// GlobalMachineID returns the global machine id. This number is limited to a number between 0 and 1023 inclusive. +func GlobalMachineID() int { + var id int + globalmachineID.RLock() + id = int(globalmachineID.id) + globalmachineID.Unlock() + return id } +// NewDefaultIDGenerator returns an *IDGenerator that uses the currently set global machine ID. +// If you change the global machine id, it will not change the id in any generators that have already been created. +func NewDefaultIDGenerator() *IDGenerator { + globalmachineID.RLock() + defer globalmachineID.RUnlock() + if globalmachineID.set { + return NewIDGenerator(WithMachineID(globalmachineID.id)) + } + return NewIDGenerator() +} + +// IDGenerator generates a new ID. type IDGenerator struct { Generator *snowflake.Generator } -func NewIDGenerator() *IDGenerator { - return &IDGenerator{ - // Maximum machine id is 1023 - Generator: snowflake.New(rand.Intn(1023)), +// IDGeneratorOp is an option for an IDGenerator. +type IDGeneratorOp func(*IDGenerator) + +// WithMachineID uses the low 12 bits of machineID to set the machine ID for the snowflake ID. +func WithMachineID(machineID int) IDGeneratorOp { + return func(g *IDGenerator) { + g.Generator = snowflake.New(machineID & 1023) + } +} + +// NewIDGenerator returns a new IDGenerator. Optionally you can use an IDGeneratorOp. +// to use a specific Generator +func NewIDGenerator(opts ...IDGeneratorOp) *IDGenerator { + gen := &IDGenerator{} + for _, f := range opts { + f(gen) + } + if gen.Generator == nil { + gen.Generator = snowflake.New(rand.Intn(1023)) } + return gen } +// ID returns the next platform.ID from an IDGenerator. func (g *IDGenerator) ID() platform.ID { id := make(platform.ID, 8) i := g.Generator.Next() diff --git a/snowflake/id_test.go b/snowflake/id_test.go index ea286292b88..76d3dd67d71 100644 --- a/snowflake/id_test.go +++ b/snowflake/id_test.go @@ -25,3 +25,26 @@ func TestToFromString(t *testing.T) { t.Errorf("id started as %x but got back %x", id, clone) } } + +func TestWithMachineID(t *testing.T) { + gen := NewIDGenerator(WithMachineID(1023)) + if gen.Generator.MachineID() != 1023 { + t.Errorf("expected machineID of %d but got %d", 1023, gen.Generator.MachineID()) + } + gen = NewIDGenerator(WithMachineID(1023)) + if gen.Generator.MachineID() != 1023 { + t.Errorf("expected machineID of %d but got %d", 1023, gen.Generator.MachineID()) + } + gen = NewIDGenerator(WithMachineID(99)) + if gen.Generator.MachineID() != 99 { + t.Errorf("expected machineID of %d but got %d", 99, gen.Generator.MachineID()) + } + gen = NewIDGenerator(WithMachineID(101376)) + if gen.Generator.MachineID() != 0 { + t.Errorf("expected machineID of %d but got %d", 0, gen.Generator.MachineID()) + } + gen = NewIDGenerator(WithMachineID(102399)) + if gen.Generator.MachineID() != 1023 { + t.Errorf("expected machineID of %d but got %d", 1023, gen.Generator.MachineID()) + } +}