3x-ui-agent/internal/agent/ping.go

package agent

import (
	"context"
	"net"
	"sync"
	"time"

	"github.com/mhsanaei/3x-ui/v3/internal/logger"
)

// Item 7 B4: latency probes, the agent side. The master pushes the
// probe task set on the config poll (probe_tasks); the agent runs one
// loop per task — resolve, TCP-connect, time the handshake — and the
// samples ride the next heartbeat back (probeResults). A port of
// monitor-probe's agent prober (agent/src/main.rs tcp_probe and its
// per-task scheduler), with the same numbers: at most 64 tasks, a
// 900 ms budget each for resolution and for every dial, at most three
// addresses tried, and a resolution overrun producing NO sample at
// all rather than a recorded miss.

// ProbeTask is one probe target as pushed by the master.
type ProbeTask struct {
	ID       int64  `json:"id"`
	Target   string `json:"target"`
	Interval int64  `json:"interval"`
}

// ProbeResult is one filed sample: the task and its handshake latency
// in milliseconds; negative means the target did not answer.
type ProbeResult struct {
	TaskID    int64 `json:"taskId"`
	LatencyMs int64 `json:"latencyMs"`
}

const (
	// maxProbeTasksRun mirrors the master's per-agent assignment cap
	// (probe.MaxProbesPerNode); an over-long list is truncated, never
	// run whole, with the loss logged once per apply.
	maxProbeTasksRun = 64
	// probeDialBudget bounds name resolution and every single dial.
	probeDialBudget = 900 * time.Millisecond
	// probeMaxAddrs caps how many resolved addresses one probe tries.
	probeMaxAddrs = 3
	// probeResultCap bounds the unsent-sample buffer; past it the
	// oldest samples drop (they describe a master outage, and their
	// timestamps would be discarded on arrival anyway).
	probeResultCap = 8192
	// Probe intervals are clamped into the master's legal range on
	// arrival as a backstop (the panel refuses out-of-range values).
	minProbeIntervalSec = 5
	maxProbeIntervalSec = 3600
)

type runningProbe struct {
	target   string
	interval int64
	cancel   context.CancelFunc
}

// probeRunner owns the running probe loops and the result buffer.
// All state is behind mu; probe loops only touch it through methods.
type probeRunner struct {
	mu            sync.Mutex
	tasks         map[int64]*runningProbe
	results       []ProbeResult
	dropped       int64
	resolveLogged map[int64]bool
	stopped       bool
}

func newProbeRunner() *probeRunner {
	return &probeRunner{
		tasks:         map[int64]*runningProbe{},
		resolveLogged: map[int64]bool{},
	}
}

// clampProbeInterval folds an interval into the legal range.
func clampProbeInterval(sec int64) int64 {
	if sec < minProbeIntervalSec {
		return minProbeIntervalSec
	}
	if sec > maxProbeIntervalSec {
		return maxProbeIntervalSec
	}
	return sec
}

// applyTasks reconciles the running loops with a freshly pushed set:
// a task whose (target, interval) is unchanged keeps its loop and its
// phase; anything else is cancelled and, if still wanted, restarted.
func (r *probeRunner) applyTasks(tasks []ProbeTask) {
	if len(tasks) > maxProbeTasksRun {
		logger.Warningf("probe: master pushed %d tasks, running only the first %d", len(tasks), maxProbeTasksRun)
		tasks = tasks[:maxProbeTasksRun]
	}
	want := make(map[int64]ProbeTask, len(tasks))
	for _, t := range tasks {
		t.Interval = clampProbeInterval(t.Interval)
		want[t.ID] = t
	}

	r.mu.Lock()
	defer r.mu.Unlock()
	if r.stopped {
		return
	}
	for id, running := range r.tasks {
		t, ok := want[id]
		if !ok || t.Target != running.target || t.Interval != running.interval {
			running.cancel()
			delete(r.tasks, id)
		}
	}
	for id, t := range want {
		if _, ok := r.tasks[id]; ok {
			continue
		}
		ctx, cancel := context.WithCancel(context.Background())
		r.tasks[id] = &runningProbe{target: t.Target, interval: t.Interval, cancel: cancel}
		go r.runLoop(ctx, t)
	}
}

