feat: Generate V6 from custom time (#172)

* Add NewV6WithTime

* Refactor generateV6

* fix NewV6WithTime doc comment

* fix: remove fmt.Println from test

---------

Co-authored-by: nicumicle <[email protected]>

time.go

@@ -45,11 +45,16 @@ func (t Time) UnixTime() (sec, nsec int64) {
 func GetTime() (Time, uint16, error) {
 	defer timeMu.Unlock()
 	timeMu.Lock()
-	return getTime()
+	return getTime(nil)
 }
 
-func getTime() (Time, uint16, error) {
-	t := timeNow()
+func getTime(customTime *time.Time) (Time, uint16, error) {
+	var t time.Time
+	if customTime == nil { // When not provided, use the current time
+		t = timeNow()
+	} else {
+		t = *customTime
+	}
 
 	// If we don't have a clock sequence already, set one.
 	if clockSeq == 0 {

time_test.go

@@ -0,0 +1,44 @@
+package uuid
+
+import (
+	"testing"
+	"time"
+)
+
+func TestGetTime(t *testing.T) {
+	now := time.Now()
+	tt := map[string]struct {
+		input        func() *time.Time
+		expectedTime int64
+	}{
+		"it should return the current time": {
+			input: func() *time.Time {
+				return nil
+			},
+			expectedTime: now.Unix(),
+		},
+		"it should return the provided time": {
+			input: func() *time.Time {
+				parsed, err := time.Parse(time.RFC3339, "2024-10-15T09:32:23Z")
+				if err != nil {
+					t.Errorf("timeParse unexpected error: %v", err)
+				}
+				return &parsed
+			},
+			expectedTime: 1728984743,
+		},
+	}
+
+	for name, tc := range tt {
+		t.Run(name, func(t *testing.T) {
+			result, _, err := getTime(tc.input())
+			if err != nil {
+				t.Errorf("getTime unexpected error: %v", err)
+			}
+			sec, _ := result.UnixTime()
+			if sec != tc.expectedTime {
+				t.Errorf("expected %v, got %v", tc.expectedTime, result)
+			}
+		})
+	}
+}

version6.go

@@ -4,7 +4,10 @@
 
 package uuid
 
-import "encoding/binary"
+import (
+	"encoding/binary"
+	"time"
+)
 
 // UUID version 6 is a field-compatible version of UUIDv1, reordered for improved DB locality.
 // It is expected that UUIDv6 will primarily be used in contexts where there are existing v1 UUIDs.
@@ -19,12 +22,32 @@ import "encoding/binary"
 // SetClockSequence then it will be set automatically. If GetTime fails to
 // return the current NewV6 returns Nil and an error.
 func NewV6() (UUID, error) {
-	var uuid UUID
 	now, seq, err := GetTime()
 	if err != nil {
-		return uuid, err
+		return Nil, err
+	}
+	return generateV6(now, seq), nil
+}
+
+// NewV6WithTime returns a Version 6 UUID based on the current NodeID, clock
+// sequence, and a specified time. It is similar to the NewV6 function, but allows
+// you to specify the time. If time is passed as nil, then the current time is used.
+//
+// There is a limit on how many UUIDs can be generated for the same time, so if you
+// are generating multiple UUIDs, it is recommended to increment the time.
+// If getTime fails to return the current NewV6WithTime returns Nil and an error.
+func NewV6WithTime(customTime *time.Time) (UUID, error) {
+	now, seq, err := getTime(customTime)
+	if err != nil {
+		return Nil, err
 	}
 
+	return generateV6(now, seq), nil
+}
+
+func generateV6(now Time, seq uint16) UUID {
+	var uuid UUID
+
 	/*
 	    0                   1                   2                   3
 	    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
@@ -56,5 +79,5 @@ func NewV6() (UUID, error) {
 	copy(uuid[10:], nodeID[:])
 	nodeMu.Unlock()
 
-	return uuid, nil
+	return uuid
 }

version6_test.go

@@ -0,0 +1,91 @@
+package uuid
+
+import (
+	"testing"
+	"time"
+)
+
+func TestNewV6WithTime(t *testing.T) {
+	testCases := map[string]string{
+		"test with current date":                      time.Now().Format(time.RFC3339),                                // now
+		"test with past date":                         time.Now().Add(-1 * time.Hour * 24 * 365).Format(time.RFC3339), // 1 year ago
+		"test with future date":                       time.Now().Add(time.Hour * 24 * 365).Format(time.RFC3339),      // 1 year from now
+		"test with different timezone":                "2021-09-01T12:00:00+04:00",
+		"test with negative timezone":                 "2021-09-01T12:00:00-12:00",
+		"test with future date in different timezone": "2124-09-23T12:43:30+09:00",
+	}
+
+	for testName, inputTime := range testCases {
+		t.Run(testName, func(t *testing.T) {
+			customTime, err := time.Parse(time.RFC3339, inputTime)
+			if err != nil {
+				t.Errorf("time.Parse returned unexpected error %v", err)
+			}
+			id, err := NewV6WithTime(&customTime)
+			if err != nil {
+				t.Errorf("NewV6WithTime returned unexpected error %v", err)
+			}
+
+			if id.Version() != 6 {
+				t.Errorf("got %d, want version 6", id.Version())
+			}
+			unixTime := time.Unix(id.Time().UnixTime())
+			// Compare the times in UTC format, since the input time might have different timezone,
+			// and the result is always in system timezone
+			if customTime.UTC().Format(time.RFC3339) != unixTime.UTC().Format(time.RFC3339) {
+				t.Errorf("got %s, want %s", unixTime.Format(time.RFC3339), customTime.Format(time.RFC3339))
+			}
+		})
+	}
+}
+
+func TestNewV6FromTimeGeneratesUniqueUUIDs(t *testing.T) {
+	now := time.Now()
+	ids := make([]string, 0)
+	runs := 26000
+
+	for i := 0; i < runs; i++ {
+		now = now.Add(time.Nanosecond) // Without this line, we can generate only 16384 UUIDs for the same timestamp
+		id, err := NewV6WithTime(&now)
+		if err != nil {
+			t.Errorf("NewV6WithTime returned unexpected error %v", err)
+		}
+		if id.Version() != 6 {
+			t.Errorf("got %d, want version 6", id.Version())
+		}
+
+		// Make sure we add only unique values
+		if !contains(t, ids, id.String()) {
+			ids = append(ids, id.String())
+		}
+	}
+
+	// Check we added all the UIDs
+	if len(ids) != runs {
+		t.Errorf("got %d UUIDs, want %d", len(ids), runs)
+	}
+}
+
+func BenchmarkNewV6WithTime(b *testing.B) {
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			now := time.Now()
+			_, err := NewV6WithTime(&now)
+			if err != nil {
+				b.Fatal(err)
+			}
+		}
+	})
+}
+
+func contains(t *testing.T, arr []string, str string) bool {
+	t.Helper()
+
+	for _, a := range arr {
+		if a == str {
+			return true
+		}
+	}
+
+	return false
+}