fix: recover durable SSH sessions stuck falsely-attached after stream timeout (#3439)

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

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) {

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()

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 {

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)
+			}
+		}
 	}
 }
 

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")
+	}
+}