package snowflake import ( "errors" "math/rand" "sync" "time" platform "github.com/influxdata/influxdb" "github.com/influxdata/influxdb/pkg/snowflake" ) 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.RUnlock() 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 holds the ID generator. type IDGenerator struct { Generator *snowflake.Generator } // 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 { var id platform.ID for !id.Valid() { id = platform.ID(g.Generator.Next()) } return id }