[live-migration] harden resume against stdio and exit wait loss

internal/controller/process/process.go

@@ -70,12 +70,12 @@ type Controller struct {
 	// exitedCh is closed when the process has exited and all cleanup is done.
 	exitedCh chan struct{}
 
-	// vsock ports restored from a migrated process, used to reattach the
-	// stdio relay on resume.
+	// vsock ports for reattaching the stdio relay on resume, captured at Save on
+	// the source and on import on the destination.
 	stdinPort, stdoutPort, stderrPort uint32
 
-	// Wait request id carried over from a migrated process, reused on resume
-	// so no duplicate wait is issued. Zero if absent.
+	// Wait request id reused on resume so no duplicate wait is issued, captured
+	// at Save on the source and on import on the destination. Zero if absent.
 	waitCallID int64
 }
 
@@ -164,14 +164,17 @@ func (c *Controller) Start(ctx context.Context, events chan interface{}) (int, e
 	c.processID = c.process.Pid()
 	c.state = StateRunning
 
-	go c.handleProcessExit(ctx, execCmd, events)
+	go c.handleProcessExit(ctx, execCmd, events, true)
 
 	return c.processID, nil
 }
 
 // handleProcessExit blocks until the process exits, cleans up IO, and
 // publishes the exit event via events channel.
