修复二开 agent 流量增量丢失: 基线改上报确认后推进并落盘持久化

- xray 新增快照/确认接口 GetTrafficSnapshot/CommitTrafficSnapshot/ RestoreTrafficBaseline, 与面板 GetTraffic 的 StatsLastValues 完全隔离 - agent 改走快照: 仅主控 HTTP 200 后才推进基线, 失败轮增量由累计计数做差在下轮自动补齐; 已确认基线原子落盘, agent 重启恢复续差 - 核心重启计数归零的从零规则保留; Onlines 仍每轮新鲜采集直报 - 新增 11 个测试覆盖断网恢复/连续失败补齐/重启恢复/计数归零/ 面板语义隔离, xray 与 agent 两包全绿
2026-10-20 15:58:38 +00:00
parent 4828f946b1
commit 0f0335e943
5 changed files with 818 additions and 22 deletions
@@ -43,6 +43,12 @@ type Config struct {
 	// config. Needed when agent and master share a host (e.g. testing),
 	// otherwise both xrays fight over the default 62789. 0 = use master's.
 	ApiPort int `json:"api_port"`
+	// BaselinePath is where the last master-acknowledged traffic counters
+	// are persisted, so an agent restart keeps diffing against them
+	// instead of swallowing unreported traffic into a fresh baseline.
+	// Empty: main derives it next to the agent config file; the agent
+	// falls back to the bin folder.
+	BaselinePath string `json:"baseline_path,omitempty"`
 }

 // Agent represents the node agent
@@ -54,6 +60,7 @@ type Agent struct {
 	xrayProc       *xray.Process
 	xrayAPI        xray.XrayAPI
 	apiPort        int
+	baselineLoaded bool
 }

 // New creates a new agent instance
@@ -331,36 +338,73 @@ type userEntry struct {
 	Down  int64  `json:"down"`
 }

