fix: stabilize flaky server API and spread gate tests
Some checks failed
CI Docker Mining Proof / Linux agent hashrate proof (push) Has been cancelled

Serialize agent WebSocket writes (ping + JSON), mutex-protect spreadGateClock,
and harden deploy-plan spread route hint setup under parallel test runs.
This commit is contained in:
AetherForge
2026-06-07 05:23:28 -07:00
parent 499d6cd8d3
commit ed9c90a420
5 changed files with 64 additions and 13 deletions

View File

@@ -92,6 +92,25 @@ func testAgentClientIP(agentID string) string {
return fmt.Sprintf("10.42.%d.%d", (sum%250)+1, ((sum/250)%250)+1)
}
func waitForHubAgents(t *testing.T, hub *WSHub, agentIDs ...string) {
t.Helper()
deadline := time.Now().Add(3 * time.Second)
for time.Now().Before(deadline) {
allConnected := true
for _, id := range agentIDs {
if !hub.isAgentConnected(id) {
allConnected = false
break
}
}
if allConnected {
return
}
time.Sleep(10 * time.Millisecond)
}
t.Fatalf("agents not connected: %v", agentIDs)
}
func connectTestAgent(t *testing.T, hub *WSHub, agentID string) *websocket.Conn {
t.Helper()
srv := httptest.NewServer(http.HandlerFunc(hub.HandleAgentWS))

View File

@@ -679,6 +679,7 @@ func TestDeployPlanIncludesSpreadRouteHint(t *testing.T) {
hub := NewWSHub(deployH.db)
connectTestAgent(t, hub, patientID)
connectTestAgent(t, hub, seedID)
waitForHubAgents(t, hub, patientID, seedID)
hub.ClearanceManager().RequestElevation(patientID, 4, "test", "test")
hub.ClearanceManager().RequestElevation(seedID, 2, "test", "test")
@@ -686,11 +687,16 @@ func TestDeployPlanIncludesSpreadRouteHint(t *testing.T) {
deployH.BindPathTracer(pathTracer)
sess := testTraceSession(2)
sess.Hops[0].AgentID = patientID
sess.Hops[0].AgentName = "PZ"
sess.Hops[0].ExternalIP = "10.9.8.7"
sess.Hops[1].AgentID = seedID
sess.Hops[1].AgentName = "Seed"
sess.Hops[1].ExternalIP = "10.9.8.9"
sess.ServiceGraph = map[string]ServiceGraphHost{
"10.9.8.20": {Host: "10.9.8.20", Subnet: "10.9.8", AgentID: seedID},
"10.9.8.20": {
Host: "10.9.8.20", Subnet: "10.9.8", AgentID: seedID,
Services: []ServiceGraphEntry{{ServiceName: "smb", Port: 445, JoinLaneCandidate: "spread_smb_unc"}},
},
}
pathTracer.mu.Lock()
pathTracer.sessions[sess.ID] = sess

View File

@@ -108,6 +108,12 @@ func (c *AgentConnection) SendJSON(v interface{}) error {
return c.Conn.WriteJSON(v)
}
func (c *AgentConnection) WriteControl(messageType int, data []byte, deadline time.Time) error {
c.mu.Lock()
defer c.mu.Unlock()
return c.Conn.WriteControl(messageType, data, deadline)
}
// DashboardConn wraps a dashboard WebSocket with its own write mutex so
// broadcastDashboard and the ping loop never race on the same connection.
type DashboardConn struct {
@@ -400,7 +406,7 @@ func (h *WSHub) runPingLoopAgent(ac *AgentConnection) {
ac.latencyMu.Lock()
ac.pingSentAt = time.Now()
ac.latencyMu.Unlock()
if err := conn.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
if err := ac.WriteControl(websocket.PingMessage, nil, time.Now().Add(10*time.Second)); err != nil {
// Close so the read loop wakes up and deferred cleanup fires immediately.
_ = conn.Close()
return
@@ -877,7 +883,7 @@ func (h *WSHub) HandleAgentWS(w http.ResponseWriter, r *http.Request) {
go h.runPingLoopAgent(ac)
}
conn.WriteJSON(Message{Type: "auth_response", Payload: mustMarshal(func() map[string]interface{} {
_ = ac.SendJSON(Message{Type: "auth_response", Payload: mustMarshal(func() map[string]interface{} {
resp := map[string]interface{}{
"success": true,
"agent_id": agentID,