2024-07-21 15:57:56 +00:00
|
|
|
|
package svc
|
|
|
|
|
|
|
|
|
|
import (
|
|
|
|
|
"context"
|
|
|
|
|
"member/internal/model"
|
|
|
|
|
"os"
|
|
|
|
|
"time"
|
2024-07-25 14:01:22 +00:00
|
|
|
|
|
|
|
|
|
"github.com/bwmarrin/snowflake"
|
2024-07-21 15:57:56 +00:00
|
|
|
|
)
|
|
|
|
|
|
2024-08-02 02:10:25 +00:00
|
|
|
|
type SnackFlow interface {
|
|
|
|
|
GetSnackFlowNode() *snowflake.Node
|
|
|
|
|
Generate() snowflake.ID
|
2024-07-21 15:57:56 +00:00
|
|
|
|
}
|
|
|
|
|
|
2024-08-02 02:10:25 +00:00
|
|
|
|
type SnackFlowNode struct {
|
|
|
|
|
*snowflake.Node
|
2024-07-21 15:57:56 +00:00
|
|
|
|
}
|
|
|
|
|
|
2024-08-02 02:10:25 +00:00
|
|
|
|
func NewSnackFlow(node model.MachineNodeModel) (SnackFlow, error) {
|
|
|
|
|
nodeID := newMachineNodeID(node)
|
|
|
|
|
ringNodeID := getMachineNodeID(nodeID)
|
|
|
|
|
s, err := snowflake.NewNode(ringNodeID)
|
|
|
|
|
if err != nil {
|
|
|
|
|
return nil, err
|
|
|
|
|
}
|
|
|
|
|
return &SnackFlowNode{Node: s}, nil
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
func newMachineNodeID(node model.MachineNodeModel) int64 {
|
2024-07-21 15:57:56 +00:00
|
|
|
|
ctx := context.Background()
|
|
|
|
|
nodeName := os.Getenv("POD_NAME")
|
2024-08-02 02:10:25 +00:00
|
|
|
|
if nodeName == "" {
|
2024-07-21 15:57:56 +00:00
|
|
|
|
nodeName = "default_node"
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
machine, err := node.FindOneByHostName(ctx, nodeName)
|
|
|
|
|
if err != nil {
|
|
|
|
|
result, err := node.Insert(ctx, &model.MachineNode{
|
|
|
|
|
CreateTime: time.Now().Unix(),
|
|
|
|
|
HostName: nodeName,
|
|
|
|
|
})
|
|
|
|
|
if err != nil {
|
|
|
|
|
return 1
|
|
|
|
|
}
|
|
|
|
|
id, err := result.LastInsertId()
|
|
|
|
|
if err != nil {
|
|
|
|
|
return 1
|
|
|
|
|
}
|
|
|
|
|
return id
|
|
|
|
|
}
|
|
|
|
|
|
|
|
|
|
return machine.Id
|
|
|
|
|
}
|
|
|
|
|
|
2024-08-02 02:10:25 +00:00
|
|
|
|
func getMachineNodeID(machineNodeID int64) int64 {
|
|
|
|
|
// Snowflake 公式,工作機器 ID 佔用 10bit,最多容納 1024 節點,
|
2024-07-21 15:57:56 +00:00
|
|
|
|
// 故用 % 1024 取餘數做 ring
|
|
|
|
|
const nodeMax = 1024
|
|
|
|
|
|
|
|
|
|
return machineNodeID % nodeMax
|
|
|
|
|
}
|
|
|
|
|
|
2024-08-02 02:10:25 +00:00
|
|
|
|
func (n *SnackFlowNode) GetSnackFlowNode() *snowflake.Node {
|
|
|
|
|
return n.Node
|
|
|
|
|
}
|
2024-07-21 15:57:56 +00:00
|
|
|
|
|
2024-08-02 02:10:25 +00:00
|
|
|
|
func (n *SnackFlowNode) Generate() snowflake.ID {
|
|
|
|
|
return n.Node.Generate()
|
2024-07-21 15:57:56 +00:00
|
|
|
|
}
|