3x-ui-agent/internal/web/job/node_heartbeat_job.go

package job

import (
	"context"
	"strconv"
	"sync"
	"time"

	"github.com/mhsanaei/3x-ui/v3/internal/database"
	"github.com/mhsanaei/3x-ui/v3/internal/database/model"
	"github.com/mhsanaei/3x-ui/v3/internal/eventbus"
	"github.com/mhsanaei/3x-ui/v3/internal/logger"
	"github.com/mhsanaei/3x-ui/v3/internal/util/common"
	"github.com/mhsanaei/3x-ui/v3/internal/web/service"
	"github.com/mhsanaei/3x-ui/v3/internal/web/websocket"
)

const (
	nodeHeartbeatConcurrency    = 32
	nodeHeartbeatRequestTimeout = 4 * time.Second
)

// agentHeartbeatStaleAfter is how long an agent-managed node may go without
// an agent heartbeat before the heartbeat job marks it offline. The node-side
// agent heartbeats every 60s (defaultHeartbeatInterval in internal/agent);
// the threshold is max(3 intervals, 120s) = 180s, so two consecutive lost
// heartbeats still don't flap the node offline.
const agentHeartbeatStaleAfter = 180 * time.Second

type NodeHeartbeatJob struct {
	nodeService service.NodeService
	running     sync.Mutex
}

func NewNodeHeartbeatJob() *NodeHeartbeatJob {
	return &NodeHeartbeatJob{}
}

func (j *NodeHeartbeatJob) Run() {
	if !j.running.TryLock() {
		return
	}
	defer j.running.Unlock()

	nodes, err := j.nodeService.GetAll()
	if err != nil {
		logger.Warning("node heartbeat: load nodes failed:", err)
		return
	}
	if len(nodes) == 0 {
		return
	}

	sem := make(chan struct{}, nodeHeartbeatConcurrency)
	var wg sync.WaitGroup
	for _, n := range nodes {
		if !n.Enable {
			continue
		}
		wg.Add(1)
		sem <- struct{}{}
		n := n
		common.GoRecover("node-heartbeat:"+n.Name, func() {
			defer wg.Done()
			defer func() { <-sem }()
			j.probeOne(n)
		})
	}
	wg.Wait()

	if !websocket.HasClients() {
		return
	}
	updated, err := j.nodeService.GetNodeTreeView()
	if err != nil {
		logger.Warning("node heartbeat: load nodes for broadcast failed:", err)
		return
	}
	websocket.BroadcastNodes(updated)
}

func (j *NodeHeartbeatJob) probeOne(n *model.Node) {
	// Agent-managed nodes are heartbeated via the agent token, not the
	// legacy panel probe. Skip the probe so we don't clobber the
	// agent-driven online status — but do expire that status when the
	// agent stops heartbeating, or an uninstalled agent would leave its
	// node showing online forever.
	var agentCount int64
	_ = database.GetDB().Model(&model.AgentNode{}).Where("node_id = ?", n.Id).Count(&agentCount).Error
	if agentCount > 0 {
		j.expireStaleAgentNode(n)
		return
	}
	ctx, cancel := context.WithTimeout(context.Background(), nodeHeartbeatRequestTimeout)
	defer cancel()
	prevStatus := n.Status
	patch, err := j.nodeService.Probe(ctx, n)
	if err != nil {
		patch.Status = "offline"
	} else {
		patch.Status = "online"
	}
	if updErr := j.nodeService.UpdateHeartbeat(n.Id, patch); updErr != nil {
		logger.Warning("node heartbeat: update node", n.Id, "failed:", updErr)
	}
	publishNodeTransition(n, prevStatus, patch)
	// Learn the nodes this node manages so the panel can surface them as
	// transitive sub-nodes (#4983). Fresh context — the probe budget above may
	// be spent. Drop them when the node is unreachable.
	if patch.Status == "online" {
		dctx, dcancel := context.WithTimeout(context.Background(), nodeHeartbeatRequestTimeout)
		j.nodeService.RefreshDescendants(dctx, n)
		dcancel()
	} else {
		j.nodeService.ClearDescendants(n.Id)
	}
}

