fix: keep load agents connected before start

2026-09-28 00:16:51 +08:00
parent 6b8ed4a2cc
commit 622fc50418
8 changed files with 352 additions and 58 deletions
@@ -0,0 +1,88 @@
+package main
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/gorilla/websocket"
+)
+
+func TestMaintainConnectionsRestoresClosedClient(t *testing.T) {
+	var connections atomic.Int64
+	server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		conn, err := (&websocket.Upgrader{}).Upgrade(w, r, nil)
+		if err != nil {
+			return
+		}
+		connections.Add(1)
+		defer conn.Close()
+		for {
+			var packet map[string]json.RawMessage
+			if err = conn.ReadJSON(&packet); err != nil {
+				return
+			}
+			for _, raw := range packet {
+				var body struct {
+					ID string `json:"id"`
+				}
+				if json.Unmarshal(raw, &body) != nil || body.ID == "" {
+					continue
+				}
+				response := map[string]any{"ctrl": map[string]any{
+					"id": body.ID, "code": 200, "text": "ok",
+					"params": map[string]any{"user": "usrReconnectTest"},
+				}}
+				if err = conn.WriteJSON(response); err != nil {
+					return
+				}
+			}
+		}
+	}))
+	defer server.Close()
+
+	account := AccountRecord{
+		Machine: "a", Index: 1, Username: "lta0001", Token: "token",
+		P2PTopic: "p2pReconnectTest", Endpoint: "ws" + strings.TrimPrefix(server.URL, "http"),
+	}
+	cfg := &Config{Tenant: "allrich", CommandTimeout: "1s", WSPingInterval: "1h"}
+	run := &agentRun{
+		cfg: cfg, machine: "a", clients: make(map[int]*WSClient),
+		recorder: newDeliveryRecorder("pending", "a", 1), ack: NewDurationHistogram(),
+	}
+	initial, err := run.connectAccountWithRetry(context.Background(), account, "grpReconnectTest", "api-key", "")
+	if err != nil {
+		t.Fatal(err)
+	}
+	run.setClient(account.Index, initial)
+	initial.Close()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	done := make(chan struct{})
+	go func() {
+		run.maintainConnections(ctx, []AccountRecord{account}, "grpReconnectTest", "api-key", "")
+		close(done)
+	}()
+	deadline := time.Now().Add(3 * time.Second)
+	for time.Now().Before(deadline) {
+		if replacement := run.client(account.Index); replacement != initial && replacement.alive() {
+			cancel()
+			<-done
+			run.close()
+			if got := connections.Load(); got < 2 {
+				t.Fatalf("websocket connections = %d, want at least 2", got)
+			}
+			return
+		}
+		time.Sleep(10 * time.Millisecond)
+	}
+	cancel()
+	<-done
+	run.close()
+	t.Fatal("closed client was not restored before deadline")
+}