72 lines
1.4 KiB
Go
72 lines
1.4 KiB
Go
package svc
|
||
|
||
import (
|
||
"context"
|
||
"member/internal/model"
|
||
"os"
|
||
"time"
|
||
|
||
"github.com/bwmarrin/snowflake"
|
||
)
|
||
|
||
type SnackFlow interface {
|
||
GetSnackFlowNode() *snowflake.Node
|
||
Generate() snowflake.ID
|
||
}
|
||
|
||
type SnackFlowNode struct {
|
||
*snowflake.Node
|
||
}
|
||
|
||
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 {
|
||
ctx := context.Background()
|
||
nodeName := os.Getenv("POD_NAME")
|
||
if nodeName == "" {
|
||
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
|
||
}
|
||
|
||
func getMachineNodeID(machineNodeID int64) int64 {
|
||
// Snowflake 公式，工作機器 ID 佔用 10bit，最多容納 1024 節點，
|
||
// 故用 % 1024 取餘數做 ring
|
||
const nodeMax = 1024
|
||
|
||
return machineNodeID % nodeMax
|
||
}
|
||
|
||
func (n *SnackFlowNode) GetSnackFlowNode() *snowflake.Node {
|
||
return n.Node
|
||
}
|
||
|
||
func (n *SnackFlowNode) Generate() snowflake.ID {
|
||
return n.Node.Generate()
|
||
}
|