// agentLastSeenMs returns the newest last_seen (Unix milliseconds, the unit
// the agent controller stamps on register/heartbeat) across the agents bound
// to the node, or 0 when no bound agent has ever checked in.
func agentLastSeenMs(nodeID int) int64 {
	var lastSeen *int64 // MAX() over zero rows is NULL
	_ = database.GetDB().Model(&model.AgentNode{}).
		Where("node_id = ?", nodeID).
		Select("MAX(last_seen)").
		Scan(&lastSeen).Error
	if lastSeen == nil {
		return 0
	}
	return *lastSeen
}

// agentHeartbeatStale decides expiry from the agent's last_seen stamp. A
// zero stamp means the token was issued but the agent never checked in: the
// node was never activated, so expiry leaves it alone (its status stays
// whatever creation left it — "unknown") instead of forcing a just-created
// node offline before its agent had a chance to come up.
func agentHeartbeatStale(lastSeenMs int64, now time.Time) bool {
	if lastSeenMs <= 0 {
		return false
	}
	return now.Sub(time.UnixMilli(lastSeenMs)) > agentHeartbeatStaleAfter
}

// expireStaleAgentNode flips an agent-managed node offline once its agent's
// heartbeats stop. Only the status changes: the patch carries the node's
// stored heartbeat fields through unchanged (including the last_heartbeat
// stamp — the truth is "last seen then", not "now"), and the transition goes
// through the same UpdateHeartbeat + broadcast path as a failed panel probe,
// so node.down subscribers see agent expiries too.
func (j *NodeHeartbeatJob) expireStaleAgentNode(n *model.Node) {
	if n.Status == "offline" {
		return
	}
	if !agentHeartbeatStale(agentLastSeenMs(n.Id), time.Now()) {
		return
	}
	patch := service.HeartbeatPatch{
		Status:        "offline",
		LastHeartbeat: n.LastHeartbeat,
		LatencyMs:     n.LatencyMs,
		XrayVersion:   n.XrayVersion,
		PanelVersion:  n.PanelVersion,
		Guid:          n.Guid,
		CpuPct:        n.CpuPct,
		MemPct:        n.MemPct,
		UptimeSecs:    n.UptimeSecs,
		NetUp:         n.NetUp,
		NetDown:       n.NetDown,
		LastError:     n.LastError,
		XrayState:     n.XrayState,
		XrayError:     n.XrayError,
	}
	if err := j.nodeService.UpdateHeartbeat(n.Id, patch); err != nil {
		logger.Warning("node heartbeat: expire agent node", n.Id, "failed:", err)
		return
	}
	publishNodeTransition(n, n.Status, patch)
}

// publishNodeTransition emits node.down / node.up only on a genuine state change.
// An "unknown"/empty previous status (fresh start) is treated as not-online, so a
// node coming up for the first time fires node.up but never a spurious node.down.
func publishNodeTransition(n *model.Node, prevStatus string, patch service.HeartbeatPatch) {
	if EventBus == nil {
		return
	}
	var eventType eventbus.EventType
	switch {
	case prevStatus == "online" && patch.Status == "offline":
		eventType = eventbus.EventNodeDown
	case prevStatus != "online" && patch.Status == "online":
		eventType = eventbus.EventNodeUp
	default:
		return
	}
	source := n.Name
	if source == "" {
		source = "node-" + strconv.Itoa(n.Id)
	}
	EventBus.Publish(eventbus.Event{
		Type:   eventType,
		Source: source,
		Data: &eventbus.NodeHealthData{
			NodeId:    n.Id,
			LatencyMs: patch.LatencyMs,
			CpuPct:    patch.CpuPct,
			MemPct:    patch.MemPct,
			XrayState: patch.XrayState,
			XrayError: patch.XrayError,
		},
	})
}