ark-member/internal/svc/machine_node.go

package svc

import (
	"context"
	sf "member/internal/lib/snackflow"
	"member/internal/model"
	"os"
	"time"

	"github.com/bwmarrin/snowflake"
)

type machineNode struct {
	MachineNodeID int64 `json:"machine_node_id"`
}

type MachineNodeCreateParams struct {
	HostName string
}

func NewMachineNode(node model.MachineNodeModel) int64 {
	ctx := context.Background()
	nodeName := os.Getenv("POD_NAME")
	if os.Getenv("POD_NAME") == "" {
		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 newSnackFlowNode(node model.MachineNodeModel) (*snowflake.Node, error) {

	nodeID := NewMachineNode(node)
	ringNodeID := GetMachineNodeID(nodeID)
	s := sf.New(sf.WithMachineNodeID(ringNodeID))
	n, err := s.NewNode()
	if err != nil {
		return nil, err
	}

	return n, nil
}