[live-migration] Fix source resume hang when cancelled before blackout

internal/controller/vm/save_lcow.go

@@ -221,47 +221,46 @@ func (c *Controller) Resume(ctx context.Context, rebuildBridge bool) error {
 		return fmt.Errorf("cannot resume from migration: VM is in state %s: %w", c.vmState, errdefs.ErrFailedPrecondition)
 	}
 
-	// On the destination, the log connection was never established.
-	// On source, the blackout dropped the source's GCS log connection, which tore
-	// down its listener and closed logOutputDone. Install a fresh signal and
-	// re-arm the listener so the resumed guest's reconnect-mode vsockexec can
-	// reconnect and host-side logs resume.
-	c.logOutputDone = make(chan struct{})
-	// We expect the reconnect to complete within the GCS connection timeout,
-	// otherwise we want to fail.
-	ctx, cancel := context.WithTimeout(ctx, timeout.GCSConnectionTimeout)
-	log.G(ctx).Debugf("using gcs connection timeout: %s\n", timeout.GCSConnectionTimeout)
-
-	g, gctx := errgroup.WithContext(ctx)
-	defer func() {
-		_ = g.Wait()
-	}()
-	defer cancel()
-
-	if err := c.setupLoggingListener(gctx, g); err != nil {
-		return fmt.Errorf("arm logging listener on resume: %w", err)
-	}
-
-	if rebuildBridge {
-		// Source rollback: arm the host GCS listener now, then accept the guest's
-		// post-blackout re-dial and swap it into the running bridge.
-		if err := c.guest.PrepareConnection(winio.VsockServiceID(prot.LinuxGcsVsockPort)); err != nil {
-			return fmt.Errorf("prepare source resume listener: %w", err)
-		}
-		if err := c.guest.ResumeConnection(ctx); err != nil {
-			return fmt.Errorf("resume source guest connection: %w", err)
+	// A source rollback before blackout never dropped the guest connection, so the
+	// live bridge and its log stream are reused instead of re-accepted.
+	reuseLiveConn := rebuildBridge && c.guest.IsBridgeConnected()
+
+	// Reuse skips the work below. Otherwise re-arm host logging and (re)establish the
+	// bridge, bounded by the GCS connection timeout; the shared tail keeps the caller ctx.
+	if !reuseLiveConn {
+		timeoutCtx, cancel := context.WithTimeout(ctx, timeout.GCSConnectionTimeout)
+		log.G(timeoutCtx).Debugf("using gcs connection timeout: %s\n", timeout.GCSConnectionTimeout)
+		g, gctx := errgroup.WithContext(timeoutCtx)
+		defer func() {
+			_ = g.Wait()
+		}()
+		defer cancel()
+
+		// Logs were dropped by the source blackout and never established back.
+		c.logOutputDone = make(chan struct{})
+		if err := c.setupLoggingListener(gctx, g); err != nil {
+			return fmt.Errorf("arm logging listener on resume: %w", err)
 		}
-	} else {
-		// Destination: reuse the connection already armed at start.
-		if err := c.guest.CreateConnection(ctx, false); err != nil {
-			return fmt.Errorf("resume destination guest connection: %w", err)
+
+		if rebuildBridge {
+			// Source rollback after blackout: accept the guest's re-dial into the bridge.
+			if err := c.guest.PrepareConnection(winio.VsockServiceID(prot.LinuxGcsVsockPort)); err != nil {
+				return fmt.Errorf("prepare source resume listener: %w", err)
+			}
+			if err := c.guest.ResumeConnection(timeoutCtx); err != nil {
+				return fmt.Errorf("resume source guest connection: %w", err)
+			}
+		} else {
+			// Destination: reuse the connection already armed at start.
+			if err := c.guest.CreateConnection(timeoutCtx, false); err != nil {
+				return fmt.Errorf("resume destination guest connection: %w", err)
+			}
 		}
-	}
 
-	// Collect any errors from establishing the log connection.
-	// If the connection is not established then we need to error out.
-	if err := g.Wait(); err != nil {
-		return err
+		// Fail if the log connection could not be established.
+		if err := g.Wait(); err != nil {
+			return err
+		}
 	}
 
 	// Clear migrating flag only now that the new transport is in place.

internal/gcs/bridge.go

@@ -54,15 +54,21 @@ type bridge struct {
 	rpcs   map[int64]*rpc
 	// conn is the transport carrying messages to and from the guest.
 	// Held atomically because the send path reads it while a migration swaps it.
-	conn      atomic.Value
-	rpcCh     chan *rpc
-	notify    notifyFunc
-	closed    bool
-	log       *logrus.Entry
-	brdgErr   error
-	waitCh    chan struct{}
+	conn    atomic.Value
+	rpcCh   chan *rpc
+	notify  notifyFunc
+	closed  bool
+	log     *logrus.Entry
+	brdgErr error
+	waitCh  chan struct{}
+
+	// Migration related fields
+	// migrating tolerates transport drops during a live-migration window.
 	migrating atomic.Bool
-	resumeCh  chan struct{}
+	// resumeCh wakes the parked recv loop when a new transport is swapped in.
+	resumeCh chan struct{}
+	// connected is true while a live transport is present and being read.
+	connected atomic.Bool
 }
 
 var ErrBridgeClosed = fmt.Errorf("bridge closed: %w", net.ErrClosed)
@@ -305,7 +311,9 @@ func (brdg *bridge) RPC(ctx context.Context, proc prot.RPCProc, req requestMessa
 
 func (brdg *bridge) recvLoopRoutine() {
 	for {
+		brdg.connected.Store(true)
 		err := brdg.recvLoop()
+		brdg.connected.Store(false)
 
 		if !brdg.migrating.Load() {
 			brdg.kill(err)

internal/gcs/guestconnection.go

@@ -244,6 +244,14 @@ func (gc *GuestConnection) SetMigrating(migrating bool) {
 	gc.brdg.SetMigrating(migrating)
 }
 
+// IsBridgeConnected reports whether a live bridge transport is currently installed.
+func (gc *GuestConnection) IsBridgeConnected() bool {
+	if gc.brdg == nil {
+		return false
+	}
+	return gc.brdg.connected.Load()
+}
+
 // ResumeOnConn resumes the bridge after swaping the bridge
 // transport without dropping outstanding RPCs.
 func (gc *GuestConnection) ResumeOnConn(ctx context.Context, conn io.ReadWriteCloser) error {

internal/vm/guestmanager/guest.go

@@ -225,6 +225,17 @@ func (gm *Guest) SetMigrating(migrating bool) {
 	gm.gc.SetMigrating(migrating)
 }
 
+// IsBridgeConnected reports whether a live bridge transport is currently installed.
+func (gm *Guest) IsBridgeConnected() bool {
+	gm.mu.RLock()
+	defer gm.mu.RUnlock()
+
+	if gm.gc == nil {
+		return false
+	}
+	return gm.gc.IsBridgeConnected()
+}
+
 // ResumeConnection accepts a fresh hvsock on the prepared listener and
 // swaps it into the existing GCS bridge, preserving in-flight RPCs.
 func (gm *Guest) ResumeConnection(ctx context.Context) error {