Browse Source

fix(traffic): reset due inbounds and clients concurrently

The periodic reset job reset every due inbound, then every due client, one
at a time, and each waited on its node: up to 10s per node inbound, and 4s
per attached node inbound for a client. A few hanging nodes stretched a
single run over hours.

Both loops now run eight at a time. With the per-client fan-out of four
that stays within the 32 concurrent node calls the other node fan-outs use.
Sanaei 14 hours ago
parent
commit
dea7cd9cc1

+ 36 - 9
internal/web/job/periodic_traffic_reset_job.go

@@ -1,12 +1,19 @@
 package job
 
 import (
+	"sync"
+	"sync/atomic"
 	"time"
 
 	"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"
 )
 
+// periodicResetConcurrency bounds how many inbounds or clients one run resets at once:
+// each waits on its node, so one at a time a few hanging nodes stretched a run for hours.
+const periodicResetConcurrency = 8
+
 // Period represents the time period for traffic resets.
 type Period string
 
@@ -35,6 +42,21 @@ func monthlyResetDue(resetDay int, now time.Time) bool {
 	return now.Day() == min(resetDay, lastDay)
 }
 
+func forEachResetBounded(n int, reset func(i int)) {
+	sem := make(chan struct{}, periodicResetConcurrency)
+	var wg sync.WaitGroup
+	for i := range n {
+		wg.Add(1)
+		sem <- struct{}{}
+		common.GoRecover("periodic-traffic-reset", func() {
+			defer wg.Done()
+			defer func() { <-sem }()
+			reset(i)
+		})
+	}
+	wg.Wait()
+}
+
 // Run resets traffic statistics for all inbounds that match the configured reset
 // period, then for the clients carrying that period on their own (#5497).
 func (j *PeriodicTrafficResetJob) Run() {
@@ -64,8 +86,9 @@ func (j *PeriodicTrafficResetJob) resetInboundsOnSchedule() {
 	}
 	logger.Infof("Running periodic traffic reset job for period: %s (%d matching inbounds)", j.period, len(inbounds))
 
-	resetCount := 0
-	for _, inbound := range inbounds {
+	var resetCount atomic.Int32
+	forEachResetBounded(len(inbounds), func(i int) {
+		inbound := inbounds[i]
 		resetInboundErr := j.inboundService.ResetInboundTraffic(inbound.Id)
 		if resetInboundErr != nil {
 			logger.Warning("Failed to reset traffic for inbound", inbound.Id, ":", resetInboundErr)
@@ -77,12 +100,12 @@ func (j *PeriodicTrafficResetJob) resetInboundsOnSchedule() {
 		}
 
 		if resetInboundErr == nil && resetClientErr == nil {
-			resetCount++
+			resetCount.Add(1)
 		}
-	}
+	})
 
-	if resetCount > 0 {
-		logger.Infof("Periodic traffic reset completed: %d inbounds reset", resetCount)
+	if count := resetCount.Load(); count > 0 {
+		logger.Infof("Periodic traffic reset completed: %d inbounds reset", count)
 	}
 }
 
@@ -115,19 +138,23 @@ func (j *PeriodicTrafficResetJob) resetClientsOnTheirOwnCycle() {
 	}
 	logger.Infof("Running periodic traffic reset job for period: %s (%d matching clients)", j.period, len(due))
 
+	var mu sync.Mutex
 	resetCount := 0
 	needRestart := false
-	for _, c := range due {
+	forEachResetBounded(len(due), func(i int) {
+		c := due[i]
 		// ResetTrafficByEmail rather than a bulk UPDATE: it is the path that also
 		// propagates to the client's node and clears the MTProto sidecar quota.
 		nr, resetErr := j.clientService.ResetTrafficByEmail(&j.inboundService, c.Email)
 		if resetErr != nil {
 			logger.Warning("Failed to reset traffic for client", c.Email, ":", resetErr)
-			continue
+			return
 		}
+		mu.Lock()
 		needRestart = needRestart || nr
 		resetCount++
-	}
+		mu.Unlock()
+	})
 	// Dropping this leaves a re-enabled client absent from the running core until
 	// something unrelated restarts it.
 	if needRestart {

+ 134 - 0
internal/web/job/periodic_traffic_reset_nodes_test.go

@@ -0,0 +1,134 @@
+package job
+
+import (
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strconv"
+	"strings"
+	"sync"
+	"sync/atomic"
+	"testing"
+	"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/web/runtime"
+	"github.com/mhsanaei/3x-ui/v3/internal/xray"
+)
+
+// resetGate holds every node reset open until released, counting how many nodes
+// the job reaches at once.
+type resetGate struct {
+	entered atomic.Int32
+	release chan struct{}
+	once    sync.Once
+}
+
+func (g *resetGate) open() { g.once.Do(func() { close(g.release) }) }
+
+func (g *resetGate) waitAll(t *testing.T, want int32) {
+	t.Helper()
+	deadline := time.Now().Add(3 * time.Second)
+	for g.entered.Load() < want {
+		if time.Now().After(deadline) {
+			t.Fatalf("periodic reset reached %d of %d hanging nodes, want all of them at once", g.entered.Load(), want)
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+}
+
+// resetNode is a node whose every traffic reset hangs until the gate opens.
+func resetNode(t *testing.T, gate *resetGate, name string) int {
+	t.Helper()
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		_, _ = io.Copy(io.Discard, r.Body)
+		if strings.Contains(r.URL.Path, "resetTraffic") {
+			gate.entered.Add(1)
+			select {
+			case <-r.Context().Done():
+			case <-gate.release:
+			}
+		}
+		w.Header().Set("Content-Type", "application/json")
+		_, _ = w.Write([]byte(`{"success":true}`))
+	}))
+	t.Cleanup(srv.Close)
+	host, port, _ := strings.Cut(strings.TrimPrefix(srv.URL, "http://"), ":")
+	portNum, _ := strconv.Atoi(port)
+	node := &model.Node{
+		Name: name, Scheme: "http", Address: host, Port: portNum, BasePath: "/", ApiToken: "tok",
+		Enable: true, Status: "online", AllowPrivateAddress: true, TlsVerifyMode: "verify",
+	}
+	if err := database.GetDB().Create(node).Error; err != nil {
+		t.Fatalf("create node: %v", err)
+	}
+	return node.Id
+}
+
+func runResetJobAgainstGate(t *testing.T, gate *resetGate, want int32) {
+	t.Helper()
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		NewPeriodicTrafficResetJob("daily", time.UTC).Run()
+	}()
+	t.Cleanup(func() { gate.open(); <-done })
+	gate.waitAll(t, want)
+}
+
+func newResetFleet(t *testing.T) *resetGate {
+	t.Helper()
+	initResetJobDB(t)
+	runtime.SetManager(runtime.NewManager(runtime.LocalDeps{APIPort: func() int { return 0 }, SetNeedRestart: func() {}}))
+	t.Cleanup(func() { runtime.SetManager(nil) })
+	return &resetGate{release: make(chan struct{})}
+}
+
+// The job reset due clients and inbounds one by one, each waiting on its node,
+// so a few hanging nodes stretched one run across hours.
+func TestPeriodicResetReachesClientNodesConcurrently(t *testing.T) {
+	gate := newResetFleet(t)
+	db := database.GetDB()
+	for i := range 3 {
+		nodeID := resetNode(t, gate, fmt.Sprintf("client-node-%d", i))
+		email := fmt.Sprintf("cycle-%d@node", i)
+		client := model.Client{Email: email, ID: fmt.Sprintf("00000000-0000-4000-8000-00000000000%d", i), Enable: true, TrafficReset: "daily"}
+		settings, _ := json.Marshal(map[string]any{"clients": []model.Client{client}})
+		ib := model.Inbound{
+			UserId: 1, Enable: true, Port: 47000 + i, Protocol: model.VLESS, NodeID: &nodeID,
+			Tag: "reset-client-" + strconv.Itoa(i), TrafficReset: "never", Settings: string(settings),
+		}
+		if err := db.Create(&ib).Error; err != nil {
+			t.Fatalf("create inbound: %v", err)
+		}
+		rec := model.ClientRecord{Email: email, UUID: client.ID, Enable: true, TrafficReset: "daily"}
+		if err := db.Create(&rec).Error; err != nil {
+			t.Fatalf("create client record: %v", err)
+		}
+		if err := db.Create(&model.ClientInbound{ClientId: rec.Id, InboundId: ib.Id}).Error; err != nil {
+			t.Fatalf("link client: %v", err)
+		}
+		if err := db.Create(&xray.ClientTraffic{InboundId: ib.Id, Email: email, Enable: true, Up: 500, Down: 700}).Error; err != nil {
+			t.Fatalf("create traffic: %v", err)
+		}
+	}
+	runResetJobAgainstGate(t, gate, 3)
+}
+
+func TestPeriodicResetReachesInboundNodesConcurrently(t *testing.T) {
+	gate := newResetFleet(t)
+	for i := range 3 {
+		nodeID := resetNode(t, gate, fmt.Sprintf("inbound-node-%d", i))
+		ib := model.Inbound{
+			UserId: 1, Enable: true, Port: 47100 + i, Protocol: model.VLESS, NodeID: &nodeID,
+			Tag: "reset-inbound-" + strconv.Itoa(i), TrafficReset: "daily", Settings: `{"clients":[]}`,
+		}
+		if err := database.GetDB().Create(&ib).Error; err != nil {
+			t.Fatalf("create inbound: %v", err)
+		}
+	}
+	runResetJobAgainstGate(t, gate, 3)
+}