Fail typed dispatches fast when their agent session drops

The agentexec read loop removed a dead session from the agent map but never
closed its done channel, so any dispatch already sent on that session kept
waiting for a result that could never arrive and only returned at its
operation timeout. That's what made the Proxmox cancellation replay test
flaky on the shared CI runner. The restarted action runner dropped its first
session right after registering, the replay had already been sent there,
and the server waited the full 30s while the runner sat reconnected on a
new session.

The read loop now signals the session done on exit, so in-flight typed
dispatches and commands report the disconnect straight away. The
integration test sends the same request again on the runner's next
session when that happens, which is the same durable replay it is there
to prove, still with exactly one mutation. It also keeps the runner's log
and prints it on failure so the cause of any future drop is visible.
This commit is contained in:
courtmanr@gmail.com 2026-09-28 10:31:58 +01:00
parent 5b3a47e305
commit 9e2ee3d535
4 changed files with 96 additions and 3 deletions

View file

@ -7792,6 +7792,12 @@ kept running on the Proxmox host). Three coupled guarantees:
answers from the durable receipt; it is never dropped as a duplicate,
because the server replays exact request IDs to recover receipts and a
dropped replay would leave the server waiting out the full timeout.
Inbound results are correlated to the exact session that carried the
request, so once that session's reader exits (socket drop or read error)
the server signals the session done and every dispatch still waiting on
it fails immediately with a `disconnected before ... receipt` error
instead of waiting out its operation timeout. The caller recovers the
outcome by replaying the same request ID on the runner's next session.
3. The unified agent's command client tracks in-flight
`execute_command`/`read_file` executions and durable host update,
storage-cleanup, Proxmox guest lifecycle, and container lifecycle/update
@ -7815,7 +7821,8 @@ Proofs: `internal/agentexec/server_websocket_test.go`
`TestExecuteCommand_AbandonedCommandSendsCancel`),
`internal/agentexec/server_websocket_test.go`
(`TestTypedOperations_AbandonedDispatchSendsExactlyOneCancel`,
`TestTypedOperation_TimeoutSendsCancelAndExpiredContextNeverDispatches`),
`TestTypedOperation_TimeoutSendsCancelAndExpiredContextNeverDispatches`,
`TestTypedOperation_SocketDropAfterSendUnblocksDispatch`),
`internal/hostagent/operation_receipt_websocket_integration_test.go`
(`TestRealServerActionRunnerCancellationPersistsAndReplaysProxmoxReceiptAfterReconnect`),
`internal/hostagent/command_client_test.go`

View file

