Fix ConnManager exit code race

internal/termpty/conn_manager.go

@@ -19,6 +19,7 @@ import (
 	"github.com/microsoft/dcp/pkg/concurrency"
 	usvc_io "github.com/microsoft/dcp/pkg/io"
 	"github.com/microsoft/dcp/pkg/osutil"
+	"github.com/microsoft/dcp/pkg/process"
 	"github.com/microsoft/dcp/pkg/resiliency"
 )
 
@@ -436,20 +437,7 @@ func (cm *ConnManager) serveConnection(conn net.Conn) {
 		defer wg.Done()
 		defer close(exitCodeCh)
 
-		// Wait for either the process to exit or the Serve loop to end.
-		select {
-		case <-ptp.ExitHandler.Exited():
-			ec := ptp.ExitHandler.ExitInfo().ExitCode
-			select {
-			case exitCodeCh <- ec:
-			case <-time.After(Hmp1ExitCodeWaitTimeout):
-				// Server has abandoned its exit-code read; close still
-				// happens via the deferred close below.
-			}
-
-		case <-serveCtx.Done():
-			// Process has not exited but we are done serving.
-		}
+		sendProcessExitCode(serveCtx, ptp.ExitHandler, exitCodeCh)
 	}()
 
 	serveErr := server.Serve(serveCtx, conn, exitCodeCh, Hmp1ServerOptions{
@@ -472,6 +460,35 @@ func (cm *ConnManager) serveConnection(conn net.Conn) {
 	wg.Wait()
 }
 
+// sendProcessExitCode is split out so tests can deterministically exercise the
+// process-exit vs. serve-cancellation timing.
+func sendProcessExitCode(
+	serveCtx context.Context,
+	exitHandler *process.ConcurrentProcessExitHandler,
+	exitCodeCh chan<- int32,
+) {
+	select {
+	case <-exitHandler.Exited():
+		sendExitCode(exitHandler, exitCodeCh)
+	case <-serveCtx.Done():
+		select {
+		case <-exitHandler.Exited():
+			sendExitCode(exitHandler, exitCodeCh)
+		default:
+			// Process has not exited but we are done serving.
+		}
+	}
+}
+
+func sendExitCode(exitHandler *process.ConcurrentProcessExitHandler, exitCodeCh chan<- int32) {
+	ec := exitHandler.ExitInfo().ExitCode
+	select {
+	case exitCodeCh <- ec:
+	case <-time.After(Hmp1ExitCodeWaitTimeout):
+		// Server has abandoned its exit-code read.
+	}
+}
+
 // shutdown() is the manager's shutdown sequence.
 // Idempotency is guaranteed via wrapping the method in a sync.Once.
 func (cm *ConnManager) shutdown() {

internal/termpty/conn_manager_test.go

@@ -602,6 +602,27 @@ func TestConnManager_ExitCodePropagatedToClient(t *testing.T) {
 	}
 }
 
+func TestSendProcessExitCodePrefersExitedProcessAfterServeCancel(t *testing.T) {
+	t.Parallel()
+
+	exitHandler := process.NewConcurrentProcessExitHandler()
+	const wantExit int32 = 99
+	exitHandler.OnProcessExited(process.Pid_t(12345), wantExit, nil)
+
+	serveCtx, serveCancel := context.WithCancel(context.Background())
+	serveCancel()
+
+	exitCodeCh := make(chan int32, 1)
+	sendProcessExitCode(serveCtx, exitHandler, exitCodeCh)
+
+	select {
+	case got := <-exitCodeCh:
+		require.Equal(t, wantExit, got)
+	default:
+		t.Fatal("expected process exit code to be sent when process exit and serve cancellation are both ready")
+	}
+}
+
 // readHmp1FrameAllowEOF reads a single HMP v1 frame from conn. It returns
 // (0, nil) on EOF / closed-connection style errors instead of failing the
 // test. Useful for shutdown path tests where the connection may close