specificy setting a machine ID on the generator

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()

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())
+	}
+}