51 lines
1.2 KiB
Go
51 lines
1.2 KiB
Go
package imtenant
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"os"
|
|
"strconv"
|
|
"sync"
|
|
|
|
"github.com/gogf/gf/v2/frame/g"
|
|
sf "github.com/tinode/snowflake"
|
|
)
|
|
|
|
const defaultSnowflakeWorkerID = 1023
|
|
|
|
var (
|
|
snowflakeOnce sync.Once
|
|
snowflakeGen *sf.SnowFlake
|
|
snowflakeErr error
|
|
)
|
|
|
|
// NextSnowflakeID returns an unencrypted Snowflake ID suitable for IM BIGINT keys.
|
|
func NextSnowflakeID(ctx context.Context) (int64, error) {
|
|
snowflakeOnce.Do(func() {
|
|
workerID := g.Cfg().MustGet(ctx, "imDirectDb.snowflakeWorkerId", defaultSnowflakeWorkerID).Int()
|
|
if value := os.Getenv("IM_DIRECT_DB_SNOWFLAKE_WORKER_ID"); value != "" {
|
|
workerID, snowflakeErr = strconv.Atoi(value)
|
|
if snowflakeErr != nil {
|
|
snowflakeErr = fmt.Errorf("invalid IM_DIRECT_DB_SNOWFLAKE_WORKER_ID: %w", snowflakeErr)
|
|
return
|
|
}
|
|
}
|
|
if workerID < 0 || workerID > 1023 {
|
|
snowflakeErr = fmt.Errorf("IM Snowflake worker ID must be between 0 and 1023")
|
|
return
|
|
}
|
|
snowflakeGen, snowflakeErr = sf.NewSnowFlake(uint32(workerID))
|
|
})
|
|
if snowflakeErr != nil {
|
|
return 0, snowflakeErr
|
|
}
|
|
id, err := snowflakeGen.Next()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if id == 0 || id > uint64(^uint64(0)>>1) {
|
|
return 0, fmt.Errorf("generated invalid IM Snowflake ID")
|
|
}
|
|
return int64(id), nil
|
|
}
|