machine.go
package snowf

type Machine interface {
	// GetID 获取机器 ID (0 <= ID < 1024)
	GetID() int64
}

machine_redis.go

package snowf

import (
	"fmt"
	"github.com/go-redis/redis"
	"sync"
	"time"
)

const (
	MachineIDLockPrefix = "machine-id-lock"
	MachineIncrPrefix   = "machine-id-incr"
)

type MachineRedis struct {
	redisClient   *redis.Client
	serverName    string
	machineID     int64
	machineIDOnce sync.Once
}

func (m *MachineRedis) GetID() int64 {
	if nil == m.redisClient || "" == m.serverName {
		panic("machine not initial")
	}

	m.machineIDOnce.Do(func() {
		redisClient := m.redisClient
		for {
			// 尝试获取 Redis 锁
			lockKey := fmt.Sprintf("%s:%s", MachineIDLockPrefix, m.serverName)
			locked, err := redisClient.SetNX(lockKey, 1, 1*time.Second).Result()
			if err != nil {
				time.Sleep(100 * time.Millisecond)
				continue
			}

			if locked {
				// 获取锁成功，执行自增操作
				incrCmd := redisClient.Incr(fmt.Sprintf("%s:%s", MachineIncrPrefix, m.serverName))
				m.machineID = incrCmd.Val() % 1024
				// 释放锁
				redisClient.Del(lockKey)
				break
			}

			// 获取锁失败，等待一段时间后重试
			time.Sleep(100 * time.Millisecond)
		}
	})

	return m.machineID
}

func NewMachineRedis(redisClient *redis.Client, serverName string) *MachineRedis {
	return &MachineRedis{
		redisClient: redisClient,
		serverName:  serverName,
	}
}

snowflake.go

package snowf

import (
	"errors"
	"github.com/bwmarrin/snowflake"
	"sync"
	"time"
)

type Snowflake struct {
	Machine Machine
	Valid   bool
	nodes   map[string]*snowflake.Node
	mu      sync.RWMutex
}

// InitPartitionNodes 初始化其他分区，防止并发情况下GetPartitionNode出现竞争
func InitPartitionNodes(partitions ...string) error {
	if !currentClient.Valid {
		return errors.New("service is not initial")
	}
	for _, partition := range partitions {
		_, err := currentClient.InitPartitionNode(partition)
		if err != nil {
			return err
		}
	}

	return nil
}

func (sf *Snowflake) InitPartitionNode(partition string) (node *snowflake.Node, err error) {
	sf.mu.Lock()
	defer sf.mu.Unlock()

	node, err = snowflake.NewNode(sf.Machine.GetID())
	if err != nil {
		return nil, err
	}

	sf.nodes[partition] = node
	return node, nil
}

func (sf *Snowflake) GetPartitionNode(partition string) (node *snowflake.Node, err error) {
	node, ok := sf.nodes[partition]
	if !ok {
		node, err = sf.InitPartitionNode(partition)
		if err != nil {
			return nil, err
		}
	}

	return node, nil
}

var currentClient *Snowflake

func InitSnowflake(machine Machine, partitions ...string) (err error) {
	currentClient = &Snowflake{
		Machine: machine,
		nodes:   make(map[string]*snowflake.Node),
	}

	partition := "default"
	if len(partitions) > 0 {
		partition = partitions[0]
	}

	_, err = currentClient.InitPartitionNode(partition)
	if err != nil {
		return err
	}
	currentClient.Valid = true
	return nil
}

// GetSnowID
// Deprecated: recommend to use GetIDInt64 method
func GetSnowID(partitions ...string) int64 {
	return int64(GetID(partitions...))
}

func GetIDInt64(partitions ...string) int64 {
	return int64(GetID(partitions...))
}

func GetID(partitions ...string) ID {
	if !currentClient.Valid {
		panic("service is not initial")
	}

	partition := "default"
	if len(partitions) > 0 {
		partition = partitions[0]
	}

	node, err := currentClient.GetPartitionNode(partition)
	if err != nil {
		// 服务降级
		return ID(time.Now().UnixNano())
	}
	return ID(node.Generate())
}

showid.go

package snowf

import (
	"database/sql/driver"
	"encoding/json"
	"fmt"
	"github.com/bwmarrin/snowflake"
	"gorm.io/gorm"
	"gorm.io/gorm/schema"
	"strconv"
	"time"
)

// ID is a custom type for IDs that should be serialized as a string in JSON.
type ID int64

// GormDataType specifies the GORM data type for ID.
func (ID) GormDataType() string {
	return "bigint" // Use the appropriate data type for your database.
}

// GormDBDataType specifies the GORM database data type for ID.
func (ID) GormDBDataType(db *gorm.DB, field *schema.Field) string {
	switch db.Dialector.Name() {
	case "sqlite", "mysql":
		return "BIGINT"
	case "postgres":
		return "BIGINT"
	// Add other databases as needed.
	default:
		return "BIGINT"
	}
}

// MarshalJSON customizes the JSON encoding of ID to be a string.
func (sid ID) MarshalJSON() ([]byte, error) {
	strID := strconv.FormatInt(int64(sid), 10)
	return []byte(`"` + strID + `"`), nil
}

// UnmarshalJSON customizes the JSON decoding of ID from a string.
func (sid *ID) UnmarshalJSON(data []byte) error {
	var s string
	if err := json.Unmarshal(data, &s); err != nil {
		return err
	}
	if s != "" {
		id, err := strconv.ParseInt(s, 10, 64)
		if err != nil {
			return fmt.Errorf("failed to parse ID from JSON: %w", err)
		}
		*sid = ID(id)
	}
	return nil
}

// Value returns the ID as a driver.Value, which is required by GORM.
func (sid ID) Value() (driver.Value, error) {
	return int64(sid), nil
}

// Scan scans the value from the database driver into ID, which is required by GORM.
func (sid *ID) Scan(value interface{}) error {
	var id int64
	switch v := value.(type) {
	case int64:
		id = v
	case []byte:
		id, _ = strconv.ParseInt(string(v), 10, 64)
	case string:
		id, _ = strconv.ParseInt(v, 10, 64)
	default:
		return fmt.Errorf("unsupported type: %T", value)
	}
	*sid = ID(id)
	return nil
}

func (sid ID) String() string {
	return strconv.FormatInt(int64(sid), 10)
}

func (sid ID) IsValid() bool {
	return int64(sid) > 0
}

func (sid ID) Time() time.Time {
	timestamp := int64(sid) >> (snowflake.NodeBits + snowflake.StepBits)
	return time.Unix(snowflake.Epoch/1000+timestamp/1000, (timestamp%1000)*1000000).UTC()
}