From bf918976b2c0ef4b1490207cede4d2636f8ba454 Mon Sep 17 00:00:00 2001 From: Serge Zaitsev Date: Tue, 6 May 2025 21:12:11 +0200 Subject: [PATCH] Chore: Add snowflake xorm tag (#104300) * add tag handler for snowflake ids * add snowflake generator * fill snowflake id back to the bean * table driven test, mockable snowflake generator * use math/rand/v2 * snowflake without time.sleep * more explicit bitwise modulo * rename snowflake to randomid --- pkg/util/xorm/core/core.go | 5 +++ pkg/util/xorm/engine.go | 1 + pkg/util/xorm/session_insert.go | 70 +++++++++++++++++++++++++++++---- pkg/util/xorm/tag.go | 7 ++++ pkg/util/xorm/xorm.go | 2 + pkg/util/xorm/xorm_test.go | 40 +++++++++++++++++++ 6 files changed, 117 insertions(+), 8 deletions(-) diff --git a/pkg/util/xorm/core/core.go b/pkg/util/xorm/core/core.go index 2aee9820793..cc1b494fd41 100644 --- a/pkg/util/xorm/core/core.go +++ b/pkg/util/xorm/core/core.go @@ -34,6 +34,7 @@ type Column struct { Indexes map[string]int IsPrimaryKey bool IsAutoIncrement bool + IsRandomID bool IsCreated bool IsUpdated bool IsDeleted bool @@ -1580,6 +1581,7 @@ type Table struct { Indexes map[string]*Index PrimaryKeys []string AutoIncrement string + RandomID string Created map[string]bool Updated string Deleted string @@ -1695,6 +1697,9 @@ func (table *Table) AddColumn(col *Column) { if col.IsAutoIncrement { table.AutoIncrement = col.Name } + if col.IsRandomID { + table.RandomID = col.Name + } if col.IsCreated { table.Created[col.Name] = true } diff --git a/pkg/util/xorm/engine.go b/pkg/util/xorm/engine.go index 164a96ac605..155a93421c7 100644 --- a/pkg/util/xorm/engine.go +++ b/pkg/util/xorm/engine.go @@ -44,6 +44,7 @@ type Engine struct { defaultContext context.Context sequenceGenerator SequenceGenerator // If not nil, this generator is used to generate auto-increment values for inserts. + randomIDGen func() int64 } // CondDeleted returns the conditions whether a record is soft deleted. diff --git a/pkg/util/xorm/session_insert.go b/pkg/util/xorm/session_insert.go index 49d75c2d6b7..d411f3d15ef 100644 --- a/pkg/util/xorm/session_insert.go +++ b/pkg/util/xorm/session_insert.go @@ -8,9 +8,12 @@ import ( "errors" "fmt" "reflect" + "slices" "sort" "strconv" "strings" + "sync" + "time" "github.com/grafana/grafana/pkg/util/xorm/core" "xorm.io/builder" @@ -306,13 +309,7 @@ func (session *Session) innerInsert(bean any) (int64, error) { // If engine has a sequence number generator, use it to produce values for auto-increment columns. if len(table.AutoIncrement) > 0 && session.engine.sequenceGenerator != nil { - var found bool - for _, col := range colNames { - if col == table.AutoIncrement { - found = true - break - } - } + found := slices.Contains(colNames, table.AutoIncrement) if !found { seq, err := session.engine.sequenceGenerator.Next(session.ctx, table.Name, table.AutoIncrement) if err != nil { @@ -322,6 +319,26 @@ func (session *Session) innerInsert(bean any) (int64, error) { colNames = append(colNames, table.AutoIncrement) args = append(args, seq) } + } else if len(table.RandomID) > 0 { + found := slices.Contains(colNames, table.RandomID) + if !found { + id := session.engine.randomIDGen() + colNames = append(colNames, table.RandomID) + args = append(args, id) + // Set random ID back to the bean. + col := table.GetColumn(table.RandomID) + if col == nil { + return 0, fmt.Errorf("column %s not found in table %s", table.RandomID, table.Name) + } + idValue, err := col.ValueOf(bean) + if err != nil { + session.engine.logger.Error(err) + } + if idValue == nil || !idValue.IsValid() || !idValue.CanSet() { + return 0, fmt.Errorf("failed to set snowflake ID to bean: %v", err) + } + idValue.Set(int64ToIntValue(id, idValue.Type())) + } } exprs := session.statement.exprColumns @@ -587,7 +604,7 @@ func (session *Session) genInsertColumns(bean any) ([]string, []any, error) { } fieldValue := *fieldValuePtr - if col.IsAutoIncrement { + if col.IsAutoIncrement || col.IsRandomID { switch fieldValue.Type().Kind() { case reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int, reflect.Int64: if fieldValue.Int() == 0 { @@ -780,3 +797,40 @@ func (session *Session) insertMap(columns []string, args []any) (int64, error) { } return affected, nil } + +type snowflake struct { + mu sync.Mutex + nodeID int64 + sequence int64 + lastTime int64 + epoch time.Time +} + +// newSnowflake creates a new instance with a random node ID (0-1023) +// It forcefully converts epoch time (in milliseconds) to monotonic time +func newSnowflake(nodeID int64) *snowflake { + const snowflakeEpoch = 1288834974657 // 2010-11-04 01:42:54.657 UTC + epoch := time.Unix(snowflakeEpoch/1000, (snowflakeEpoch%1000)*1000000) + now := time.Now() + return &snowflake{nodeID: nodeID & 0x3ff, epoch: now.Add(epoch.Sub(now))} +} + +func (s *snowflake) Generate() int64 { + s.mu.Lock() + defer s.mu.Unlock() + currentTime := time.Since(s.epoch).Milliseconds() + if currentTime == s.lastTime { + s.sequence = (s.sequence + 1) & 0xfff + if s.sequence == 0 { + // wait for next millisecond, we are not using time.Sleep() here due to its low resolution (often >4ms) + for currentTime <= s.lastTime { + currentTime = time.Since(s.epoch).Milliseconds() + } + } + } else { + s.sequence = 0 + } + s.lastTime = currentTime + id := (currentTime << 22) | (s.nodeID << 12) | s.sequence + return id +} diff --git a/pkg/util/xorm/tag.go b/pkg/util/xorm/tag.go index bc473985e4d..0bed6548da9 100644 --- a/pkg/util/xorm/tag.go +++ b/pkg/util/xorm/tag.go @@ -38,6 +38,7 @@ var ( "NULL": NULLTagHandler, "NOT": IgnoreTagHandler, "AUTOINCR": AutoIncrTagHandler, + "RANDOMID": RandomIDTagHandler, "DEFAULT": DefaultTagHandler, "CREATED": CreatedTagHandler, "UPDATED": UpdatedTagHandler, @@ -88,6 +89,12 @@ func AutoIncrTagHandler(ctx *tagContext) error { return nil } +// RandomIDTagHandler describes snowflake id tag handler +func RandomIDTagHandler(ctx *tagContext) error { + ctx.col.IsRandomID = true + return nil +} + // DefaultTagHandler describes default tag handler func DefaultTagHandler(ctx *tagContext) error { if len(ctx.params) > 0 { diff --git a/pkg/util/xorm/xorm.go b/pkg/util/xorm/xorm.go index 359212dbe52..bf2e4bc7ea5 100644 --- a/pkg/util/xorm/xorm.go +++ b/pkg/util/xorm/xorm.go @@ -11,6 +11,7 @@ import ( "context" "database/sql" "fmt" + "math/rand/v2" "os" "reflect" "runtime" @@ -95,6 +96,7 @@ func NewEngine(driverName string, dataSourceName string) (*Engine, error) { tagHandlers: defaultTagHandlers, defaultContext: context.Background(), timestampFormat: "2006-01-02 15:04:05", + randomIDGen: newSnowflake(rand.Int64N(1024)).Generate, } switch uri.DbType { diff --git a/pkg/util/xorm/xorm_test.go b/pkg/util/xorm/xorm_test.go index 8b8d5aeb03b..f2939fbf9a8 100644 --- a/pkg/util/xorm/xorm_test.go +++ b/pkg/util/xorm/xorm_test.go @@ -48,3 +48,43 @@ type TestStruct struct { Comment string Json json.RawMessage } + +func TestRandomID(t *testing.T) { + type RandomIDRecord struct { + ID int64 `xorm:"'id' pk randomid"` + Comment string + } + + eng, err := NewEngine("sqlite3", ":memory:") + require.NoError(t, err) + require.NoError(t, eng.Sync(new(RandomIDRecord))) + + // Test sequence of different snowflake values + testCases := []struct { + name string + id int64 + comment string + }{ + {"first insert", 42, "first comment"}, + {"second insert", 123, "second comment"}, + {"third insert", 1337, "third comment"}, + } + + for _, tc := range testCases { + t.Run(tc.name, func(t *testing.T) { + eng.randomIDGen = func() int64 { return tc.id } + + obj := &RandomIDRecord{Comment: tc.comment} + _, err := eng.Insert(obj) + require.NoError(t, err) + require.Equal(t, tc.id, obj.ID, "ID should match current snowflake value") + + // Verify database entry + var retrieved RandomIDRecord + has, err := eng.ID(tc.id).Get(&retrieved) + require.NoError(t, err) + require.True(t, has) + require.Equal(t, tc.comment, retrieved.Comment) + }) + } +}