@ -1580,6 +1580,10 @@ func (s *Server) readLoop(ac *agentConn) {
for _, ch := range closeChs {
close(ch)
}
// Results are correlated to this exact session, so nothing can answer
// a request it carried once its reader is gone. Release in-flight
// dispatches now instead of leaving them to wait out their timeout.
ac.signalDone()
if err := ac.conn.Close(); err != nil {
log.Debug().Err(err).Str("agent_id", agentID).Msg("Failed to close connection during read-loop cleanup")
}

View file

@ -2152,6 +2152,38 @@ func TestTypedOperation_TimeoutSendsCancelAndExpiredContextNeverDispatches(t *te
})
}
// A typed request is correlated to the session that carried it, so once that
// socket drops nothing can deliver its receipt. The dispatch must report the
// disconnect instead of waiting out the operation timeout.
func TestTypedOperation_SocketDropAfterSendUnblocksDispatch(t *testing.T) {
admission := AgentAdmission{TokenID: "typed-token", AgentID: "typed-a1", Hostname: "typed-host1", RuntimeRole: RuntimeRoleActionRunner, ActionCapability: ActionCapabilityTypedV1}
s := NewServerWithAdmissionValidator(func(token, _, _ string) (AgentAdmission, bool) {
return admission, token == admission.TokenID
}, func(AgentAdmission) bool { return true })
ts := newWSServer(t, s)
defer ts.Close()
agent := registerTypedCancelTestAgent(t, s, ts.URL)
request := ProxmoxGuestLifecyclePayload{RequestID: "pve.dropped.1", ActionID: "pve", Operation: "shutdown", GuestKind: "vm", VMID: 101, ExpectedStatus: "running", Timeout: 30}
if err := BindProxmoxGuestLifecyclePayload(&request); err != nil {
t.Fatal(err)
}
done := make(chan error, 1)
go func() { _, err := s.ExecuteProxmoxGuestLifecycle(context.Background(), "typed-a1", request); done <- err }()
if sent, ok := agent.nextMessage(3 * time.Second); !ok || sent.Type != MsgTypeProxmoxGuestLifecycle {
t.Fatalf("dispatched request = %+v", sent)
}
_ = agent.conn.Close()
select {
case err := <-done:
if err == nil || !strings.Contains(err.Error(), "disconnected before") {
t.Fatalf("dropped-session dispatch error = %v", err)
}
case <-time.After(5 * time.Second):
t.Fatal("dispatch kept waiting on a session whose socket already closed")
}
}
// Probe the scheduling window after publication but before the registration
// acknowledgement. No sleeps or repeated runs are needed to expose this order.
func TestRegistrationReservesFirstFrameBeforePublishingSession(t *testing.T) {

View file

@ -1,6 +1,7 @@
package hostagent
import (
"bytes"
"context"
"encoding/json"
"errors"
@ -253,8 +254,23 @@ func TestRealServerActionRunnerCancellationPersistsAndReplaysProxmoxReceiptAfter
}))
defer httpServer.Close()
// A runner that drops its session must come back well inside the replay
// window rather than after the production reconnect backoff.
origReconnectDelay := reconnectDelay
reconnectDelay = 50 * time.Millisecond
t.Cleanup(func() { reconnectDelay = origReconnectDelay })
stateDir := t.TempDir()
logger := zerolog.Nop()
// Keep runner logs so a failure shows why a session dropped. They are only
// printed from the cleanup below, never from runner goroutines that can
// outlive the test.
runnerLogs := &lockedLogBuffer{}
t.Cleanup(func() {
if t.Failed() {
t.Logf("action runner log:\n%s", runnerLogs.String())
}
})
logger := zerolog.New(runnerLogs).With().Timestamp().Logger()
mutationStarted := make(chan struct{})
var startOnce sync.Once
var mutationMu sync.Mutex
@ -373,7 +389,23 @@ func TestRealServerActionRunnerCancellationPersistsAndReplaysProxmoxReceiptAfter
second, cancelSecond, secondDone := startRunner(t)
defer stopRunner(t, second, cancelSecond, secondDone)
replayed, err := server.ExecuteProxmoxGuestLifecycle(context.Background(), admission.AgentID, request)
// The server publishes the session before the runner finishes activating,
// so the replay can land on a session the runner then drops. That request
// can never be answered; send the same request again on the session the
// runner reconnects with. Every attempt must replay the durable receipt.
var replayed *agentexec.ProxmoxGuestLifecycleResultPayload
var err error
for attempt := 1; ; attempt++ {
replayed, err = server.ExecuteProxmoxGuestLifecycle(context.Background(), admission.AgentID, request)
if err == nil || attempt == 3 || !strings.Contains(err.Error(), "disconnected before") {
break
}
t.Logf("replay attempt %d lost its session (%v), retrying after reconnect", attempt, err)
deadline := time.Now().Add(3 * time.Second)
for !server.IsAgentConnected(admission.AgentID) && time.Now().Before(deadline) {
time.Sleep(10 * time.Millisecond)
}
}
if err != nil || replayed == nil || !replayed.MutationStarted || replayed.MutationCompleted || replayed.Error != canceled.Error {
t.Fatalf("replayed result = %+v, err=%v", replayed, err)
}
@ -404,3 +436,21 @@ func TestRealServerActionRunnerCancellationPersistsAndReplaysProxmoxReceiptAfter
t.Fatalf("cross-agent receipt query error = %v", err)
}
}
// lockedLogBuffer collects zerolog output from concurrent runner goroutines.
type lockedLogBuffer struct {
mu sync.Mutex
buf bytes.Buffer
}
func (b *lockedLogBuffer) Write(p []byte) (int, error) {
b.mu.Lock()
defer b.mu.Unlock()
return b.buf.Write(p)
}
func (b *lockedLogBuffer) String() string {
b.mu.Lock()
defer b.mu.Unlock()
return b.buf.String()
}