pangolin/server/internal/agentd/heartbeat.go

package agentd

import (
	"context"
	"net/http"
	"time"

	agentv1 "github.com/wangjia/pangolin/server/internal/pb/agentv1"
)

// LoadSource provides the runtime metrics reported in each heartbeat. The
// production implementation can read sing-box stats + host CPU; the default only
// reports the online peer count (number of provisioned credentials), which never
// reveals identities.
type LoadSource interface {
	Load() (onlinePeers int32, upBps, downBps int64, cpuPercent float64)
}

// defaultLoadSource reports peer count from the SingBox table and zeros for the
// rest. Bandwidth/CPU collection is a documented extension point.
type defaultLoadSource struct{ sb *SingBox }

func (d defaultLoadSource) Load() (int32, int64, int64, float64) {
	return int32(d.sb.OnlinePeers()), 0, 0, 0
}

// HealthSource reports whether the local data plane (sing-box) is healthy, sent
// in each heartbeat so the control plane can mark a node down when its sing-box
// is dead even though the agent's gRPC stream is still up.
type HealthSource interface{ Healthy() bool }

// clashHealthSource probes sing-box's clash_api /version. A reachable, 200-OK API
// means the sing-box process is alive and its config loaded. Applies a 2-strike
// debounce: a single failed probe is tolerated (covers the sing-box restart
// window); two consecutive failures report unhealthy.
type clashHealthSource struct {
	probe func() bool // single probe, true = healthy; injectable for tests
	fails int
}

func newClashHealthSource() *clashHealthSource {
	client := &http.Client{Timeout: 2 * time.Second}
	return &clashHealthSource{probe: func() bool {
		req, err := http.NewRequestWithContext(
			context.Background(), http.MethodGet, "http://"+clashAPIAddr+"/version", nil)
		if err != nil {
			return false
		}
		if clashAPISecret != "" {
			req.Header.Set("Authorization", "Bearer "+clashAPISecret)
		}
		resp, err := client.Do(req)
		if err != nil {
			return false
		}
		defer resp.Body.Close()
		return resp.StatusCode == http.StatusOK
	}}
}

func (h *clashHealthSource) Healthy() bool {
	if h.probe() {
		h.fails = 0
		return true
	}
	h.fails++
	return h.fails < 2 // tolerate 1 failure; unhealthy from the 2nd consecutive
}

// runHeartbeat sends a heartbeat every cfg.HeartbeatInterval until ctx is
// cancelled or an RPC error occurs. A need_full_resync response returns
// errResync so the session re-Registers and overwrites local state.
func (a *Agent) runHeartbeat(ctx context.Context, client agentv1.AgentServiceClient) error {
	ticker := time.NewTicker(a.cfg.HeartbeatInterval)
	defer ticker.Stop()
	for {
		if err := a.sendHeartbeat(ctx, client); err != nil {
			return err
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case <-ticker.C:
		}
	}
}

func (a *Agent) sendHeartbeat(ctx context.Context, client agentv1.AgentServiceClient) error {
	peers, up, down, cpu := a.load.Load()
	resp, err := client.Heartbeat(ctx, &agentv1.HeartbeatRequest{
		NodeUUID:         a.NodeUUID(),
		ConfigVersion:    a.sb.ConfigVersion(),
		OnlinePeers:      peers,
		BandwidthUpBps:   up,
		BandwidthDownBps: down,
		CPUPercent:       cpu,
		TimestampUnix:    a.clock().Unix(),
		DataPlaneHealthy: a.health.Healthy(),
	})
	if err != nil {
		return err
	}
	if resp.NeedFullResync {
		logf("control plane requested full resync")
		return errResync
	}
	return nil
}