From f5be30148bbc327629e69a43a60e3ddc7a14cc14 Mon Sep 17 00:00:00 2001
From: ccfuncy <ccfuncy@163.com>
Date: Sat, 6 May 2023 15:47:06 +0800
Subject: [PATCH] fix clock backward bug

---
 README.md         |  6 ++++++
 sonyflake.go      | 41 ++++++++++++++++++++++++++++++++++++-----
 sonyflake_test.go | 32 ++++++++++++++++++++++++++++++++
 3 files changed, 74 insertions(+), 5 deletions(-)

diff --git a/README.md b/README.md
index 9c9ebeb..39ab7ed 100644
--- a/README.md
+++ b/README.md
@@ -71,6 +71,12 @@ func (sf *Sonyflake) NextID() (uint64, error)
 NextID can continue to generate IDs for about 174 years from StartTime.
 But after the Sonyflake time is over the limit, NextID returns an error.
 
+if you want to use the mono time, you can call the method NextIDMono.
+
+```go
+func (sf *Sonyflake) NextIDMono() (uint64, error)
+```
+NextIDMono can avoid the clock backwards.
 > **Note:**
 > Sonyflake currently does not use the most significant bit of IDs,
 > so you can convert Sonyflake IDs from `uint64` to `int64` safely.
diff --git a/sonyflake.go b/sonyflake.go
index 1c1402e..d4394fe 100644
--- a/sonyflake.go
+++ b/sonyflake.go
@@ -45,11 +45,12 @@ type Settings struct {
 
 // Sonyflake is a distributed unique ID generator.
 type Sonyflake struct {
-	mutex       *sync.Mutex
-	startTime   int64
-	elapsedTime int64
-	sequence    uint16
-	machineID   uint16
+	mutex         *sync.Mutex
+	startTime     int64
+	startTimeMono time.Time
+	elapsedTime   int64
+	sequence      uint16
+	machineID     uint16
 }
 
 var defaultInterfaceAddrs = net.InterfaceAddrs
@@ -69,8 +70,10 @@ func NewSonyflake(st Settings) *Sonyflake {
 	}
 	if st.StartTime.IsZero() {
 		sf.startTime = toSonyflakeTime(time.Date(2014, 9, 1, 0, 0, 0, 0, time.UTC))
+		sf.startTimeMono = time.Now()
 	} else {
 		sf.startTime = toSonyflakeTime(st.StartTime)
+		sf.startTimeMono = st.StartTime
 	}
 
 	var err error
@@ -110,6 +113,31 @@ func (sf *Sonyflake) NextID() (uint64, error) {
 	return sf.toID()
 }
 
+// NextIDMono generates a next unique ID.
+// After the Sonyflake time overflows, NextID returns an error.
+// This function use mono time to avoid clock backwards
+func (sf *Sonyflake) NextIDMono() (uint64, error) {
+	const maskSequence = uint16(1<<BitLenSequence - 1)
+
+	sf.mutex.Lock()
+	defer sf.mutex.Unlock()
+
+	current := currentElapsedTimeMono(sf.startTimeMono)
+	if sf.elapsedTime < current {
+		sf.elapsedTime = current
+		sf.sequence = 0
+	} else {
+		sf.sequence = (sf.sequence + 1) & maskSequence
+		if sf.sequence == 0 {
+			sf.elapsedTime++
+			overtime := sf.elapsedTime - current
+			time.Sleep(sleepTime((overtime)))
+		}
+	}
+
+	return sf.toID()
+}
+
 const sonyflakeTimeUnit = 1e7 // nsec, i.e. 10 msec
 
 func toSonyflakeTime(t time.Time) int64 {
@@ -119,6 +147,9 @@ func toSonyflakeTime(t time.Time) int64 {
 func currentElapsedTime(startTime int64) int64 {
 	return toSonyflakeTime(time.Now()) - startTime
 }
+func currentElapsedTimeMono(startTimeMono time.Time) int64 {
+	return time.Now().Sub(startTimeMono).Nanoseconds() / sonyflakeTimeUnit
+}
 
 func sleepTime(overtime int64) time.Duration {
 	return time.Duration(overtime*sonyflakeTimeUnit) -
diff --git a/sonyflake_test.go b/sonyflake_test.go
index db58778..7ba4dc6 100644
--- a/sonyflake_test.go
+++ b/sonyflake_test.go
@@ -282,3 +282,35 @@ func TestSonyflakeTimeUnit(t *testing.T) {
 		t.Errorf("unexpected time unit")
 	}
 }
+
+func nextIDMono(t *testing.T) uint64 {
+	id, err := sf.NextIDMono()
+	if err != nil {
+		t.Fatal("id not generated")
+	}
+	return id
+}
+func TestSonyflake_NextIDMono(t *testing.T) {
+	sleepTime := time.Duration(50 * sonyflakeTimeUnit)
+	time.Sleep(sleepTime)
+
+	id := nextIDMono(t)
+
+	actualTime := ElapsedTime(id)
+	if actualTime < sleepTime || actualTime > sleepTime+sonyflakeTimeUnit {
+		t.Errorf("unexpected time: %d", actualTime)
+	}
+
+	actualSequence := SequenceNumber(id)
+	if actualSequence != 0 {
+		t.Errorf("unexpected sequence: %d", actualSequence)
+	}
+
+	actualMachineID := MachineID(id)
+	if actualMachineID != machineID {
+		t.Errorf("unexpected machine id: %d", actualMachineID)
+	}
+
+	fmt.Println("sonyflake id:", id)
+	fmt.Println("decompose:", Decompose(id))
+}