ark-member/internal/svc/machine_node.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()
}