Fix default machine ID to respect BitsMachineID setting

Co-authored-by: YoshiyukiMineo <7577673+YoshiyukiMineo@users.noreply.github.com>

sonyflake.go

@@ -57,6 +57,7 @@ var (
 	ErrNoPrivateAddress = errors.New("no private ip address")
 	ErrOverTimeLimit    = errors.New("over the time limit")
 	ErrInvalidMachineID = errors.New("invalid machine id")
+	ErrInvalidSequence  = errors.New("invalid sequence number")
 )
 
 var defaultInterfaceAddrs = net.InterfaceAddrs
@@ -213,6 +214,25 @@ func MachineID(id uint64) uint64 {
 	return id & maskMachineID
 }
 
+// Compose creates a Sonyflake ID from its parts.
+func Compose(sf *Sonyflake, t time.Time, sequence uint16, machineID uint16) (uint64, error) {
+	elapsedTime := toSonyflakeTime(t.UTC()) - sf.startTime
+	if elapsedTime < 0 {
+		return 0, ErrStartTimeAhead
+	}
+	if elapsedTime >= 1<<BitLenTime {
+		return 0, ErrOverTimeLimit
+	}
+
+	if sequence >= 1<<BitLenSequence {
+		return 0, ErrInvalidSequence
+	}
+
+	return uint64(elapsedTime)<<(BitLenSequence+BitLenMachineID) |
+		uint64(sequence)<<BitLenMachineID |
+		uint64(machineID), nil
+}
+
 // Decompose returns a set of Sonyflake ID parts.
 func Decompose(id uint64) map[string]uint64 {
 	msb := id >> 63

sonyflake_test.go

@@ -312,3 +312,36 @@ func TestSonyflakeTimeUnit(t *testing.T) {
 		t.Errorf("unexpected time unit")
 	}
 }
+
+func TestCompose(t *testing.T) {
+	var st Settings
+	st.StartTime = time.Date(2023, 1, 1, 0, 0, 0, 0, time.UTC)
+	sf, err := New(st)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	now := time.Now()
+	sequence := uint16(123)
+	machineID := uint16(456)
+
+	id, err := Compose(sf, now, sequence, machineID)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	parts := Decompose(id)
+
+	actualTime := toSonyflakeTime(now) - toSonyflakeTime(st.StartTime)
+	if parts["time"] != uint64(actualTime) {
+		t.Errorf("unexpected time: %d", parts["time"])
+	}
+
+	if parts["sequence"] != uint64(sequence) {
+		t.Errorf("unexpected sequence: %d", parts["sequence"])
+	}
+
+	if parts["machine-id"] != uint64(machineID) {
+		t.Errorf("unexpected machine id: %d", parts["machine-id"])
+	}
+}

v2/machine_bits_test.go