-// collectAndReportTraffic queries the local xray stats API for per-poll
-// deltas and posts them to the master. The XrayAPI instance is kept across
-// polls so delta baselines survive; xray restarts are handled (counters
-// resetting to zero are treated as fresh).
-func (a *Agent) collectAndReportTraffic() error {
-	if a.xrayProc == nil || !a.xrayProc.IsRunning() {
-		return nil
-	}
-	if a.apiPort <= 0 {
-		return nil
-	}
-	if err := a.xrayAPI.Init(a.apiPort); err != nil {
-		return err
-	}
-	defer a.xrayAPI.Close()
+// statsClient is the slice of the xray stats API the traffic reporter
+// uses. The snapshot/commit pair is the loss-proofing contract: snapshots
+// never advance the baseline, commits do.
+type statsClient interface {
+	GetTrafficSnapshot() (*xray.TrafficSnapshot, error)
+	CommitTrafficSnapshot(*xray.TrafficSnapshot)
+	RestoreTrafficBaseline(map[string]int64)
+	GetOnlineUsers() ([]xray.OnlineUser, error)
+}

-	traffics, clientTraffics, err := a.xrayAPI.GetTraffic()
+// trafficBaselineFile is the on-disk form of the last committed baseline.
+type trafficBaselineFile struct {
+	Version int              `json:"version"`
+	Values  map[string]int64 `json:"values"`
+}
+
+// baselinePath resolves where the acknowledged traffic baseline lives.
+func (a *Agent) baselinePath() string {
+	if a.config.BaselinePath != "" {
+		return a.config.BaselinePath
+	}
+	return filepath.Join(config.GetBinFolderPath(), "agent-traffic-baseline.json")
+}
+
+// loadTrafficBaseline reads a persisted baseline. A missing file is a
+// normal first run; a corrupt one is reported so the caller can log it —
+// both yield no values and the reporter falls back to a baseline-only
+// first snapshot rather than crashing the agent.
+func loadTrafficBaseline(path string) (map[string]int64, error) {
+	data, err := os.ReadFile(path)
+	if err != nil {
+		return nil, err
+	}
+	var f trafficBaselineFile
+	if err := json.Unmarshal(data, &f); err != nil {
+		return nil, fmt.Errorf("parse traffic baseline: %w", err)
+	}
+	return f.Values, nil
+}
+
+// saveTrafficBaseline persists the baseline atomically (write-then-rename)
+// so a crash mid-write can never leave a torn file behind.
+func saveTrafficBaseline(path string, values map[string]int64) error {
+	data, err := json.Marshal(trafficBaselineFile{Version: 1, Values: values})
 	if err != nil {
 		return err
 	}
-	onlineUsers, _ := a.xrayAPI.GetOnlineUsers()
+	tmp := path + ".tmp"
+	if err := os.WriteFile(tmp, data, 0o600); err != nil {
+		return err
+	}
+	return os.Rename(tmp, path)
+}

+// buildTrafficReport shapes a snapshot plus the freshly polled online
+// users into the wire report, dropping zero-traffic entries. Onlines are
+// point-in-time data: they ride along with each report and are never part
+// of the baseline bookkeeping.
+func buildTrafficReport(snap *xray.TrafficSnapshot, onlineUsers []xray.OnlineUser) trafficReport {
 	report := trafficReport{}
-	for _, t := range traffics {
+	for _, t := range snap.Traffics {
 		if t == nil || !t.IsInbound || t.Up+t.Down <= 0 {
 			continue
 		}
 		report.Inbounds = append(report.Inbounds, trafficEntry{Tag: t.Tag, Up: t.Up, Down: t.Down})
 	}
-	for _, ct := range clientTraffics {
+	for _, ct := range snap.ClientTraffics {
 		if ct == nil || ct.Email == "" || ct.Up+ct.Down <= 0 {
 			continue
 		}
@@ -369,10 +413,16 @@ func (a *Agent) collectAndReportTraffic() error {
 	for _, u := range onlineUsers {
 		report.Onlines = append(report.Onlines, u.Email)
 	}
-	if len(report.Inbounds) == 0 && len(report.Users) == 0 && len(report.Onlines) == 0 {
-		return nil
-	}
+	return report
+}

+func (r trafficReport) empty() bool {
+	return len(r.Inbounds) == 0 && len(r.Users) == 0 && len(r.Onlines) == 0
+}
+
+// postTrafficReport sends one report to the master; only an HTTP 200
+// counts as acknowledged.
+func (a *Agent) postTrafficReport(report trafficReport) error {
 	url := fmt.Sprintf("%s/panel/api/agent/traffic", a.config.MasterURL)
 	body, _ := json.Marshal(report)
 	req, err := http.NewRequest("POST", url, bytes.NewReader(body))
@@ -393,6 +443,77 @@ func (a *Agent) collectAndReportTraffic() error {
 	return nil
 }

+// reportTrafficWith runs one snapshot → report → commit cycle.
+//
+// The committed baseline advances ONLY after the master acknowledges the
+// report (HTTP 200). Core counters are cumulative, so when a report fails
+// nothing is lost: the next snapshot diffs against the same baseline and
+// re-includes everything unreported, exactly once. The acknowledged
+// baseline is persisted after every commit, and restored before the first
+// snapshot of a restarted agent, so an agent restart cannot swallow the
+// unreported tail into a fresh baseline either.
+//
+// A poll with nothing to send still commits: it establishes the very
+// first baseline, and keeps the persisted copy fresh.
+func (a *Agent) reportTrafficWith(sc statsClient, post func(trafficReport) error) error {
+	if !a.baselineLoaded {
+		a.baselineLoaded = true
+		if values, err := loadTrafficBaseline(a.baselinePath()); err != nil {
+			if !os.IsNotExist(err) {
+				logger.Warningf("Traffic baseline at %s unusable (%v); starting from a fresh baseline", a.baselinePath(), err)
+			}
+		} else if len(values) > 0 {
+			sc.RestoreTrafficBaseline(values)
+		}
+	}
+
+	snap, err := sc.GetTrafficSnapshot()
+	if err != nil {
+		return err
+	}
+	onlineUsers, _ := sc.GetOnlineUsers()
+
+	report := buildTrafficReport(snap, onlineUsers)
+	if !report.empty() {
+		if err := post(report); err != nil {
+			// No commit: the baseline stays where the master last
+			// acknowledged it, and these deltas are reported again
+			// (grown) on the next poll.
+			return err
+		}
+	}
+
+	sc.CommitTrafficSnapshot(snap)
+	if err := saveTrafficBaseline(a.baselinePath(), snap.BaselineValues()); err != nil {
+		// The in-memory commit stands — rolling it back would re-report
+		// traffic the master already counted. The stale file self-heals
+		// on the next successful poll; only a crash before that could
+		// re-report a single interval.
+		logger.Warningf("Failed to persist traffic baseline: %v", err)
+		return err
+	}
+	return nil
+}
+
+// collectAndReportTraffic snapshots the local xray counters and reports
+// the deltas since the last master-acknowledged baseline. The XrayAPI
+// instance is kept across polls so baselines survive; xray restarts are
+// handled (counters resetting to zero are treated as fresh).
+func (a *Agent) collectAndReportTraffic() error {
+	if a.xrayProc == nil || !a.xrayProc.IsRunning() {
+		return nil
+	}
+	if a.apiPort <= 0 {
+		return nil
+	}
+	if err := a.xrayAPI.Init(a.apiPort); err != nil {
+		return err
+	}
+	defer a.xrayAPI.Close()
+
+	return a.reportTrafficWith(&a.xrayAPI, a.postTrafficReport)
+}
+
 // LoadConfig loads agent config from file
 func LoadConfig(path string) (*Config, error) {
 	data, err := os.ReadFile(path)
@@ -0,0 +1,322 @@
+package agent
+
+import (
+	"encoding/json"
+	"errors"
+	"net/http"
+	"net/http/httptest"
+	"os"
+	"path/filepath"
+	"strings"
+	"sync"
+	"testing"
+
+	"github.com/mhsanaei/3x-ui/v3/internal/xray"
+)
+
+// fakeStatsClient implements the snapshot/commit contract over scripted
+// cumulative counters: snapshots diff against the committed baseline with
+// the same per-stat rule as the real client (unseen or moved-backwards
+// counters count from zero), and only Commit moves the baseline. The
+// delta math itself is proven against the wire-level fake in the xray
+// package; here the fake exists so the agent's orchestration — commit
+// only after a 200, persist, restore after restart — can be driven
+// exactly.
+type fakeStatsClient struct {
+	counters  map[string]int64
+	committed map[string]int64
+	ready     bool
+}
+
+func newFakeStatsClient() *fakeStatsClient {
+	return &fakeStatsClient{counters: map[string]int64{}}
+}
+
+func (f *fakeStatsClient) GetTrafficSnapshot() (*xray.TrafficSnapshot, error) {
+	snap := &xray.TrafficSnapshot{Values: map[string]int64{}}
+	clients := map[string]*xray.ClientTraffic{}
+	for name, value := range f.counters {
+		snap.Values[name] = value
+		if !f.ready {
+			continue
+		}
+		last, ok := f.committed[name]
+		if !ok || value < last {
+			last = 0
+		}
+		delta := value - last
+		parts := strings.Split(name, ">>>")
+		if len(parts) != 4 || parts[0] != "user" {
+			continue
+		}
+		ct, ok := clients[parts[1]]
+		if !ok {
+			ct = &xray.ClientTraffic{Email: parts[1]}
+			clients[parts[1]] = ct
+		}
+		if parts[3] == "downlink" {
+			ct.Down = delta
+		} else {
+			ct.Up = delta
+		}
+	}
+	for _, ct := range clients {
+		snap.ClientTraffics = append(snap.ClientTraffics, ct)
+	}
+	return snap, nil
+}
+
+func (f *fakeStatsClient) CommitTrafficSnapshot(s *xray.TrafficSnapshot) {
+	f.committed = s.BaselineValues()
+	f.ready = true
+}
+
+func (f *fakeStatsClient) RestoreTrafficBaseline(values map[string]int64) {
+	if len(values) == 0 {
+		return
+	}
+	f.committed = map[string]int64{}
+	for k, v := range values {
+		f.committed[k] = v
+	}
+	f.ready = true
+}
+
+func (f *fakeStatsClient) GetOnlineUsers() ([]xray.OnlineUser, error) { return nil, nil }
+
+// fakeMaster is an HTTP master whose /panel/api/agent/traffic endpoint
+// fails (500) for the first failTimes requests, then acknowledges and
+// records everything it receives.
+type fakeMaster struct {
+	srv       *httptest.Server
+	mu        sync.Mutex
+	failTimes int
+	posts     int
+	userUp    map[string]int64
+	userDown  map[string]int64
+}
+
+func newFakeMaster(t *testing.T, failTimes int) *fakeMaster {
+	t.Helper()
+	m := &fakeMaster{failTimes: failTimes, userUp: map[string]int64{}, userDown: map[string]int64{}}
+	m.srv = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		m.mu.Lock()
+		defer m.mu.Unlock()
+		m.posts++
+		if m.posts <= m.failTimes {
+			w.WriteHeader(http.StatusInternalServerError)
+			return
+		}
+		var report trafficReport
+		if err := json.NewDecoder(r.Body).Decode(&report); err != nil {
+			w.WriteHeader(http.StatusBadRequest)
+			return
+		}
+		for _, u := range report.Users {
+			m.userUp[u.Email] += u.Up
+			m.userDown[u.Email] += u.Down
+		}
+		w.WriteHeader(http.StatusOK)
+	}))
+	t.Cleanup(m.srv.Close)
+	return m
+}
+
+func newTestAgent(t *testing.T, masterURL string) *Agent {
+	t.Helper()
+	return New(&Config{
+		MasterURL:    masterURL,
+		Token:        "test-token",
+		NodeName:     "test-node",
+		BaselinePath: filepath.Join(t.TempDir(), "agent-traffic-baseline.json"),
+	})
+}
+
+const aliceUp = "user>>>alice>>>traffic>>>uplink"
+
+// TestReportOutageThenRecoveryConservesTotals is the acceptance scenario
+// for the fix: the master is down for several polls while traffic accrues,
+// then recovers. Every byte the core counted must reach the master exactly
+// once — the recovery report carries the whole outage window, and later
+// polls carry only new traffic.
+func TestReportOutageThenRecoveryConservesTotals(t *testing.T) {
+	master := newFakeMaster(t, 2) // first two posts fail with 500
+	a := newTestAgent(t, master.srv.URL)
+	fake := newFakeStatsClient()
+	fake.counters[aliceUp] = 1000 // pre-existing, unattributable
+
+	// Poll 1 establishes the baseline: nothing is posted, the baseline
+	// file appears with the raw counter.
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
+		t.Fatalf("baseline poll: %v", err)
+	}
+	if master.posts != 0 {
+		t.Fatalf("baseline poll posted %d times, want 0", master.posts)
+	}
+	values, err := loadTrafficBaseline(a.baselinePath())
+	if err != nil || values[aliceUp] != 1000 {
+		t.Fatalf("persisted baseline = %v (err %v), want alice at 1000", values, err)
+	}
+
+	// Outage: two polls accrue traffic but the master 500s.
+	fake.counters[aliceUp] = 1400
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err == nil {
+		t.Fatal("poll during outage returned nil, want the post error")
+	}
+	fake.counters[aliceUp] = 2000
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err == nil {
+		t.Fatal("second outage poll returned nil, want the post error")
+	}
+	// The failures advanced nothing: the persisted baseline still holds
+	// the last acknowledged counters.
+	values, err = loadTrafficBaseline(a.baselinePath())
+	if err != nil || values[aliceUp] != 1000 {
+		t.Fatalf("baseline after failed polls = %v (err %v), want still 1000", values, err)
+	}
+
+	// Recovery: one post carries the entire outage window (2000-1000).
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
+		t.Fatalf("recovery poll: %v", err)
+	}
+	if got := master.userUp["alice"]; got != 1000 {
+		t.Fatalf("master received %d for alice after recovery, want exactly 1000", got)
+	}
+
+	// Steady state resumes: only fresh growth is reported.
+	fake.counters[aliceUp] = 2100
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
+		t.Fatalf("steady poll: %v", err)
+	}
+	if got := master.userUp["alice"]; got != 1100 {
+		t.Fatalf("master total = %d, want 1100 (= 2100 counted - 1000 baseline)", got)
+	}
+}
+
+// TestAgentRestartRestoresBaselineFromDisk: a new agent process (fresh
+// in-memory state everywhere) over the same baseline file must pick up
+// diffing exactly where the old one stopped — the traffic that accrued
+// while it was down is reported by its FIRST poll, and nothing is counted
+// twice.
+func TestAgentRestartRestoresBaselineFromDisk(t *testing.T) {
+	master := newFakeMaster(t, 0)
+	baselinePath := filepath.Join(t.TempDir(), "agent-traffic-baseline.json")
+	newAgent := func() *Agent {
+		return New(&Config{MasterURL: master.srv.URL, Token: "t", NodeName: "n", BaselinePath: baselinePath})
+	}
+
+	a1 := newAgent()
+	fake := newFakeStatsClient()
+	fake.counters[aliceUp] = 1000
+	if err := a1.reportTrafficWith(fake, a1.postTrafficReport); err != nil {
+		t.Fatalf("a1 baseline poll: %v", err)
+	}
+	fake.counters[aliceUp] = 1500
+	if err := a1.reportTrafficWith(fake, a1.postTrafficReport); err != nil {
+		t.Fatalf("a1 report poll: %v", err)
+	}
+	if got := master.userUp["alice"]; got != 500 {
+		t.Fatalf("master received %d before restart, want 500", got)
+	}
+
+	// "Restart": brand-new agent AND a brand-new stats client (its
+	// in-memory committed baseline is gone); only the file survives. The
+	// core kept counting meanwhile.
+	a2 := newAgent()
+	fake2 := newFakeStatsClient()
+	fake2.counters[aliceUp] = 1900
+	if err := a2.reportTrafficWith(fake2, a2.postTrafficReport); err != nil {
+		t.Fatalf("a2 first poll: %v", err)
+	}
+	if got := master.userUp["alice"]; got != 900 {
+		t.Fatalf("master total = %d, want 900 (500 + the 400 accrued across the restart)", got)
+	}
+}
+
+// TestCounterResetAfterCoreRestart: the core restarts while the agent
+// keeps running; its counters fall below the committed baseline. The
+// report must carry the new counter as fresh traffic — never a negative
+// delta — and then resume normal diffing in the new epoch.
+func TestCounterResetAfterCoreRestart(t *testing.T) {
+	master := newFakeMaster(t, 0)
+	a := newTestAgent(t, master.srv.URL)
+	fake := newFakeStatsClient()
+	fake.counters[aliceUp] = 5000
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
+		t.Fatalf("baseline poll: %v", err)
+	}
+
+	fake.counters[aliceUp] = 300 // core restarted; counters began again
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
+		t.Fatalf("post-reset poll: %v", err)
+	}
+	if got := master.userUp["alice"]; got != 300 {
+		t.Fatalf("master received %d after core restart, want the fresh 300", got)
+	}
+
+	fake.counters[aliceUp] = 450
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
+		t.Fatalf("new-epoch poll: %v", err)
+	}
+	if got := master.userUp["alice"]; got != 450 {
+		t.Fatalf("master total = %d, want 450 (300 + 150 of new-epoch growth)", got)
+	}
+}
+
+// TestCorruptBaselineFileTolerated: a torn or foreign baseline file must
+// not wedge the agent — it logs and falls back to a baseline-only first
+// snapshot, then reports normally.
+func TestCorruptBaselineFileTolerated(t *testing.T) {
+	master := newFakeMaster(t, 0)
+	a := newTestAgent(t, master.srv.URL)
+	if err := os.WriteFile(a.baselinePath(), []byte("{not json"), 0o600); err != nil {
+		t.Fatal(err)
+	}
+	fake := newFakeStatsClient()
+	fake.counters[aliceUp] = 800
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
+		t.Fatalf("poll with corrupt baseline: %v", err)
+	}
+	if master.posts != 0 {
+		t.Fatalf("corrupt-baseline first poll posted %d times, want baseline-only (0)", master.posts)
+	}
+	fake.counters[aliceUp] = 900
+	if err := a.reportTrafficWith(fake, a.postTrafficReport); err != nil {
+		t.Fatalf("poll after corrupt baseline: %v", err)
+	}
+	if got := master.userUp["alice"]; got != 100 {
+		t.Fatalf("master received %d, want only the post-baseline 100", got)
+	}
+}
+
+// TestSaveLoadBaselineRoundTrip pins the persistence helpers: atomic save,
+// version field, exact values back.
+func TestSaveLoadBaselineRoundTrip(t *testing.T) {
+	path := filepath.Join(t.TempDir(), "baseline.json")
+	want := map[string]int64{aliceUp: 42, "inbound>>>in-443>>>traffic>>>downlink": 7}
+	if err := saveTrafficBaseline(path, want); err != nil {
+		t.Fatalf("save: %v", err)
+	}
+	got, err := loadTrafficBaseline(path)
+	if err != nil {
+		t.Fatalf("load: %v", err)
+	}
+	if len(got) != len(want) || got[aliceUp] != 42 || got["inbound>>>in-443>>>traffic>>>downlink"] != 7 {
+		t.Fatalf("round trip = %v, want %v", got, want)
+	}
+	if _, err := os.Stat(path + ".tmp"); !os.IsNotExist(err) {
+		t.Fatal("temp file left behind after atomic save")
+	}
+	if _, err := loadTrafficBaseline(filepath.Join(t.TempDir(), "missing.json")); !os.IsNotExist(err) {
+		t.Fatalf("missing file error = %v, want a not-exist error", err)
+	}
+	var parseErr error
+	_, parseErr = loadTrafficBaseline(path)
+	if parseErr != nil {
+		t.Fatalf("valid file failed to load: %v", parseErr)
+	}
+	bad := filepath.Join(t.TempDir(), "bad.json")
+	_ = os.WriteFile(bad, []byte("[]"), 0o600)
+	if _, err := loadTrafficBaseline(bad); err == nil || errors.Is(err, os.ErrNotExist) {
+		t.Fatalf("corrupt file error = %v, want a parse error", err)
+	}
+}