-func (c *Controller) handleProcessExit(ctx context.Context, execCmd *cmd.Cmd, events chan interface{}) {
+// In case of source rollback, there would be an existing instance of
+// handleProcessExit which would report the exit. Therefore, for the
+// duplicate call, we would exit early post cmd cleanup via cmd.Wait.
+func (c *Controller) handleProcessExit(ctx context.Context, execCmd *cmd.Cmd, events chan interface{}, reportExit bool) {
 	// Detach from the caller's context so upstream cancellation does
 	// not abort the background teardown.
 	ctx = context.WithoutCancel(ctx)
@@ -182,6 +185,12 @@ func (c *Controller) handleProcessExit(ctx context.Context, execCmd *cmd.Cmd, ev
 		log.G(ctx).WithError(err).Warn("process exit wait failed")
 	}
 
+	// A source rollback's re-attached relay only needs draining; the watcher
+	// started with the process reports the exit.
+	if !reportExit {
+		return
+	}
+
 	exitCode := execCmd.ExitState.ExitCode()
 
 	// Record the exit status under the lock.

internal/controller/process/process_test.go

@@ -13,6 +13,7 @@ import (
 	"github.com/opencontainers/runtime-spec/specs-go"
 	"go.uber.org/mock/gomock"
 
+	"github.com/Microsoft/hcsshim/internal/cmd"
 	"github.com/Microsoft/hcsshim/internal/controller/process/mocks"
 	hcs "github.com/Microsoft/hcsshim/internal/hcs/v2"
 )
@@ -256,6 +257,49 @@ func TestStart_HostCreateProcessFails(t *testing.T) {
 	}
 }
 
+// TestHandleProcessExit_DrainOnly verifies that with reportExit=false — a source
+// rollback's re-attached relay — handleProcessExit drains its command but leaves
+// exit reporting (state transition, upstream IO close, and the exit event) to the
+// watcher started with the process.
+func TestHandleProcessExit_DrainOnly(t *testing.T) {
+	t.Parallel()
+	mockCtrl, _, mockIO, controller := newSetup(t)
+	controller.upstreamIO = mockIO
+	controller.state = StateRunning
+	mockProc := mocks.NewMockProcess(mockCtrl)
+
+	// cmd.Attach reads Pid (for logging) and Stdio; nil IO means no relay goroutines.
+	mockProc.EXPECT().Pid().Return(testPID)
+	mockProc.EXPECT().Stdio().Return(nil, nil, nil)
+	// execCmd.Wait drives Process.Wait, ExitCode, and Close exactly once.
+	mockProc.EXPECT().Wait().Return(nil)
+	mockProc.EXPECT().ExitCode().Return(0, nil)
+	mockProc.EXPECT().Close().Return(nil)
+
+	execCmd, err := cmd.Attach(context.WithoutCancel(t.Context()), mockProc, nil, nil, nil)
+	if err != nil {
+		t.Fatalf("Attach() = %v; want nil", err)
+	}
+
+	// No upstreamIO.Close is expected: the unset mock would fail if it were called.
+	events := make(chan interface{}, 1)
+	controller.handleProcessExit(t.Context(), execCmd, events, false)
+
+	if controller.State() != StateRunning {
+		t.Errorf("state = %s; want unchanged StateRunning", controller.State())
+	}
+	select {
+	case <-controller.exitedCh:
+		t.Error("exitedCh was closed; want left to the original watcher")
+	default:
+	}
+	select {
+	case ev := <-events:
+		t.Errorf("published event %v; want none", ev)
+	default:
+	}
+}
+
 // TestKill_NotCreatedState verifies that Kill on a process that was never
 // created transitions it directly to StateTerminated without error. Because
 // upstreamIO has not been populated yet, abortInternal must tolerate a nil

internal/controller/process/save.go

@@ -48,6 +48,10 @@ func (c *Controller) Save(ctx context.Context) (*anypb.Any, error) {
 		ms := c.process.MigrationState()
 		state.StdinPort, state.StdoutPort, state.StderrPort = ms.StdinPort, ms.StdoutPort, ms.StderrPort
 		state.WaitCallID = ms.WaitCallID
+
+		// Retain them so a source rollback resume re-opens IO like the destination.
+		c.stdinPort, c.stdoutPort, c.stderrPort = ms.StdinPort, ms.StdoutPort, ms.StderrPort
+		c.waitCallID = ms.WaitCallID
 	}
 
 	// Exec processes carry their OCI spec; init processes leave it unset.
@@ -179,8 +183,9 @@ func (c *Controller) Patch(ctx context.Context, containerID string, opts *Create
 
 // Resume returns a migrating process to the running state. On the destination
 // it reattaches the patched process to its live guest counterpart, wires up the
-// stdio relay, and begins watching for exit. On the source it simply lifts the
-// freeze that Save applied, since the live process and IO are still intact.
+// stdio relay, and begins watching for exit. On the source it re-opens the IO
+// the blackout dropped and resumes the relay, since the live process is intact
+// but its IO connections are not.
 // Pass events=nil for an init process, whose exit is reported by its owning
 // container instead.
 func (c *Controller) Resume(ctx context.Context, gcsContainer *gcs.Container, events chan interface{}) error {
@@ -192,19 +197,16 @@ func (c *Controller) Resume(ctx context.Context, gcsContainer *gcs.Container, ev
 		return nil
 	}
 
-	// Source rollback: the live process and IO are intact, so just lift the
-	// freeze that Save applied.
-	if c.state == StateSourceMigrating {
-		c.state = StateRunning
-		return nil
-	}
-
-	if c.state != StateDestinationMigrating {
+	if c.state != StateDestinationMigrating && c.state != StateSourceMigrating {
 		return fmt.Errorf("process %q in container %q is in state %s; cannot resume: %w", c.execID, c.containerID, c.state, errdefs.ErrFailedPrecondition)
 	}
 
-	// Reopen the live process on its preserved IO ports and wait id.
-	gcsProc, err := gcsContainer.OpenProcessWithIO(ctx, uint32(c.processID), c.stdinPort, c.stdoutPort, c.stderrPort, c.waitCallID)
+	// Flag to determine if the resume is happening on destination.
+	isDestination := c.state == StateDestinationMigrating
+
+	// Reopen the process on its preserved IO ports and wait id. A source rollback
+	// reuses the still-outstanding wait, so it does not start a second one.
+	gcsProc, err := gcsContainer.OpenProcessWithIO(ctx, uint32(c.processID), c.stdinPort, c.stdoutPort, c.stderrPort, c.waitCallID, isDestination)
 	if err != nil {
 		return fmt.Errorf("open gcs process pid %d in container %q: %w", c.processID, c.containerID, err)
 	}
@@ -223,9 +225,10 @@ func (c *Controller) Resume(ctx context.Context, gcsContainer *gcs.Container, ev
 	// Ports are single-use; clear them now that IO is reattached.
 	c.stdinPort, c.stdoutPort, c.stderrPort = 0, 0, 0
 
-	// Watch for exit in the background, mirroring a freshly started process.
-	go c.handleProcessExit(ctx, execCmd, events)
+	// The destination owns exit reporting; a source rollback leaves that to the
+	// watcher from Start, so this handler only drains the re-attached relay.
+	go c.handleProcessExit(ctx, execCmd, events, isDestination)
 
-	log.G(ctx).WithField(logfields.ProcessID, c.processID).Debug("resumed migrated process on destination")
+	log.G(ctx).WithField(logfields.ProcessID, c.processID).Debug("resumed migrated process")
 	return nil
 }

internal/controller/process/save_test.go

@@ -116,6 +116,14 @@ func TestSave_Succeeds(t *testing.T) {
 			if controller.state != StateSourceMigrating {
 				t.Errorf("state = %s; want StateSourceMigrating", controller.state)
 			}
+			// The ports and wait id are retained on the controller so a source
+			// rollback resume re-opens IO the same way the destination does.
+			if controller.stdinPort != testStdinPort || controller.stdoutPort != testStdoutPort || controller.stderrPort != testStderrPort {
+				t.Errorf("controller ports = (%d,%d,%d); want (%d,%d,%d)", controller.stdinPort, controller.stdoutPort, controller.stderrPort, testStdinPort, testStdoutPort, testStderrPort)
+			}
+			if controller.waitCallID != testWaitCallID {
+				t.Errorf("controller waitCallID = %d; want %d", controller.waitCallID, testWaitCallID)
+			}
 		})
 	}
 }
@@ -334,22 +342,6 @@ func TestResume_WrongState(t *testing.T) {
 	}
 }
 
