diff --git a/cmd/test-streammanager/main-test-streammanager.go b/cmd/test-streammanager/main-test-streammanager.go index 4e6702e790..128c079ece 100644 --- a/cmd/test-streammanager/main-test-streammanager.go +++ b/cmd/test-streammanager/main-test-streammanager.go @@ -209,8 +209,9 @@ type BrokerDataSender struct { broker *streamclient.Broker } -func (s *BrokerDataSender) SendData(dataPk wshrpc.CommandStreamData) { +func (s *BrokerDataSender) SendData(dataPk wshrpc.CommandStreamData) error { s.broker.SendData(dataPk) + return nil } // MetricsWriter wraps an io.Writer and records bytes written to metrics diff --git a/pkg/jobcontroller/jobcontroller.go b/pkg/jobcontroller/jobcontroller.go index e41d77585c..c45e74232b 100644 --- a/pkg/jobcontroller/jobcontroller.go +++ b/pkg/jobcontroller/jobcontroller.go @@ -370,38 +370,44 @@ func handleRouteEvent(event *wps.WaveEvent, newStatus string) { sendBlockJobStatusEventByJob(ctx, job) if newStatus == JobConnStatus_Disconnected && job != nil && isJobManagerRunning(job) { - if shouldAttemptAutoReconnect(jobId) { - go attemptAutoReconnect(jobId, job.Connection) + if delay, ok := nextAutoReconnectDelay(jobId); ok { + go attemptAutoReconnect(jobId, job.Connection, delay) } } } } } -func shouldAttemptAutoReconnect(jobId string) bool { +// nextAutoReconnectDelay returns how long to wait before auto-reconnecting jobId, +// or false if an attempt is already scheduled. A route down inside the cooldown is +// deferred to the end of the cooldown rather than dropped, otherwise a job whose +// route goes down twice in quick succession would stay disconnected. +func nextAutoReconnectDelay(jobId string) (time.Duration, bool) { now := time.Now().Unix() + cooldownSecs := int64(AutoReconnectCooldown / time.Second) lastAttempt, exists := lastAutoReconnectAttempt.GetEx(jobId) - if !exists { + if !exists || now-lastAttempt >= cooldownSecs { lastAutoReconnectAttempt.Set(jobId, now) - return true + return AutoReconnectDelay, true } - timeSinceLastAttempt := time.Duration(now-lastAttempt) * time.Second - if timeSinceLastAttempt >= AutoReconnectCooldown { - lastAutoReconnectAttempt.Set(jobId, now) - return true + if lastAttempt > now { + // a deferred attempt is already scheduled + return 0, false } - return false + nextAttempt := lastAttempt + cooldownSecs + lastAutoReconnectAttempt.Set(jobId, nextAttempt) + return time.Duration(nextAttempt-now)*time.Second + AutoReconnectDelay, true } -func attemptAutoReconnect(jobId string, connName string) { +func attemptAutoReconnect(jobId string, connName string, delay time.Duration) { defer func() { panichandler.PanicHandler("jobcontroller:attemptAutoReconnect", recover()) }() - time.Sleep(AutoReconnectDelay) + time.Sleep(delay) isConnected, err := conncontroller.IsConnected(connName) if err != nil || !isConnected { diff --git a/pkg/jobcontroller/jobcontroller_test.go b/pkg/jobcontroller/jobcontroller_test.go new file mode 100644 index 0000000000..d25750a4b4 --- /dev/null +++ b/pkg/jobcontroller/jobcontroller_test.go @@ -0,0 +1,44 @@ +// Copyright 2025, Command Line Inc. +// SPDX-License-Identifier: Apache-2.0 + +package jobcontroller + +import ( + "testing" + "time" +) + +func TestNextAutoReconnectDelayDefersInsideCooldown(t *testing.T) { + jobId := "test-job-cooldown" + defer lastAutoReconnectAttempt.Delete(jobId) + + delay, ok := nextAutoReconnectDelay(jobId) + if !ok || delay != AutoReconnectDelay { + t.Fatalf("Expected immediate attempt, got delay=%v ok=%v", delay, ok) + } + + // a second route down inside the cooldown must be deferred, not dropped + delay, ok = nextAutoReconnectDelay(jobId) + if !ok { + t.Fatal("Expected a deferred attempt inside the cooldown") + } + if delay <= AutoReconnectDelay || delay > AutoReconnectCooldown+AutoReconnectDelay { + t.Errorf("Expected delay until the end of the cooldown, got %v", delay) + } + + // further route downs while the deferred attempt is pending schedule nothing + if _, ok = nextAutoReconnectDelay(jobId); ok { + t.Error("Expected no new attempt while a deferred attempt is pending") + } +} + +func TestNextAutoReconnectDelayAfterCooldown(t *testing.T) { + jobId := "test-job-expired" + defer lastAutoReconnectAttempt.Delete(jobId) + + lastAutoReconnectAttempt.Set(jobId, time.Now().Add(-AutoReconnectCooldown).Unix()) + delay, ok := nextAutoReconnectDelay(jobId) + if !ok || delay != AutoReconnectDelay { + t.Errorf("Expected immediate attempt after cooldown, got delay=%v ok=%v", delay, ok) + } +} diff --git a/pkg/jobmanager/jobmanager.go b/pkg/jobmanager/jobmanager.go index dd58bccc52..82bbb3d294 100644 --- a/pkg/jobmanager/jobmanager.go +++ b/pkg/jobmanager/jobmanager.go @@ -182,8 +182,8 @@ func (jm *JobManager) connectToStreamHelper_withlock(mainServerConn *MainServerC jm.connectedStreamClient = nil } dataSender := &routedDataSender{ - wshRpc: mainServerConn.WshRpc, - route: streamMeta.ReaderRouteId, + msc: mainServerConn, + route: streamMeta.ReaderRouteId, } serverSeq, err := jm.StreamManager.ClientConnected( streamMeta.Id, diff --git a/pkg/jobmanager/mainserverconn.go b/pkg/jobmanager/mainserverconn.go index 33bb10cdfb..51a3158c66 100644 --- a/pkg/jobmanager/mainserverconn.go +++ b/pkg/jobmanager/mainserverconn.go @@ -37,17 +37,22 @@ func (msc *MainServerConn) Close() { } type routedDataSender struct { - wshRpc *wshutil.WshRpc - route string + msc *MainServerConn + route string } -func (rds *routedDataSender) SendData(dataPk wshrpc.CommandStreamData) { +func (rds *routedDataSender) SendData(dataPk wshrpc.CommandStreamData) error { // log.Printf("SendData: sending seq=%d, len=%d, eof=%t, error=%s, route=%s", // dataPk.Seq, len(dataPk.Data64), dataPk.Eof, dataPk.Error, rds.route) - err := wshclient.StreamDataCommand(rds.wshRpc, dataPk, &wshrpc.RpcOpts{NoResponse: true, Route: rds.route}) + err := wshclient.StreamDataCommand(rds.msc.WshRpc, dataPk, &wshrpc.RpcOpts{NoResponse: true, Route: rds.route}) if err != nil { - log.Printf("SendData: error sending stream data: %v\n", err) + log.Printf("SendData: error sending stream data, closing connection: %v\n", err) + // close the connection so the main server sees the job route go down + // and reconnects, resuming the stream from its last received seq + rds.msc.Close() + return err } + return nil } func (msc *MainServerConn) authenticateSelfToServer(jobAuthToken string) error { diff --git a/pkg/jobmanager/streammanager.go b/pkg/jobmanager/streammanager.go index 4d77ed5acc..fde88e920c 100644 --- a/pkg/jobmanager/streammanager.go +++ b/pkg/jobmanager/streammanager.go @@ -21,7 +21,7 @@ const ( ) type DataSender interface { - SendData(dataPk wshrpc.CommandStreamData) + SendData(dataPk wshrpc.CommandStreamData) error } type streamTerminalEvent struct { @@ -170,7 +170,10 @@ func (sm *StreamManager) GetStreamDoneInfo() (done bool, streamError string) { func (sm *StreamManager) ClientDisconnected() { sm.lock.Lock() defer sm.lock.Unlock() + sm.clientDisconnected_withlock() +} +func (sm *StreamManager) clientDisconnected_withlock() { if !sm.connected { return } @@ -352,8 +355,28 @@ func (sm *StreamManager) senderLoop() { if pkt == nil { continue } - sender.SendData(*pkt) + err := sender.SendData(*pkt) + if err != nil { + sm.handleSendError(sender, pkt.Seq, err) + } + } +} + +// handleSendError drops back to DISCONNECTED (buffering) mode after a failed send. +// The failed packet is already counted in sentNotAcked and is never retransmitted, +// so staying connected would leave a gap the client can never ACK past, and the +// stream would stall once the window fills. The client resumes from its last +// received seq when it reconnects. +func (sm *StreamManager) handleSendError(sender DataSender, seq int64, err error) { + sm.lock.Lock() + defer sm.lock.Unlock() + + if !sm.connected || sm.dataSender != sender { + // a new client connected after this packet was prepared, leave it alone + return } + log.Printf("handleSendError: send failed at seq=%d, switching to disconnected mode: %v", seq, err) + sm.clientDisconnected_withlock() } func (sm *StreamManager) prepareNextPacket() (done bool, pkt *wshrpc.CommandStreamData, sender DataSender) { diff --git a/pkg/jobmanager/streammanager_test.go b/pkg/jobmanager/streammanager_test.go index 9a0e3c895e..223f1d8acd 100644 --- a/pkg/jobmanager/streammanager_test.go +++ b/pkg/jobmanager/streammanager_test.go @@ -5,6 +5,7 @@ package jobmanager import ( "encoding/base64" + "errors" "io" "strings" "sync" @@ -19,10 +20,11 @@ type testWriter struct { packets []wshrpc.CommandStreamData } -func (tw *testWriter) SendData(pkt wshrpc.CommandStreamData) { +func (tw *testWriter) SendData(pkt wshrpc.CommandStreamData) error { tw.mu.Lock() defer tw.mu.Unlock() tw.packets = append(tw.packets, pkt) + return nil } func (tw *testWriter) GetPackets() []wshrpc.CommandStreamData { @@ -328,6 +330,106 @@ func TestTerminalEventOrdering(t *testing.T) { sm.Close() } +type failingWriter struct { + mu sync.Mutex + calls int +} + +func (fw *failingWriter) SendData(pkt wshrpc.CommandStreamData) error { + fw.mu.Lock() + defer fw.mu.Unlock() + fw.calls++ + return errors.New("timeout sending request") +} + +func (fw *failingWriter) GetCalls() int { + fw.mu.Lock() + defer fw.mu.Unlock() + return fw.calls +} + +func TestSendErrorFallsBackToDisconnected(t *testing.T) { + sm := MakeStreamManager() + + data := strings.Repeat("a", MaxPacketSize) + strings.Repeat("b", MaxPacketSize) + err := sm.AttachReader(strings.NewReader(data)) + if err != nil { + t.Fatalf("AttachReader failed: %v", err) + } + + fw := &failingWriter{} + _, err = sm.ClientConnected("1", fw, CwndSize, 0) + if err != nil { + t.Fatalf("ClientConnected failed: %v", err) + } + + time.Sleep(100 * time.Millisecond) + + if calls := fw.GetCalls(); calls != 1 { + t.Errorf("Expected sending to stop after the first failed send, got %d calls", calls) + } + + // the failed packet must not be counted as in flight, a new client should get all data from seq 0 + tw := &testWriter{} + _, err = sm.ClientConnected("2", tw, CwndSize, 0) + if err != nil { + t.Fatalf("ClientConnected after send error failed: %v", err) + } + + time.Sleep(100 * time.Millisecond) + + packets := tw.GetPackets() + if len(packets) == 0 || packets[0].Seq != 0 { + t.Fatalf("Expected resend starting at seq 0, got %d packets", len(packets)) + } + allData := "" + for _, pkt := range packets { + if pkt.Data64 != "" { + allData += decodeData(pkt.Data64) + } + } + if allData != data { + t.Errorf("Expected %d bytes after reconnect, got %d", len(data), len(allData)) + } + + sm.Close() +} + +func TestStaleSendErrorKeepsNewClient(t *testing.T) { + sm := MakeStreamManager() + + err := sm.AttachReader(strings.NewReader("hello")) + if err != nil { + t.Fatalf("AttachReader failed: %v", err) + } + + oldWriter := &testWriter{} + _, err = sm.ClientConnected("1", oldWriter, CwndSize, 0) + if err != nil { + t.Fatalf("ClientConnected failed: %v", err) + } + time.Sleep(50 * time.Millisecond) + sm.ClientDisconnected() + + newWriter := &testWriter{} + _, err = sm.ClientConnected("2", newWriter, CwndSize, 0) + if err != nil { + t.Fatalf("ClientConnected failed: %v", err) + } + + // a late error from the old client's sender must not disconnect the new client + sm.handleSendError(oldWriter, 0, errors.New("timeout sending request")) + + sm.lock.Lock() + connected := sm.connected + sm.lock.Unlock() + if !connected { + t.Error("Expected new client to stay connected after a stale send error") + } + + sm.Close() +} + type slowReader struct { data []byte pos int diff --git a/pkg/wshutil/wshrpc.go b/pkg/wshutil/wshrpc.go index 8a29e204bd..48d6c5634a 100644 --- a/pkg/wshutil/wshrpc.go +++ b/pkg/wshutil/wshrpc.go @@ -772,7 +772,12 @@ func (w *WshRpc) SendComplexRequest(command string, data any, opts *wshrpc.RpcOp timeoutMs = DefaultTimeoutMs } defer func() { - panichandler.PanicHandler("SendComplexRequest", recover()) + // sending on OutputCh panics if the connection closed while we were blocked on it. + // callers dereference the handler unless an error is returned + if panicErr := panichandler.PanicHandler("SendComplexRequest", recover()); panicErr != nil { + rtnHandler = nil + rtnErr = panicErr + } }() if command == "" { return nil, fmt.Errorf("command cannot be empty") diff --git a/pkg/wshutil/wshrpc_test.go b/pkg/wshutil/wshrpc_test.go new file mode 100644 index 0000000000..f9d204bbb5 --- /dev/null +++ b/pkg/wshutil/wshrpc_test.go @@ -0,0 +1,26 @@ +// Copyright 2025, Command Line Inc. +// SPDX-License-Identifier: Apache-2.0 + +package wshutil + +import ( + "testing" + + "github.com/wavetermdev/waveterm/pkg/wshrpc" +) + +func TestSendOnClosedOutputReturnsError(t *testing.T) { + outputCh := make(chan []byte) + w := MakeWshRpcWithChannels(nil, outputCh, wshrpc.RpcContext{}, nil, "test") + // same state a sender sees when the connection closes while it is blocked on OutputCh + close(outputCh) + + err := w.SendCommand("test", nil, nil) + if err == nil { + t.Error("Expected an error when sending on a closed connection") + } + _, err = w.SendRpcRequest("test", nil, nil) + if err == nil { + t.Error("Expected an error when sending on a closed connection") + } +}