init xboard node

2026-03-15 09:27:24 +08:00
commit 37e6c016cb
40 changed files with 11011 additions and 0 deletions
@@ -0,0 +1,275 @@
+package tracker
+
+import (
+	"log/slog"
+
+	"github.com/cedar2025/xboard-node/internal/kernel"
+)
+
+const (
+	// DefaultMaxConnState caps how many connection-state entries we maintain.
+	// Prevents unbounded memory growth from rogue / DDoS connection floods.
+	DefaultMaxConnState = 100_000
+
+	// DefaultMaxAliveIPsPerUser caps the per-user IP set tracked per push interval.
+	// 200 IPs per user is generous for any legitimate use case.
+	DefaultMaxAliveIPsPerUser = 200
+)
+
+// Tracker tracks per-connection traffic deltas and accumulates per-user totals.
+//
+// Thread safety model: Tracker is designed for SINGLE-GOROUTINE use only.
+// All methods (Process, Flush*, Restore*, CurrentOnline, speed accessors) are
+// called exclusively from the service goroutine's select loop. No internal
+// locking is provided. Do NOT call Tracker methods from multiple goroutines.
+//
+// Hot-path optimisation (②): activeConns is allocated once and cleared in-place
+// every Process() call instead of being reallocated, saving GC pressure at high
+// connection counts.
+//
+// OOM protection (④): connState and aliveIPs are capped at configurable maximums.
+type Tracker struct {
+	connState map[string][2]int64     // conn_id → last [upload, download]
+	traffic   map[int][2]int64        // user_id → accumulated [upload, download]
+	aliveIPs  map[int]map[string]bool // user_id → set of source IPs
+	online    map[int]int             // user_id → current active connection count
+
+	// ② Reused buffer: cleared each Process() call, not reallocated.
+	activeConns map[string]struct{}
+
+	// ④ Limits (set at construction, never mutated).
+	maxConnState       int
+	maxAliveIPsPerUser int
+
+	// lastActiveConns stores the last observed number of active connections.
+	// This is updated on each Process() call and read by the service when
+	// building node metrics for panel reporting.
+	lastActiveConns int
+	totalConns      int64 // cumulative number of connections seen since startup
+
+	// lastUserCount is used to pre-size intervalTraffic each cycle, avoiding
+	// repeated map rehashing under heavy user counts.
+	lastUserCount int
+
+	inSpeed  int64
+	outSpeed int64
+}
+
+func New() *Tracker {
+	return NewWithLimits(DefaultMaxConnState, DefaultMaxAliveIPsPerUser)
+}
+
+func NewWithLimits(maxConnState, maxAliveIPsPerUser int) *Tracker {
+	return &Tracker{
+		connState:          make(map[string][2]int64),
+		traffic:            make(map[int][2]int64),
+		aliveIPs:           make(map[int]map[string]bool),
+		online:             make(map[int]int),
+		activeConns:        make(map[string]struct{}),
+		maxConnState:       maxConnState,
+		maxAliveIPsPerUser: maxAliveIPsPerUser,
+	}
+}
+
+// Process calculates per-user traffic deltas for this interval.
+// Internally accumulates totals for panel push and tracks alive IPs.
+// Returns per-user interval traffic [upload, download] for the limiter.
+func (t *Tracker) Process(conns []kernel.Connection) map[int][2]int64 {
+	// Pre-size using last cycle's user count to avoid rehashing under heavy load.
+	intervalTraffic := make(map[int][2]int64, t.lastUserCount)
+
+	// Clear activeConns in-place (no alloc, backing array retained).
+	for k := range t.activeConns {
+		delete(t.activeConns, k)
+	}
+
+	// Rebuild aliveIPs each cycle so only currently-active IPs are reported.
+	// Previously, IPs accumulated across cycles leading to stale entries.
+	for uid := range t.aliveIPs {
+		for ip := range t.aliveIPs[uid] {
+			delete(t.aliveIPs[uid], ip)
+		}
+	}
+
+	// Clear online counts
+	for uid := range t.online {
+		delete(t.online, uid)
+	}
+
+	for _, conn := range conns {
+		t.activeConns[conn.ID] = struct{}{}
+
+		if conn.UserID == 0 {
+			continue
+		}
+
+		t.online[conn.UserID]++
+
+		var deltaUp, deltaDown int64
+		prev, exists := t.connState[conn.ID]
+		if exists {
+			deltaUp = conn.Upload - prev[0]
+			deltaDown = conn.Download - prev[1]
+			if deltaUp < 0 {
+				deltaUp = conn.Upload
+			}
+			if deltaDown < 0 {
+				deltaDown = conn.Download
+			}
+		} else {
+			// ④ New connection: only track if below the state cap.
+			if len(t.connState) >= t.maxConnState {
+				slog.Debug("connState cap reached, skipping new connection",
+					"cap", t.maxConnState, "conn_id", conn.ID)
+				continue
+			}
+			deltaUp = conn.Upload
+			deltaDown = conn.Download
+			t.totalConns++
+		}
+		t.connState[conn.ID] = [2]int64{conn.Upload, conn.Download}
+
+		if deltaUp > 0 || deltaDown > 0 {
+			cur := t.traffic[conn.UserID]
+			cur[0] += deltaUp
+			cur[1] += deltaDown
+			t.traffic[conn.UserID] = cur
+
+			it := intervalTraffic[conn.UserID]
+			it[0] += deltaUp
+			it[1] += deltaDown
+			intervalTraffic[conn.UserID] = it
+		}
+
+		// ④ Track alive IPs, capped per user.
+		if conn.SourceIP != "" {
+			ips := t.aliveIPs[conn.UserID]
+			if ips == nil {
+				ips = make(map[string]bool)
+				t.aliveIPs[conn.UserID] = ips
+			}
+			if len(ips) < t.maxAliveIPsPerUser {
+				ips[conn.SourceIP] = true
+			}
+		}
+	}
+
+	// Cleanup stale connState entries for connections that are no longer active.
+	for id := range t.connState {
+		if _, active := t.activeConns[id]; !active {
+			delete(t.connState, id)
+		}
+	}
+
+	// Snapshot current active connection count for metrics reporting.
+	t.lastActiveConns = len(t.connState)
+
+	// Derive speed directly from intervalTraffic deltas accumulated this cycle.
+	// This avoids the negative-diff bug that occurs when connState entries for
+	// closed connections are removed, which made cumulative totals drop and
+	// clamped speed to 0 for one cycle after a mass disconnect.
+	var cycleIn, cycleOut int64
+	for _, it := range intervalTraffic {
+		cycleOut += it[0]
+		cycleIn += it[1]
+	}
+	t.inSpeed = cycleIn
+	t.outSpeed = cycleOut
+
+	// Update hint for next cycle's map pre-sizing.
+	if n := len(intervalTraffic); n > 0 {
+		t.lastUserCount = n
+	}
+
+	return intervalTraffic
+}
+
+// FlushTraffic returns accumulated per-user traffic and resets the counter.
+func (t *Tracker) FlushTraffic() map[int][2]int64 {
+	data := t.traffic
+	t.traffic = make(map[int][2]int64, len(data))
+	return data
+}
+
+// RestoreTraffic adds traffic back (used when push to panel fails).
+func (t *Tracker) RestoreTraffic(data map[int][2]int64) {
+	for uid, d := range data {
+		cur := t.traffic[uid]
+		cur[0] += d[0]
+		cur[1] += d[1]
+		t.traffic[uid] = cur
+	}
+}
+
+// HasTraffic returns true if there is accumulated traffic to report.
+func (t *Tracker) HasTraffic() bool {
+	return len(t.traffic) > 0
+}
+
+// FlushAliveIPs returns per-user alive IPs and resets the tracker.
+func (t *Tracker) FlushAliveIPs() map[int][]string {
+	data := make(map[int][]string, len(t.aliveIPs))
+	for uid, ips := range t.aliveIPs {
+		ipList := make([]string, 0, len(ips))
+		for ip := range ips {
+			ipList = append(ipList, ip)
+		}
+		data[uid] = ipList
+	}
+	t.aliveIPs = make(map[int]map[string]bool, len(t.aliveIPs))
+	return data
+}
+
+// CurrentOnline returns a map of user_id to active connection count.
+func (t *Tracker) CurrentOnline() map[int]int {
+	return t.online
+}
+
+// RestoreAliveIPs merges alive IPs back in (used when push to panel fails).
+func (t *Tracker) RestoreAliveIPs(data map[int][]string) {
+	for uid, ipList := range data {
+		ips := t.aliveIPs[uid]
+		if ips == nil {
+			ips = make(map[string]bool, len(ipList))
+			t.aliveIPs[uid] = ips
+		}
+		for _, ip := range ipList {
+			if len(ips) < t.maxAliveIPsPerUser {
+				ips[ip] = true
+			}
+		}
+	}
+}
+
+// LogStats logs current tracking statistics.
+func (t *Tracker) LogStats() {
+	slog.Debug("tracker stats",
+		"active_connections", len(t.connState),
+		"users_with_traffic", len(t.traffic),
+		"users_with_ips", len(t.aliveIPs),
+	)
+}
+
+// ActiveConnections returns the last observed active connection count.
+// This is cheap and safe because Tracker is only mutated from the service
+// goroutine; reads happen from the same goroutine when pushing reports.
+func (t *Tracker) ActiveConnections() int {
+	return t.lastActiveConns
+}
+
+// TotalConnections returns the cumulative connection count since startup.
+func (t *Tracker) TotalConnections() int64 {
+	return t.totalConns
+}
+
+// InboundSpeed returns the last observed inbound (download) speed in bytes/second.
+// Process() is called every 10 seconds, so we divide the 10-second total by 10.
+func (t *Tracker) InboundSpeed() int64 {
+	return t.inSpeed / 10
+}
+
+// OutboundSpeed returns the last observed outbound (upload) speed in bytes/second.
+// Process() is called every 10 seconds, so we divide the 10-second total by 10.
+func (t *Tracker) OutboundSpeed() int64 {
+	return t.outSpeed / 10
+}
@@ -0,0 +1,278 @@
+package tracker
+
+import (
+	"testing"
+
+	"github.com/cedar2025/xboard-node/internal/kernel"
+)
+
+func TestProcess_InitialTraffic(t *testing.T) {
+	tr := New()
+	conns := []kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200, SourceIP: "1.1.1.1"},
+		{ID: "c2", UserID: 2, Upload: 50, Download: 80, SourceIP: "2.2.2.2"},
+	}
+
+	interval := tr.Process(conns)
+
+	if interval[1] != [2]int64{100, 200} {
+		t.Errorf("user 1 interval: got %v", interval[1])
+	}
+	if interval[2] != [2]int64{50, 80} {
+		t.Errorf("user 2 interval: got %v", interval[2])
+	}
+}
+
+func TestProcess_DeltaCalculation(t *testing.T) {
+	tr := New()
+
+	// First tick
+	conns := []kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200},
+	}
+	tr.Process(conns)
+
+	// Second tick — counters advanced
+	conns[0].Upload = 300
+	conns[0].Download = 500
+	interval := tr.Process(conns)
+
+	if interval[1] != [2]int64{200, 300} {
+		t.Errorf("delta: got %v, want [200,300]", interval[1])
+	}
+}
+
+func TestProcess_CounterReset(t *testing.T) {
+	tr := New()
+
+	// First tick with high counters
+	conns := []kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 1000, Download: 2000},
+	}
+	tr.Process(conns)
+
+	// Counter reset (new value < old value) — treat as fresh
+	conns[0].Upload = 50
+	conns[0].Download = 80
+	interval := tr.Process(conns)
+
+	if interval[1] != [2]int64{50, 80} {
+		t.Errorf("counter reset: got %v, want [50,80]", interval[1])
+	}
+}
+
+func TestProcess_ConnectionCleanup(t *testing.T) {
+	tr := New()
+
+	// Tick 1: connection exists
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200},
+	})
+
+	// Tick 2: connection gone — should clean up state
+	tr.Process([]kernel.Connection{})
+
+	// Tick 3: same conn ID reappears — should count full bytes (not delta)
+	interval := tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 50, Download: 30},
+	})
+
+	if interval[1] != [2]int64{50, 30} {
+		t.Errorf("after cleanup: got %v, want [50,30]", interval[1])
+	}
+}
+
+func TestProcess_ZeroUserID(t *testing.T) {
+	tr := New()
+	interval := tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 0, Upload: 100, Download: 200},
+	})
+
+	if len(interval) != 0 {
+		t.Errorf("expected no interval traffic for userID=0, got %v", interval)
+	}
+}
+
+func TestProcess_MultipleConnectionsSameUser(t *testing.T) {
+	tr := New()
+	conns := []kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200},
+		{ID: "c2", UserID: 1, Upload: 50, Download: 80},
+	}
+
+	interval := tr.Process(conns)
+
+	if interval[1] != [2]int64{150, 280} {
+		t.Errorf("aggregated: got %v, want [150,280]", interval[1])
+	}
+}
+
+func TestFlushTraffic(t *testing.T) {
+	tr := New()
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200},
+		{ID: "c2", UserID: 2, Upload: 50, Download: 80},
+	})
+
+	flushed := tr.FlushTraffic()
+	if flushed[1] != [2]int64{100, 200} {
+		t.Errorf("flushed user 1: got %v", flushed[1])
+	}
+	if flushed[2] != [2]int64{50, 80} {
+		t.Errorf("flushed user 2: got %v", flushed[2])
+	}
+
+	// After flush, should be empty
+	if tr.HasTraffic() {
+		t.Error("expected no traffic after flush")
+	}
+	flushed2 := tr.FlushTraffic()
+	if len(flushed2) != 0 {
+		t.Errorf("expected empty after second flush, got %v", flushed2)
+	}
+}
+
+func TestRestoreTraffic(t *testing.T) {
+	tr := New()
+
+	// Generate some traffic
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200},
+	})
+
+	flushed := tr.FlushTraffic()
+
+	// Restore (simulating failed push)
+	tr.RestoreTraffic(flushed)
+
+	if !tr.HasTraffic() {
+		t.Error("expected traffic after restore")
+	}
+
+	restored := tr.FlushTraffic()
+	if restored[1] != [2]int64{100, 200} {
+		t.Errorf("restored: got %v", restored[1])
+	}
+}
+
+func TestRestoreTraffic_Additive(t *testing.T) {
+	tr := New()
+
+	// Generate traffic
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200},
+	})
+
+	// Generate more traffic
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 200, Download: 400},
+	})
+
+	// Flush and restore
+	first := tr.FlushTraffic()
+	tr.RestoreTraffic(first)
+
+	// Generate additional traffic
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 250, Download: 500},
+	})
+
+	final := tr.FlushTraffic()
+	// Should be restored(200,400) + delta(50,100) = (250,500)
+	if final[1][0] != 250 || final[1][1] != 500 {
+		t.Errorf("additive restore: got %v", final[1])
+	}
+}
+
+func TestFlushAliveIPs(t *testing.T) {
+	tr := New()
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200, SourceIP: "1.1.1.1"},
+		{ID: "c2", UserID: 1, Upload: 50, Download: 80, SourceIP: "2.2.2.2"},
+		{ID: "c3", UserID: 2, Upload: 30, Download: 40, SourceIP: "3.3.3.3"},
+	})
+
+	aliveIPs := tr.FlushAliveIPs()
+	if len(aliveIPs[1]) != 2 {
+		t.Errorf("user 1 IPs: got %d, want 2", len(aliveIPs[1]))
+	}
+	if len(aliveIPs[2]) != 1 {
+		t.Errorf("user 2 IPs: got %d, want 1", len(aliveIPs[2]))
+	}
+
+	// After flush, should be empty
+	aliveIPs2 := tr.FlushAliveIPs()
+	if len(aliveIPs2) != 0 {
+		t.Errorf("expected empty after second flush, got %v", aliveIPs2)
+	}
+}
+
+func TestFlushAliveIPs_DedupSameIP(t *testing.T) {
+	tr := New()
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200, SourceIP: "1.1.1.1"},
+		{ID: "c2", UserID: 1, Upload: 50, Download: 80, SourceIP: "1.1.1.1"},
+	})
+
+	aliveIPs := tr.FlushAliveIPs()
+	if len(aliveIPs[1]) != 1 {
+		t.Errorf("expected dedup to 1 IP, got %d", len(aliveIPs[1]))
+	}
+}
+
+func TestHasTraffic(t *testing.T) {
+	tr := New()
+	if tr.HasTraffic() {
+		t.Error("new tracker should not have traffic")
+	}
+
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200},
+	})
+	if !tr.HasTraffic() {
+		t.Error("should have traffic after process")
+	}
+
+	tr.FlushTraffic()
+	if tr.HasTraffic() {
+		t.Error("should not have traffic after flush")
+	}
+}
+
+func TestProcess_NoTrafficDelta(t *testing.T) {
+	tr := New()
+
+	conns := []kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200},
+	}
+	tr.Process(conns)
+
+	// Same values — no delta
+	interval := tr.Process(conns)
+	if len(interval) != 0 {
+		t.Errorf("expected no interval traffic for zero delta, got %v", interval)
+	}
+}
+
+func TestTrafficAccumulation(t *testing.T) {
+	tr := New()
+
+	// Tick 1
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 100, Download: 200},
+	})
+	// Tick 2
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 300, Download: 500},
+	})
+	// Tick 3
+	tr.Process([]kernel.Connection{
+		{ID: "c1", UserID: 1, Upload: 350, Download: 600},
+	})
+
+	flushed := tr.FlushTraffic()
+	// Total: 100+200+50 = 350 upload, 200+300+100 = 600 download
+	if flushed[1] != [2]int64{350, 600} {
+		t.Errorf("accumulated: got %v, want [350,600]", flushed[1])
+	}
+}