Initial commit

common/genId/lid/lid.go

@@ -0,0 +1,102 @@
+package lid
+
+import (
+	"context"
+	"stat/common/mongos"
+	"time"
+
+	"go.mongodb.org/mongo-driver/bson"
+	"go.mongodb.org/mongo-driver/mongo"
+)
+
+const notFindMsg = "mongo: no documents in result"
+
+type IdModel struct {
+	BusinessId  string    `bson:"business_id"` //业务id
+	Step        int64     `bson:"step"`        //步长
+	MaxId       int64     `bson:"max_id"`      //最大id
+	Description string    `bson:"description"` //描述
+	CreatedAt   time.Time `bson:"createdAt"`   //创建时间
+	UpdatedAt   time.Time `bson:"updatedAt"`   //更改时间
+}
+
+type Lid struct {
+	db         *mongo.Collection // 数据库连接
+	businessId string            // 业务id
+	ch         chan int64        // id缓冲池
+	min, max   int64             // id段最小值，最大值
+}
+
+// NewLid 创建一个lid,db:数据库连接；businessId：业务id;len：缓冲池大小
+func NewLid(db *mongo.Collection, businessId string, len int) (*Lid, error) {
+	lid := Lid{
+		db:         db,
+		businessId: businessId,
+		ch:         make(chan int64, len),
+	}
+	go lid.productId()
+	return &lid, nil
+}
+
+// Get 获取自增id
+func (l *Lid) Get() int64 {
+	return <-l.ch
+}
+
+// productId 生产id
+func (l *Lid) productId() {
+	l.reset()
+	for {
+		if l.min >= l.max {
+			l.reset()
+
+		}
+
+		l.min++
+		l.ch <- l.min
+	}
+}
+
+// reset 不断尝试从数据库获取，直到成功
+func (l *Lid) reset() {
+	for {
+		err := l.getFromDB()
+		if err == nil {
+			return
+		}
+		time.Sleep(time.Second)
+		continue
+	}
+}
+
+// getFromDB 从数据库获取id段
+func (l *Lid) getFromDB() error {
+	var (
+		maxId   int64
+		step    int64
+		idModel IdModel
+	)
+	txErr := mongos.Transaction(context.Background(), func(sctx mongo.SessionContext) (err error) {
+		err = l.db.FindOne(sctx, bson.M{"business_id": l.businessId}).Decode(&idModel)
+		if err != nil {
+			if err.Error() == notFindMsg {
+				_, err = l.db.InsertOne(sctx, bson.M{"business_id": l.businessId, "max_id": 0, "step": 10, "description": "id自增器", "createdAt": time.Now(), "updatedAt": time.Now()})
+				if err != nil {
+					return
+				}
+			} else {
+				return
+			}
+		}
+		maxId = idModel.MaxId
+		step = idModel.Step
+		_, err = l.db.UpdateOne(sctx, bson.M{"business_id": l.businessId}, bson.M{"$set": bson.M{"max_id": maxId + step}})
+		if err != nil {
+			return
+		}
+		l.min = maxId
+		l.max = maxId + step
+		return
+	})
+	return txErr
+}