// stop cancels every probe loop (agent shutdown).
func (r *probeRunner) stop() {
	r.mu.Lock()
	defer r.mu.Unlock()
	r.stopped = true
	for id, running := range r.tasks {
		running.cancel()
		delete(r.tasks, id)
	}
}

// runLoop probes immediately, then on the task's interval. A tick
// that arrives while a probe is still running is dropped rather than
// queued (the hub's MissedTickBehavior::Delay; a Go ticker behaves
// the same way).
func (r *probeRunner) runLoop(ctx context.Context, task ProbeTask) {
	r.probeOnce(task)
	ticker := time.NewTicker(time.Duration(task.Interval) * time.Second)
	defer ticker.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-ticker.C:
			r.probeOnce(task)
		}
	}
}

func (r *probeRunner) probeOnce(task ProbeTask) {
	ms, sampled := probeTCP(task.Target, func() {
		r.noteResolveTimeout(task.ID, task.Target)
	})
	if !sampled {
		return
	}
	r.addResult(ProbeResult{TaskID: task.ID, LatencyMs: ms})
}

// noteResolveTimeout logs a skipped round once per task per process:
// a target whose name stops resolving would otherwise be silent
// forever (nothing is filed for it), and one line per interval per
// target would flood the log on a DNS outage (hub discipline).
func (r *probeRunner) noteResolveTimeout(taskID int64, target string) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if r.resolveLogged[taskID] {
		return
	}
	r.resolveLogged[taskID] = true
	logger.Warningf("probe: resolving %s overran %s; round skipped (further rounds log nothing)", target, probeDialBudget)
}

func (r *probeRunner) addResult(res ProbeResult) {
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.results) >= probeResultCap {
		r.results = r.results[1:]
		r.dropped++
	}
	r.results = append(r.results, res)
}

// drainResults takes the buffered samples for one heartbeat. The
// caller requeues them when the heartbeat fails, so a master outage
// loses nothing that still fits the buffer.
func (r *probeRunner) drainResults() []ProbeResult {
	r.mu.Lock()
	defer r.mu.Unlock()
	if len(r.results) == 0 {
		return nil
	}
	out := r.results
	r.results = nil
	return out
}

// requeueResults puts a failed heartbeat's batch back ahead of
// anything probed since, trimming the oldest if the buffer overflows.
func (r *probeRunner) requeueResults(batch []ProbeResult) {
	if len(batch) == 0 {
		return
	}
	r.mu.Lock()
	defer r.mu.Unlock()
	merged := make([]ProbeResult, 0, len(batch)+len(r.results))
	merged = append(merged, batch...)
	merged = append(merged, r.results...)
	if len(merged) > probeResultCap {
		merged = merged[len(merged)-probeResultCap:]
	}
	r.results = merged
}

// probeTCP measures one target's TCP handshake latency in
// milliseconds. It returns (ms, true) for a sample to file — the
// handshake time, or -1 when the target did not answer — and
// (0, false) when the round produced no sample: name resolution
// overran its budget, which the hub treats as "nothing happened"
// rather than a miss, because filing it would blame the target for
// the resolver's state. A malformed target is -1: it is re-probed
// every interval and reads as total loss, which is what it is.
func probeTCP(target string, onResolveTimeout func()) (int64, bool) {
	host, port, err := net.SplitHostPort(target)
	if err != nil {
		return -1, true
	}
	rctx, cancel := context.WithTimeout(context.Background(), probeDialBudget)
	addrs, err := (&net.Resolver{}).LookupIPAddr(rctx, host)
	timedOut := rctx.Err() == context.DeadlineExceeded
	cancel()
	if err != nil {
		if timedOut {
			if onResolveTimeout != nil {
				onResolveTimeout()
			}
			return 0, false
		}
		return -1, true
	}
	if len(addrs) == 0 {
		return -1, true
	}
	dialer := &net.Dialer{}
	tried := 0
	for _, addr := range addrs {
		if tried >= probeMaxAddrs {
			break
		}
		tried++
		dctx, dcancel := context.WithTimeout(context.Background(), probeDialBudget)
		start := time.Now()
		conn, derr := dialer.DialContext(dctx, "tcp", net.JoinHostPort(addr.IP.String(), port))
		elapsed := time.Since(start)
		dcancel()
		if derr == nil {
			conn.Close()
			return elapsed.Milliseconds(), true
		}
	}
	return -1, true
}