ossService/util/heartbeat.go

package util

import (
	"context"
	"encoding/json"
	"github.com/go-redis/redis/v8"
	"github.com/gogf/gf/v2/frame/g"
	"github.com/gogf/gf/v2/os/gctx"
	"log"
	"time"
)

var ctx = context.Background()
var rdb *redis.Client

// InitRedis 初始化redis客户端
func InitRedis() {
	ctx := gctx.New()
	rdb = redis.NewClient(&redis.Options{
		Addr:     g.Config().MustGet(ctx, "redis.address").String(),
		Password: g.Config().MustGet(ctx, "redis.password").String(),
		DB:       0,
	})
}

// NodeInfo 描述节点心跳信息
type NodeInfo struct {
	NodeName  string `json:"node_name"`
	Timestamp int64  `json:"timestamp"`
}

// SendHeartbeat 将节点心跳信息写入redis，过期时间为10秒
func SendHeartbeat() {
	nodeName := g.Config().MustGet(ctx, "node.node_name").String()
	node := NodeInfo{
		NodeName:  nodeName,
		Timestamp: time.Now().Unix(),
	}
	data, err := json.Marshal(node)
	if err != nil {
		log.Println("Heartbeat marshal error", err)
		return
	}
	// key 为 "node:<nodeName>"
	rdb.Set(ctx, "node:"+nodeName, data, 10*time.Second)
}

// CheckOnlineNodes 检查在线节点数
func CheckOnlineNodes() ([]NodeInfo, error) {
	keys, err := rdb.Keys(ctx, "node:*").Result()
	if err != nil {
		return nil, err
	}
	var nodes []NodeInfo
	for _, key := range keys {
		data, err := rdb.Get(ctx, key).Result()
		if err != nil {
			continue
		}
		var node NodeInfo
		json.Unmarshal([]byte(data), &node)
		nodes = append(nodes, node)
	}
	return nodes, nil
}

// GetOnlineNodesJSON 返回在线节点信息的JSON格式（供HTTP接口调用）
func GetOnlineNodesJSON() ([]byte, error) {
	nodes, err := CheckOnlineNodes()
	if err != nil {
		return nil, err
	}
	return json.Marshal(nodes)
}