-// TestResume_SourceRollback verifies that resuming a source-migrating process
-// lifts the freeze and returns it to running without touching the host.
-func TestResume_SourceRollback(t *testing.T) {
-	t.Parallel()
-	_, _, _, controller := newSetup(t)
-	controller.state = StateSourceMigrating
-
-	// nil host/events are unused: the live process and IO stay intact.
-	if err := controller.Resume(t.Context(), nil, nil); err != nil {
-		t.Fatalf("Resume() = %v; want nil", err)
-	}
-	if controller.state != StateRunning {
-		t.Errorf("state = %s; want StateRunning", controller.state)
-	}
-}
-
 // TestResume_IdempotentWhenRunning verifies that resuming an already-resumed
 // process is a no-op, so a retry after a completed resume is safe.
 func TestResume_IdempotentWhenRunning(t *testing.T) {

internal/gcs/bridge.go

@@ -572,8 +572,10 @@ func (brdg *bridge) PreregisterRPC(id int64, proc prot.RPCProc, resp responseMes
 	if brdg.rpcs == nil {
 		return nil, ErrBridgeClosed
 	}
-	if _, dup := brdg.rpcs[id]; dup {
-		return nil, fmt.Errorf("preregister rpc: id %d already in use", id)
+	if existing, dup := brdg.rpcs[id]; dup {
+		// A source rollback re-opens a process whose wait is still outstanding;
+		// hand back that call.
+		return existing, nil
 	}
 	brdg.rpcs[id] = call
 	return call, nil

internal/gcs/bridge_test.go

@@ -256,3 +256,24 @@ func TestRPCErrorUnwrapHCSCode(t *testing.T) {
 		t.Fatalf("hcs.IsNotExist(wrapped) = false; want true (err=%v)", wrapped)
 	}
 }
+
+// TestPreregisterRPCReusesOutstanding verifies that pre-registering an id that
+// is already outstanding hands back the existing call (a source rollback
+// re-opens a process whose wait is still pending) rather than failing.
+func TestPreregisterRPCReusesOutstanding(t *testing.T) {
+	s, _ := pipeConn()
+	b := newBridge(s, nil, logrus.NewEntry(logrus.StandardLogger()))
+
+	first, err := b.PreregisterRPC(7, prot.RPCWaitForProcess, &testResp{})
+	if err != nil {
+		t.Fatalf("first PreregisterRPC = %v; want nil", err)
+	}
+
+	again, err := b.PreregisterRPC(7, prot.RPCWaitForProcess, &testResp{})
+	if err != nil {
+		t.Fatalf("duplicate PreregisterRPC = %v; want nil", err)
+	}
+	if again != first {
+		t.Errorf("duplicate PreregisterRPC returned a new call; want the outstanding one")
+	}
+}

internal/gcs/container.go

@@ -125,8 +125,12 @@ func (c *Container) CreateProcess(ctx context.Context, config interface{}) (_ co
 // [Container.CreateProcess]: it attaches to a process already running
 // in this container, re-listens on the supplied vsock ports, and
 // pre-registers the source bridge's WaitForProcess id so the guest's
-// still-outstanding response is routed without arming a duplicate wait.
-func (c *Container) OpenProcessWithIO(ctx context.Context, pid uint32, stdinPort, stdoutPort, stderrPort uint32, waitCallID int64) (_ *Process, err error) {
+// still-outstanding response is routed.
+// watchExit launches the background exit wait, first probing for an already-
+// exited process whose adopted response the guest never redelivers; pass false
+// when the caller already has one outstanding (a source rollback reuses the
+// live process's).
+func (c *Container) OpenProcessWithIO(ctx context.Context, pid uint32, stdinPort, stdoutPort, stderrPort uint32, waitCallID int64, watchExit bool) (_ *Process, err error) {
 	ctx, span := ot.StartSpan(ctx, "gcs::Container::OpenProcessWithIO", ot.WithClientSpanKind)
 	defer span.End()
 	defer func() { ot.SetSpanStatus(span, err) }()
@@ -173,11 +177,21 @@ func (c *Container) OpenProcessWithIO(ctx context.Context, pid uint32, stdinPort
 		return nil, err
 	}
 
-	p.waitCall, err = c.gc.brdg.PreregisterRPC(waitCallID, prot.RPCWaitForProcess, &p.waitResp)
+	p.waitResp = &prot.ContainerWaitForProcessResponse{}
+	p.waitCall, err = c.gc.brdg.PreregisterRPC(waitCallID, prot.RPCWaitForProcess, p.waitResp)
 	if err != nil {
 		return nil, fmt.Errorf("preregister wait for pid %d in container %s (id %d): %w", pid, c.id, waitCallID, err)
 	}
-	go p.waitBackground()
+	// A reused outstanding wait carries its own response; adopt it so the
+	// reported exit code reflects the guest's reply and not this unused one.
+	if resp, ok := p.waitCall.resp.(*prot.ContainerWaitForProcessResponse); ok {
+		p.waitResp = resp
+	}
+	// Establish exit reporting unless the caller already watches this process
+	// (a reuse that already has a wait outstanding, e.g. a rollback).
+	if watchExit {
+		p.startExitWatch(ctx)
+	}
 	log.G(ctx).WithField("pid", p.id).Debug("opened existing process with IO")
 	return p, nil
 }

internal/gcs/process.go

@@ -22,6 +22,12 @@ import (
 
 const (
 	hrNotFound = 0x80070490
+
+	// exitProbeTimeoutMs bounds the wait used to detect a process that already
+	// exited when it is reopened (see startExitWatch). The finite timeout lets
+	// the guest-side waiter self-clean when the process is still running, so it
+	// never lingers if the process is reopened again later.
+	exitProbeTimeoutMs = 500
 )
 
 // Process represents a process in a container or container host.
@@ -30,7 +36,7 @@ type Process struct {
 	cid                   string
 	id                    uint32
 	waitCall              *rpc
-	waitResp              prot.ContainerWaitForProcessResponse
+	waitResp              *prot.ContainerWaitForProcessResponse
 	stdin, stdout, stderr *ioChannel
 	stdinCloseWriteOnce   sync.Once
 	stdinCloseWriteErr    error
@@ -120,7 +126,8 @@ func (gc *GuestConnection) exec(ctx context.Context, cid string, params interfac
 		ProcessID:   p.id,
 		TimeoutInMs: 0xffffffff,
 	}
-	p.waitCall, err = gc.brdg.AsyncRPC(ctx, prot.RPCWaitForProcess, &waitReq, &p.waitResp)
+	p.waitResp = &prot.ContainerWaitForProcessResponse{}
+	p.waitCall, err = gc.brdg.AsyncRPC(ctx, prot.RPCWaitForProcess, &waitReq, p.waitResp)
 	if err != nil {
 		return nil, fmt.Errorf("failed to wait on process, leaking process: %w", err)
 	}
@@ -322,3 +329,40 @@ func (p *Process) waitBackground() {
 	log.G(ctx).WithField("exitCode", ec).Debug("process exited")
 	ot.SetSpanStatus(span, err)
 }
+
+// startExitWatch arranges for a reopened process's exit to be reported.
+// OpenProcessWithIO has already pre-registered p.waitCall — the WaitForProcess
+// call the previous owner left outstanding in the guest — which completes when
+// the process exits. That suffices for a process still running at reopen, but
+// one that exited while no bridge was connected had its response dropped on the
+// severed connection and never resent, so waiting on p.waitCall alone would
+// hang.
+//
+// To cover that, startExitWatch issues its own bounded WaitForProcess. If the
+// process has already exited the guest returns the retained exit code at once,
+// and that completed call replaces p.waitCall so Wait and ExitCode report
+// through the normal path. Otherwise the probe times out and self-cleans (no
+// lingering guest-side waiter) and p.waitCall is watched in the background.
+func (p *Process) startExitWatch(ctx context.Context) {
+	req := prot.ContainerWaitForProcess{
+		RequestBase: makeRequest(ctx, p.cid),
+		ProcessID:   p.id,
+		TimeoutInMs: exitProbeTimeoutMs,
+	}
+	resp := &prot.ContainerWaitForProcessResponse{}
+	if probe, err := p.gc.brdg.AsyncRPC(ctx, prot.RPCWaitForProcess, &req, resp); err == nil {
+		probe.Wait()
+		// A clean response means the process already exited; a still-running one
+		// makes the guest return a timeout error.
+		if probe.Err() == nil {
+			// Make this completed probe the process's wait so its retained exit
+			// code flows through the normal Wait and ExitCode path.
+			p.waitCall = probe
+			p.waitResp = resp
+			return
+		}
+	}
+	// Still running (or the probe could not be issued): watch the pre-registered
+	// wait (p.waitCall) in the background.
+	go p.waitBackground()
+}