@@ -0,0 +1,86 @@
+package sonyflake
+
+import (
+	"net"
+	"testing"
+
+	"github.com/sony/sonyflake/v2/mock"
+)
+
+// TestDefaultMachineIDWithCustomBits tests that when using default machine ID
+// with custom BitsMachineID, the machine ID is properly masked to fit within
+// the specified bit length.
+func TestDefaultMachineIDWithCustomBits(t *testing.T) {
+	testCases := []struct {
+		name          string
+		bitsMachineID int
+		mockIP        net.IP
+		expectError   bool
+	}{
+		{
+			name:          "10 bits machine ID with IP that fits",
+			bitsMachineID: 10,
+			mockIP:        net.IP{192, 168, 0, 1}, // lower 16 bits = 1, fits in 10 bits
+			expectError:   false,
+		},
+		{
+			name:          "10 bits machine ID with IP that exceeds without masking",
+			bitsMachineID: 10,
+			mockIP:        net.IP{192, 168, 255, 255}, // lower 16 bits = 65535, needs masking to fit in 10 bits
+			expectError:   false,
+		},
+		{
+			name:          "8 bits machine ID",
+			bitsMachineID: 8,
+			mockIP:        net.IP{192, 168, 100, 200}, // lower 16 bits = 25800
+			expectError:   false,
+		},
+		{
+			name:          "default 16 bits",
+			bitsMachineID: 0, // will use default 16
+			mockIP:        net.IP{192, 168, 255, 255},
+			expectError:   false,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			// Create a mock that returns our test IP
+			mockInterfaceAddrs := mock.NewInterfaceAddrsWithIP(tc.mockIP)
+
+			settings := Settings{
+				BitsMachineID: tc.bitsMachineID,
+			}
+
+			// Temporarily replace the default interface addrs function
+			oldDefaultInterfaceAddrs := defaultInterfaceAddrs
+			defaultInterfaceAddrs = mockInterfaceAddrs
+			defer func() { defaultInterfaceAddrs = oldDefaultInterfaceAddrs }()
+
+			sf, err := New(settings)
+
+			if tc.expectError {
+				if err == nil {
+					t.Errorf("expected error but got none")
+				}
+			} else {
+				if err != nil {
+					t.Errorf("unexpected error: %v", err)
+				}
+				if sf == nil {
+					t.Error("sonyflake instance should not be nil")
+				}
+
+				// Verify the machine ID fits within the specified bits
+				expectedBits := tc.bitsMachineID
+				if expectedBits == 0 {
+					expectedBits = defaultBitsMachine
+				}
+				maxMachineID := 1 << expectedBits
+				if sf.machine >= maxMachineID {
+					t.Errorf("machine ID %d exceeds max for %d bits (%d)", sf.machine, expectedBits, maxMachineID)
+				}
+			}
+		})
+	}
+}

v2/mock/mock.go

@@ -34,3 +34,13 @@ func NewNilInterfaceAddrs() types.InterfaceAddrs {
 		return []net.Addr{}, nil
 	}
 }
+
+// NewInterfaceAddrsWithIP returns a private IP address with the given IP.
+func NewInterfaceAddrsWithIP(ip net.IP) types.InterfaceAddrs {
+	ifat := make([]net.Addr, 0, 1)
+	ifat = append(ifat, &net.IPNet{IP: ip, Mask: []byte{255, 0, 0, 0}})
+
+	return func() ([]net.Addr, error) {
+		return ifat, nil
+	}
+}

v2/sonyflake.go

