103 lines
2.3 KiB
Go
103 lines
2.3 KiB
Go
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
|
|
}
|