[live-migration] restore container stdio on source rollback

When a live migration rolls back to the source, the VM resumes but the
container's stdout/stderr were dropped during blackout and never restored,
so anything watching the container's output saw it stop for good.

Resume now brings those streams back the same way the destination already
does, so a rolled-back container keeps streaming its output as if the
migration had never been attempted.

Signed-off-by: Harsh Rawat <harshrawat@microsoft.com>

Author: rawahars

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,10 @@ 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.
+// startWait launches the background exit wait; 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, startWait bool) (_ *Process, err error) {
 	ctx, span := ot.StartSpan(ctx, "gcs::Container::OpenProcessWithIO", ot.WithClientSpanKind)
 	defer span.End()
 	defer func() { ot.SetSpanStatus(span, err) }()
@@ -177,7 +179,9 @@ func (c *Container) OpenProcessWithIO(ctx context.Context, pid uint32, stdinPort
 	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()
+	if startWait {
+		go p.waitBackground()
+	}
 	log.G(ctx).WithField("pid", p.id).Debug("opened existing process with IO")
 	return p, nil
 }