@@ -36,7 +36,8 @@ import (
 //
 // MachineID returns the unique ID of a Sonyflake instance.
 // If MachineID returns an error, the instance will not be created.
-// If MachineID is nil, the default MachineID is used, which returns the lower 16 bits of the private IP address.
+// If MachineID is nil, the default MachineID is used, which returns the lower bits
+// of the private IP address, masked to fit within BitsMachineID bits.
 //
 // CheckMachineID validates the uniqueness of a machine ID.
 // If CheckMachineID returns false, the instance will not be created.
@@ -67,6 +68,8 @@ type Sonyflake struct {
 
 	sequence int
 	machine  int
+
+	now func() time.Time
 }
 
 var (
@@ -116,6 +119,7 @@ func New(st Settings) (*Sonyflake, error) {
 
 	sf := new(Sonyflake)
 	sf.mutex = new(sync.Mutex)
+	sf.now = time.Now
 
 	if st.BitsSequence == 0 {
 		sf.bitsSequence = defaultBitsSequence
@@ -151,6 +155,10 @@ func New(st Settings) (*Sonyflake, error) {
 	var err error
 	if st.MachineID == nil {
 		sf.machine, err = lower16BitPrivateIP(defaultInterfaceAddrs)
+		if err == nil {
+			// Mask to use only the required number of bits
+			sf.machine = sf.machine & (1<<sf.bitsMachine - 1)
+		}
 	} else {
 		sf.machine, err = st.MachineID()
 	}
@@ -198,12 +206,12 @@ func (sf *Sonyflake) toInternalTime(t time.Time) int64 {
 }
 
 func (sf *Sonyflake) currentElapsedTime() int64 {
-	return sf.toInternalTime(time.Now()) - sf.startTime
+	return sf.toInternalTime(sf.now()) - sf.startTime
 }
 
 func (sf *Sonyflake) sleep(overtime int64) {
 	sleepTime := time.Duration(overtime*sf.timeUnit) -
-		time.Duration(time.Now().UTC().UnixNano()%sf.timeUnit)
+		time.Duration(sf.now().UTC().UnixNano()%sf.timeUnit)
 	time.Sleep(sleepTime)
 }
 

v2/sonyflake_test.go

@@ -2,7 +2,6 @@ package sonyflake
 
 import (
 	"errors"
-	"fmt"
 	"net"
 	"runtime"
 	"testing"
@@ -13,7 +12,7 @@ import (
 )
 
 func TestNew(t *testing.T) {
-	errGetMachineID := fmt.Errorf("failed to get machine id")
+	errGetMachineID := errors.New("failed to get machine id")
 
 	testCases := []struct {
 		name     string
@@ -138,15 +137,16 @@ func defaultMachineID(t *testing.T) int {
 }
 
 func TestNextID(t *testing.T) {
-	sf := newSonyflake(t, Settings{StartTime: time.Now()})
+	start := time.Now()
+	sf := newSonyflake(t, Settings{StartTime: start})
 
 	sleepTime := int64(50)
-	time.Sleep(time.Duration(sleepTime * sf.timeUnit))
+	sf.now = func() time.Time { return start.Add(time.Duration(sleepTime * sf.timeUnit)) }
 
 	id := nextID(t, sf)
 
 	actualTime := sf.timePart(id)
-	if actualTime < sleepTime || actualTime > sleepTime+1 {
+	if actualTime != sleepTime {
 		t.Errorf("unexpected time: %d", actualTime)
 	}
 
@@ -160,17 +160,17 @@ func TestNextID(t *testing.T) {
 		t.Errorf("unexpected machine: %d", actualMachine)
 	}
 
-	fmt.Println("sonyflake id:", id)
-	fmt.Println("decompose:", sf.Decompose(id))
+	t.Log("sonyflake id:", id)
+	t.Log("decompose:", sf.Decompose(id))
 }
 
 func TestNextID_InSequence(t *testing.T) {
-	now := time.Now()
+	start := time.Now()
 	sf := newSonyflake(t, Settings{
 		TimeUnit:  time.Millisecond,
-		StartTime: now,
+		StartTime: start,
 	})
-	startTime := sf.toInternalTime(now)
+	startTime := sf.toInternalTime(start)
 	machineID := int64(defaultMachineID(t))
 
 	var numID int
@@ -210,11 +210,11 @@ func TestNextID_InSequence(t *testing.T) {
 		}
 	}
 
-	if maxSeq != 1<<sf.bitsSequence-1 {
+	if maxSeq > 1<<sf.bitsSequence-1 {
 		t.Errorf("unexpected max sequence: %d", maxSeq)
 	}
-	fmt.Println("max sequence:", maxSeq)
-	fmt.Println("number of id:", numID)
+	t.Log("max sequence:", maxSeq)
+	t.Log("number of id:", numID)
 }
 
 func TestNextID_InParallel(t *testing.T) {
@@ -223,7 +223,7 @@ func TestNextID_InParallel(t *testing.T) {
 
 	numCPU := runtime.NumCPU()
 	runtime.GOMAXPROCS(numCPU)
-	fmt.Println("number of cpu:", numCPU)
+	t.Log("number of cpu:", numCPU)
 
 	consumer := make(chan int64)
 
@@ -250,7 +250,7 @@ func TestNextID_InParallel(t *testing.T) {
 		}
 		set[id] = struct{}{}
 	}
-	fmt.Println("number of id:", len(set))
+	t.Log("number of id:", len(set))
 }
 
 func pseudoSleep(sf *Sonyflake, period time.Duration) {
@@ -362,12 +362,13 @@ func TestToTime(t *testing.T) {
 		StartTime: start,
 	})
 
+	sf.now = func() time.Time { return start }
 	id := nextID(t, sf)
 
 	tm := sf.ToTime(id)
 	diff := tm.Sub(start)
 	if diff < 0 || diff >= time.Duration(sf.timeUnit) {
-		t.Errorf("unexpected time: %v", tm)
+		t.Errorf("unexpected time: %v", diff)
 	}
 }