From 32e00f9438ed29fc26351f65cf7d98eefd1d838e Mon Sep 17 00:00:00 2001
From: zhangqian <zhangqian@123.com>
Date: 星期五, 08 十二月 2023 14:26:21 +0800
Subject: [PATCH] 换个雪花算法

---
 pkg/snowflake/snowflake.go |  140 +++++++++++++++++++++++++++++++++++++---------
 1 files changed, 112 insertions(+), 28 deletions(-)

diff --git a/pkg/snowflake/snowflake.go b/pkg/snowflake/snowflake.go
index 049b307..f0f6442 100644
--- a/pkg/snowflake/snowflake.go
+++ b/pkg/snowflake/snowflake.go
@@ -2,26 +2,18 @@
 
 import (
 	"apsClient/pkg/logx"
+	"errors"
 	"fmt"
-	"github.com/bwmarrin/snowflake"
 	"hash/fnv"
 	"net"
+	"strconv"
+	"sync"
+	"time"
 )
 
-var Node *snowflake.Node
+var idGenerater *IdWorker
 
-func Init(nodeId int64) error {
-	// Create a new Node with a Node number of 0~1023
-	var err error
-	Node, err = snowflake.NewNode(nodeId)
-	if err != nil {
-		logx.Errorf("snowflake NewNode error:%v", err)
-		return err
-	}
-	return nil
-}
-
-func InitWithIP() {
+func init() {
 	// 浣跨敤 LookupIP 鑾峰彇涓绘満鐨� IP 鍦板潃鍒楄〃
 	// 鑾峰彇鏈満鎵�鏈夌綉缁滄帴鍙�
 	interfaces, err := net.Interfaces()
@@ -69,20 +61,11 @@
 		logx.Errorf("snowflake can not generate, init error, ip to number error :%v", err)
 		panic(fmt.Sprintf("snowflake can not generate, init error, ip to number error :%v", err))
 	}
-	err = Init(ipNumber)
+	idGenerater, err = NewIdWorker(ipNumber)
 	if err != nil {
 		logx.Errorf("snowflake can not generate, init error :%v", err)
 		panic(fmt.Sprintf("snowflake can not generate, init error :%v", err))
 	}
-}
-
-func StringToNumber(input string) (int64, error) {
-	hash := fnv.New32a()
-	_, err := hash.Write([]byte(input))
-	if err != nil {
-		return 0, err
-	}
-	return int64(hash.Sum32() % 1024), nil // 鍙栦綑鏁扮‘淇濈粨鏋滃湪 1 鍒� 1023 涔嬮棿
 }
 
 func ipToNumber(ip string) (int64, error) {
@@ -98,9 +81,110 @@
 	return int64(hash.Sum32() % 1024), nil // 鍙栦綑鏁扮‘淇濈粨鏋滃湪 1 鍒� 1023 涔嬮棿
 }
 
-func GenerateID() int64 {
-	if Node == nil {
-		InitWithIP()
+const (
+	CEpoch         = 1474802888000
+	CWorkerIdBits  = 10 // Num of WorkerId Bits
+	CSenquenceBits = 12 // Num of Sequence Bits
+
+	CWorkerIdShift  = 12
+	CTimeStampShift = 22
+
+	CSequenceMask = 0xfff // equal as getSequenceMask()
+	CMaxWorker    = 0x3ff // equal as getMaxWorkerId()
+)
+
+type IdWorker struct {
+	workerId      int64
+	lastTimeStamp int64
+	sequence      int64
+	maxWorkerId   int64
+	lock          *sync.Mutex
+}
+
+func NewIdWorker(workerId int64) (iw *IdWorker, err error) {
+	iw = new(IdWorker)
+
+	iw.maxWorkerId = getMaxWorkerId()
+
+	if workerId > iw.maxWorkerId || workerId < 0 {
+		return nil, errors.New("worker not fit")
 	}
-	return int64(Node.Generate())
+	iw.workerId = workerId
+	iw.lastTimeStamp = -1
+	iw.sequence = 0
+	iw.lock = new(sync.Mutex)
+	return iw, nil
+}
+
+func getMaxWorkerId() int64 {
+	return -1 ^ -1<<CWorkerIdBits
+}
+
+func getSequenceMask() int64 {
+	return -1 ^ -1<<CSenquenceBits
+}
+
+// return in ms
+func (iw *IdWorker) timeGen() int64 {
+	return time.Now().UnixNano() / 1000 / 1000
+}
+
+func (iw *IdWorker) timeReGen(last int64) int64 {
+	ts := time.Now().UnixNano() / 1000 / 1000
+	for {
+		if ts <= last {
+			ts = iw.timeGen()
+		} else {
+			break
+		}
+	}
+	return ts
+}
+
+func (iw *IdWorker) NextId() (ts int64, err error) {
+	iw.lock.Lock()
+	defer iw.lock.Unlock()
+	ts = iw.timeGen()
+	if ts == iw.lastTimeStamp {
+		iw.sequence = (iw.sequence + 1) & CSequenceMask
+		if iw.sequence == 0 {
+			ts = iw.timeReGen(ts)
+		}
+	} else {
+		iw.sequence = 0
+	}
+
+	if ts < iw.lastTimeStamp {
+		err = errors.New("Clock moved backwards, Refuse gen id")
+		return 0, err
+	}
+	iw.lastTimeStamp = ts
+	ts = (ts-CEpoch)<<CTimeStampShift | iw.workerId<<CWorkerIdShift | iw.sequence
+	return ts, nil
+}
+
+func ParseId(id int64) (t time.Time, ts int64, workerId int64, seq int64) {
+	seq = id & CSequenceMask
+	workerId = (id >> CWorkerIdShift) & CMaxWorker
+	ts = (id >> CTimeStampShift) + CEpoch
+	t = time.Unix(ts/1000, (ts%1000)*1000000)
+	return
+}
+
+func GenerateID() int64 {
+start:
+	id, err := idGenerater.NextId()
+	if err != nil {
+		goto start
+	}
+	return id
+}
+
+func GenerateIdStr() string {
+start:
+	id, err := idGenerater.NextId()
+	if err != nil {
+		goto start
+	}
+	return strconv.FormatInt(id, 10)
 }

--
Gitblit v1.8.0