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..302e8d225b 100644 --- a/pkg/jobcontroller/jobcontroller.go +++ b/pkg/jobcontroller/jobcontroller.go @@ -71,6 +71,12 @@ const JobOutputFileName = "term" const AutoReconnectDelay = 1 * time.Second const AutoReconnectCooldown = 30 * time.Second +// StreamStallTimeout is how long a job's output stream may go without any +// data/EOF before the connection is considered stalled. When tripped, the +// remote job manager is asked to terminate its (dead) client connection so a +// fresh attach + stream reconnect can proceed -- see upstream issue #3439. +const StreamStallTimeout = 45 * time.Second + type connState struct { actual bool processed bool @@ -820,7 +826,7 @@ func runOutputLoop(ctx context.Context, jobId string, streamId string, reader *s log.Printf("[job:%s] [stream:%s] output loop started", jobId, streamId) buf := make([]byte, 4096) for { - n, err := reader.Read(buf) + n, err := readWithStallTimeout(ctx, jobId, streamId, reader, buf) currentStreamId, _ := jobStreamIds.GetEx(jobId) if currentStreamId != streamId { log.Printf("[job:%s] [stream:%s] stream superseded by [stream:%s], exiting output loop", jobId, streamId, currentStreamId) @@ -833,6 +839,12 @@ func runOutputLoop(ctx context.Context, jobId string, streamId string, reader *s } } + if err == errStreamStalled { + log.Printf("[job:%s] [stream:%s] stream stalled (no data for %v), forcing job manager reconnect", jobId, streamId, StreamStallTimeout) + handleStalledStream(jobId) + break + } + if err == io.EOF { log.Printf("[job:%s] stream ended (EOF)", jobId) updateErr := wstore.DBUpdateFn(ctx, jobId, func(job *waveobj.Job) { @@ -861,6 +873,102 @@ func runOutputLoop(ctx context.Context, jobId string, streamId string, reader *s } } +var errStreamStalled = fmt.Errorf("stream stalled: no data within timeout") + +// readWithStallTimeout wraps reader.Read with a stall watchdog. The remote +// durable-session job manager can lose its RPC/stream path without the job +// route ever going down (upstream issue #3439); in that case Read blocks +// forever while the shield keeps showing "Attached". If no data/EOF arrives +// within StreamStallTimeout, errStreamStalled is returned so the caller can +// force a reconnect. +func readWithStallTimeout(ctx context.Context, jobId string, streamId string, reader *streamclient.Reader, buf []byte) (int, error) { + type readResult struct { + n int + err error + } + resultCh := make(chan readResult, 1) + go func() { + n, err := reader.Read(buf) + resultCh <- readResult{n: n, err: err} + }() + + timer := time.NewTimer(StreamStallTimeout) + defer timer.Stop() + + select { + case res := <-resultCh: + return res.n, res.err + case <-timer.C: + return 0, errStreamStalled + case <-ctx.Done(): + return 0, ctx.Err() + } +} + +// handleStalledStream recovers a job whose output stream went stale. It asks +// the remote connection to drop the job manager's dead client connection, +// waits for the job route to re-register, then restarts streaming from the +// last persisted offset. The remote shell and its scrollback survive. +func handleStalledStream(jobId string) { + defer func() { + panichandler.PanicHandler("jobcontroller:handleStalledStream", recover()) + }() + + ctx, cancelFn := context.WithTimeout(context.Background(), 30*time.Second) + defer cancelFn() + + job, err := wstore.DBMustGet[*waveobj.Job](ctx, jobId) + if err != nil { + log.Printf("[job:%s] stalled stream: failed to load job: %v", jobId, err) + return + } + + if job.JobManagerStatus != JobManagerStatus_Running { + log.Printf("[job:%s] stalled stream: job manager not running, skipping recovery", jobId) + return + } + + // prevent overlapping recovery attempts for the same job + lastAttempt, exists := lastAutoReconnectAttempt.GetEx(jobId) + now := time.Now().Unix() + if exists && time.Duration(now-lastAttempt)*time.Second < AutoReconnectCooldown { + log.Printf("[job:%s] stalled stream: recovery attempted recently, skipping", jobId) + return + } + lastAutoReconnectAttempt.Set(jobId, now) + + isConnected, err := conncontroller.IsConnected(job.Connection) + if err != nil || !isConnected { + log.Printf("[job:%s] stalled stream: connection %q is down, cannot recover", jobId, job.Connection) + return + } + + log.Printf("[job:%s] stalled stream: requesting remote disconnect of stale client", jobId) + disconnectData := wshrpc.CommandRemoteDisconnectFromJobManagerData{ + JobId: jobId, + } + rpcOpts := &wshrpc.RpcOpts{ + Route: wshutil.MakeConnectionRouteId(job.Connection), + Timeout: 5000, + } + err = wshclient.RemoteDisconnectFromJobManagerCommand(wshclient.GetBareRpcClient(), disconnectData, rpcOpts) + if err != nil { + log.Printf("[job:%s] stalled stream: remote disconnect failed: %v", jobId, err) + return + } + + // the route should drop and re-register as the job manager re-attaches; + // reconnect restarts streaming from the persisted offset + SetJobConnStatus(jobId, JobConnStatus_Disconnected) + reconnectCtx, reconnectCancel := context.WithTimeout(context.Background(), 15*time.Second) + defer reconnectCancel() + if err := ReconnectJob(reconnectCtx, jobId, nil); err != nil { + log.Printf("[job:%s] stalled stream: reconnect failed: %v", jobId, err) + } else { + log.Printf("[job:%s] stalled stream: reconnect succeeded", jobId) + } +} + func HandleCmdJobExited(ctx context.Context, jobId string, data wshrpc.CommandJobCmdExitedData) error { var updatedJob *waveobj.Job err := wstore.DBUpdateFn(ctx, jobId, func(job *waveobj.Job) { diff --git a/pkg/jobmanager/jobmanager.go b/pkg/jobmanager/jobmanager.go index dd58bccc52..e2af443f57 100644 --- a/pkg/jobmanager/jobmanager.go +++ b/pkg/jobmanager/jobmanager.go @@ -55,6 +55,9 @@ func SetupJobManager(clientId string, jobId string, publicKeyBytes []byte, jobAu WshCmdJobManager.JobAuthToken = jobAuthToken WshCmdJobManager.StreamManager = MakeStreamManager() WshCmdJobManager.InputQueue = utilds.MakeQuickReorderQueue[wshrpc.CommandJobInputData](JobInputQueueSize, JobInputQueueTimeout) + WshCmdJobManager.StreamManager.OnSendError = func(err error) { + go WshCmdJobManager.handleStreamSendError(err) + } err := wavejwt.SetPublicKey(publicKeyBytes) if err != nil { return fmt.Errorf("failed to set public key: %w", err) @@ -208,6 +211,23 @@ func (jm *JobManager) disconnectFromStreamHelper(mainServerConn *MainServerConn) jm.connectedStreamClient = nil } +// handleStreamSendError is invoked (async) when the stream manager fails to +// push data to the main server. The attached client connection is stale at +// this point: keep the job + buffered stream alive, but close the stale +// socket so the main server observes the route drop and auto-reconnects. +func (jm *JobManager) handleStreamSendError(sendErr error) { + jm.lock.Lock() + client := jm.connectedStreamClient + if client != nil { + jm.connectedStreamClient = nil + } + jm.lock.Unlock() + if client != nil { + log.Printf("handleStreamSendError: closing stale client connection after send error: %v\n", sendErr) + client.Close() + } +} + func (jm *JobManager) SetAttachedClient(msc *MainServerConn) { jm.lock.Lock() defer jm.lock.Unlock() diff --git a/pkg/jobmanager/mainserverconn.go b/pkg/jobmanager/mainserverconn.go index 33bb10cdfb..1c27970f67 100644 --- a/pkg/jobmanager/mainserverconn.go +++ b/pkg/jobmanager/mainserverconn.go @@ -41,13 +41,15 @@ type routedDataSender struct { 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}) if err != nil { log.Printf("SendData: error sending stream data: %v\n", err) + 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..bc452ed715 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 { @@ -60,6 +60,10 @@ type StreamManager struct { // terminal state - once true, stream is complete terminalEventAcked bool closed bool + + // OnSendError, if set, is called (asynchronously) when a SendData call fails. + // Used to tear down the stale client connection so a fresh attach can proceed. + OnSendError func(err error) } func MakeStreamManager() *StreamManager { @@ -352,7 +356,14 @@ func (sm *StreamManager) senderLoop() { if pkt == nil { continue } - sender.SendData(*pkt) + err := sender.SendData(*pkt) + if err != nil { + log.Printf("senderLoop: send error (seq=%d): %v -- marking client disconnected\n", pkt.Seq, err) + sm.ClientDisconnected() + if sm.OnSendError != nil { + sm.OnSendError(err) + } + } } } diff --git a/pkg/jobmanager/streammanager_test.go b/pkg/jobmanager/streammanager_test.go index 9a0e3c895e..e6b3edbf04 100644 --- a/pkg/jobmanager/streammanager_test.go +++ b/pkg/jobmanager/streammanager_test.go @@ -17,12 +17,23 @@ import ( type testWriter struct { mu sync.Mutex packets []wshrpc.CommandStreamData + sendErr error } -func (tw *testWriter) SendData(pkt wshrpc.CommandStreamData) { +func (tw *testWriter) SendData(pkt wshrpc.CommandStreamData) error { tw.mu.Lock() defer tw.mu.Unlock() + if tw.sendErr != nil { + return tw.sendErr + } tw.packets = append(tw.packets, pkt) + return nil +} + +func (tw *testWriter) SetSendError(err error) { + tw.mu.Lock() + defer tw.mu.Unlock() + tw.sendErr = err } func (tw *testWriter) GetPackets() []wshrpc.CommandStreamData { @@ -346,3 +357,94 @@ func (sr *slowReader) Read(p []byte) (n int, err error) { return n, nil } + +// failingWriter always fails SendData, to simulate a stale RPC/stream path. +type failingWriter struct { + err error +} + +func (fw *failingWriter) SendData(pkt wshrpc.CommandStreamData) error { + return fw.err +} + +// TestSendErrorDisconnectsClient verifies that when the data path to the main +// server fails (e.g. RPC stream timeout), the stream manager drops the client +// back to disconnected mode (buffering) instead of staying falsely attached. +func TestSendErrorDisconnectsClient(t *testing.T) { + tw := &testWriter{} + sm := MakeStreamManager() + defer sm.Close() + + // pipe stays open (no EOF) so the sender loop remains active + pr, pw := io.Pipe() + defer pr.Close() + defer pw.Close() + + err := sm.AttachReader(pr) + if err != nil { + t.Fatalf("AttachReader failed: %v", err) + } + + _, err = sm.ClientConnected("stream-1", tw, CwndSize, 0) + if err != nil { + t.Fatalf("ClientConnected failed: %v", err) + } + + if _, err := pw.Write([]byte("first chunk")); err != nil { + t.Fatalf("pipe write failed: %v", err) + } + time.Sleep(100 * time.Millisecond) + if len(tw.GetPackets()) == 0 { + t.Fatal("expected packets to flow after ClientConnected") + } + + // now the link goes stale: every send fails with a timeout-like error + tw.SetSendError(io.ErrClosedPipe) + + onSendErrCh := make(chan error, 1) + sm.OnSendError = func(err error) { + select { + case onSendErrCh <- err: + default: + } + } + + // new output forces another SendData call, which now fails + if _, err := pw.Write([]byte("second chunk")); err != nil { + t.Fatalf("pipe write failed: %v", err) + } + + deadline := time.Now().Add(2 * time.Second) + for { + sm.lock.Lock() + connected := sm.connected + sm.lock.Unlock() + if !connected { + break + } + if time.Now().After(deadline) { + t.Fatal("stream manager did not mark client disconnected after send error") + } + time.Sleep(5 * time.Millisecond) + } + + select { + case err := <-onSendErrCh: + if err == nil { + t.Fatal("expected non-nil error from OnSendError") + } + default: + t.Fatal("OnSendError callback was not invoked") + } + + // a new client must be able to attach and receive the buffered data + tw2 := &testWriter{} + _, err = sm.ClientConnected("stream-2", tw2, CwndSize, 0) + if err != nil { + t.Fatalf("reconnect after send error failed: %v", err) + } + time.Sleep(100 * time.Millisecond) + if len(tw2.GetPackets()) == 0 { + t.Fatal("expected buffered data to be delivered to reconnected client") + } +}