65 lines
1.4 KiB
Go
65 lines
1.4 KiB
Go
package module
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"git.hlsq.asia/mmorpg/service-common/db/etcd"
|
|
"git.hlsq.asia/mmorpg/service-common/utils"
|
|
clientv3 "go.etcd.io/etcd/client/v3"
|
|
"go.etcd.io/etcd/client/v3/concurrency"
|
|
"math/rand"
|
|
)
|
|
|
|
// Snowflake 雪花模块
|
|
type Snowflake struct {
|
|
DefaultModule
|
|
snowflakeSession *concurrency.Session
|
|
}
|
|
|
|
func (m *Snowflake) Init() error {
|
|
node, session, err := acquire(context.Background(), 1, 1000)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
m.snowflakeSession = session
|
|
utils.InitSnowflake(node)
|
|
return nil
|
|
}
|
|
|
|
func (m *Snowflake) Stop() error {
|
|
_ = m.snowflakeSession.Close()
|
|
return nil
|
|
}
|
|
|
|
func acquire(ctx context.Context, min, max int) (int64, *concurrency.Session, error) {
|
|
nums := rand.Perm(max - min + 1)
|
|
for i := range nums {
|
|
nums[i] += min
|
|
}
|
|
|
|
for _, n := range nums {
|
|
key := fmt.Sprintf("node/num/%d", n)
|
|
|
|
session, err := concurrency.NewSession(
|
|
etcd.GetClient().Raw(),
|
|
concurrency.WithContext(ctx),
|
|
)
|
|
if err != nil {
|
|
return 0, nil, utils.ErrorsWrap(fmt.Errorf("etcd NewSession error: %v", err))
|
|
}
|
|
|
|
txnResp, _ := etcd.GetClient().Raw().Txn(ctx).
|
|
If(clientv3.Compare(clientv3.CreateRevision(key), "=", 0)).
|
|
Then(clientv3.OpPut(key, "", clientv3.WithLease(session.Lease()))).
|
|
Commit()
|
|
|
|
if txnResp.Succeeded {
|
|
return int64(n), session, nil
|
|
} else {
|
|
_ = session.Close()
|
|
}
|
|
}
|
|
return 0, nil, utils.ErrorsWrap(errors.New("etcd num empty"), "acquire error")
|
|
}
|