Add LCOW live migration support across the controller stacks

internal/controller/device/plan9/save.go

@@ -0,0 +1,20 @@
+//go:build windows && lcow
+
+package plan9
+
+import (
+	"fmt"
+)
+
+// Save is not yet supported for the Plan9 sub-controller; any tracked state
+// indicates a live-migration scenario the controller cannot represent.
+func (c *Controller) Save() error {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	if len(c.sharesByHostPath) > 0 || len(c.reservations) > 0 {
+		return fmt.Errorf("plan9 controller save not supported: %d shares, %d reservations", len(c.sharesByHostPath), len(c.reservations))
+	}
+
+	return nil
+}

internal/controller/device/plan9/save_test.go

@@ -0,0 +1,33 @@
+//go:build windows && lcow
+
+package plan9
+
+import (
+	"testing"
+
+	"github.com/Microsoft/go-winio/pkg/guid"
+
+	"github.com/Microsoft/hcsshim/internal/controller/device/plan9/share"
+)
+
+func TestSave_EmptyOK(t *testing.T) {
+	c := &Controller{
+		reservations:     map[guid.GUID]*reservation{},
+		sharesByHostPath: map[string]*share.Share{},
+	}
+
+	if err := c.Save(); err != nil {
+		t.Fatalf("Save on empty controller: %v", err)
+	}
+}
+
+func TestSave_NonEmptyErrors(t *testing.T) {
+	c := &Controller{
+		reservations:     map[guid.GUID]*reservation{{}: {hostPath: "/h"}},
+		sharesByHostPath: map[string]*share.Share{},
+	}
+
+	if err := c.Save(); err == nil {
+		t.Fatal("expected Save to error when reservations are present")
+	}
+}

internal/controller/device/scsi/controller.go

@@ -34,7 +34,7 @@ import (
 // it succeeds to release the reservation and all resources.
 type Controller struct {
 	// mu serializes all public operations on the Controller.
-	mu sync.Mutex
+	mu sync.RWMutex
 
 	// vm is the host-side interface for adding and removing SCSI disks.
 	// Immutable after construction.
@@ -58,6 +58,10 @@ type Controller struct {
 	//   ControllerID = index / numLUNsPerController
 	//   LUN          = index % numLUNsPerController
 	controllerSlots []*disk.Disk
+
+	// isMigrating rejects all public ops while set: true once a snapshot has
+	// been taken or imported, until migration is resumed. Guarded by mu.
+	isMigrating bool
 }
 
 // New creates a new [Controller] for the given number of SCSI controllers and
@@ -78,18 +82,22 @@ func New(numControllers int, vm VMSCSIOps, guest GuestSCSIOps) *Controller {
 // once per controller and lun location, and must be called before any calls to
 // Reserve() to ensure the rootfs reservation is not evicted by a dynamic
 // reservation.
-func (c *Controller) ReserveForRootfs(ctx context.Context, controller, lun uint) error {
+func (c *Controller) ReserveForRootfs(ctx context.Context, controller, lun uint, cfg disk.Config) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	if c.isMigrating {
+		return fmt.Errorf("SCSI controller is migrating; call Resume first")
+	}
+
 	slot := int(controller*numLUNsPerController + lun)
 	if slot >= len(c.controllerSlots) {
 		return fmt.Errorf("invalid controller %d or lun %d", controller, lun)
 	}
 	if c.controllerSlots[slot] != nil {
 		return fmt.Errorf("slot for controller %d and lun %d is already reserved", controller, lun)
 	}
-	c.controllerSlots[slot] = disk.NewReserved(controller, lun, disk.Config{})
+	c.controllerSlots[slot] = disk.NewReserved(controller, lun, cfg)
 	return nil
 }
 
@@ -103,6 +111,10 @@ func (c *Controller) Reserve(ctx context.Context, diskConfig disk.Config, mountC
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	if c.isMigrating {
+		return guid.GUID{}, fmt.Errorf("SCSI controller is migrating; call Resume first")
+	}
+
 	ctx, _ = log.WithContext(ctx, logrus.WithFields(logrus.Fields{
 		logfields.HostPath:  diskConfig.HostPath,
 		logfields.Partition: mountConfig.Partition,
@@ -178,6 +190,10 @@ func (c *Controller) MapToGuest(ctx context.Context, id guid.GUID) (string, erro
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	if c.isMigrating {
+		return "", fmt.Errorf("SCSI controller is migrating; call Resume first")
+	}
+
 	r, ok := c.reservations[id]
 	if !ok {
 		return "", fmt.Errorf("reservation %s not found", id)
@@ -212,6 +228,10 @@ func (c *Controller) UnmapFromGuest(ctx context.Context, id guid.GUID) error {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
+	if c.isMigrating {
+		return fmt.Errorf("SCSI controller is migrating; call Resume first")
+	}
+
 	ctx, _ = log.WithContext(ctx, logrus.WithField("reservation", id.String()))
 
 	r, ok := c.reservations[id]

internal/controller/device/scsi/controller_test.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"errors"
 	"fmt"
+	"strings"
 	"testing"
 
 	"github.com/Microsoft/hcsshim/internal/controller/device/scsi/disk"
@@ -70,6 +71,17 @@ func mappedController(t *testing.T) (*Controller, guid.GUID) {
 	return c, id
 }
 
+func attachmentsContainPath(att map[string]hcsschema.Scsi, path string) bool {
+	for _, s := range att {
+		for _, a := range s.Attachments {
+			if a.Path == path {
+				return true
+			}
+		}
+	}
+	return false
+}
+
 // --- Tests: New ---
 
 func TestNew(t *testing.T) {
@@ -397,3 +409,111 @@ func TestUnmapFromGuest_RetryAfterDetachFailure(t *testing.T) {
 		t.Fatalf("re-reserve after retry: %v", err)
 	}
 }
+
+// --- Tests: ReserveForRootfs ---
+
+func TestReserveForRootfs_Success(t *testing.T) {
+	c := newController(&mockVMOps{}, newMockGuestOps())
+	cfg := defaultDiskConfig()
+	if err := c.ReserveForRootfs(context.Background(), 0, 0, cfg); err != nil {
+		t.Fatalf("unexpected error: %v", err)
+	}
+	// The reserved rootfs disk surfaces in the VM topology with its config.
+	if !attachmentsContainPath(c.HCSAttachments(), cfg.HostPath) {
+		t.Errorf("expected rootfs path %q in HCS attachments", cfg.HostPath)
+	}
+}
+
+func TestReserveForRootfs_InvalidLocation(t *testing.T) {
+	c := newController(&mockVMOps{}, newMockGuestOps())
+	// Controller index beyond the single configured controller.
+	if err := c.ReserveForRootfs(context.Background(), 1, 0, defaultDiskConfig()); err == nil {
+		t.Fatal("expected error for out-of-range location")
+	}
+}
+
+func TestReserveForRootfs_AlreadyReserved(t *testing.T) {
+	c := newController(&mockVMOps{}, newMockGuestOps())
+	if err := c.ReserveForRootfs(context.Background(), 0, 0, defaultDiskConfig()); err != nil {
+		t.Fatalf("first reserve: %v", err)
+	}
+	if err := c.ReserveForRootfs(context.Background(), 0, 0, defaultDiskConfig()); err == nil {
+		t.Fatal("expected error reserving an occupied location")
+	}
+}
+
+// --- Tests: migration guard ---
+
+func TestPublicOps_RejectedWhileMigrating(t *testing.T) {
+	ctx := t.Context()
+	ops := []struct {
+		name string
+		call func(*Controller) error
+	}{
+		{"ReserveForRootfs", func(c *Controller) error {
+			return c.ReserveForRootfs(ctx, 0, 0, defaultDiskConfig())
+		}},
+		{"Reserve", func(c *Controller) error {
+			_, err := c.Reserve(ctx, defaultDiskConfig(), defaultMountConfig())
+			return err
+		}},
+		{"MapToGuest", func(c *Controller) error {
+			_, err := c.MapToGuest(ctx, guid.GUID{})
+			return err
+		}},
+		{"UnmapFromGuest", func(c *Controller) error {
+			return c.UnmapFromGuest(ctx, guid.GUID{})
+		}},
+	}
+	for _, op := range ops {
+		t.Run(op.name, func(t *testing.T) {
+			// Saving the source blocks further operations until migration resumes.
+			src := New(1, &mockVMOps{}, newMockGuestOps())
+			env, err := src.Save(ctx)
+			if err != nil {
+				t.Fatalf("Save: %v", err)
+			}
+			if err := op.call(src); err == nil || !strings.Contains(err.Error(), "migrating") {
+				t.Fatalf("source: expected migrating error, got %v", err)
+			}
+
+			// A freshly imported controller is also mid-migration until resumed.
+			c, err := Import(ctx, env)
+			if err != nil {
+				t.Fatalf("Import: %v", err)
+			}
+			if err := op.call(c); err == nil || !strings.Contains(err.Error(), "migrating") {
+				t.Fatalf("imported: expected migrating error, got %v", err)
+			}
+		})
+	}
+}
+
+func TestResume_LiftsMigrationGuard(t *testing.T) {
+	ctx := t.Context()
+	// Snapshot a controller holding a reservation, then import it.
+	src := newController(&mockVMOps{}, newMockGuestOps())
+	if _, err := src.Reserve(ctx, defaultDiskConfig(), defaultMountConfig()); err != nil {
+		t.Fatalf("setup Reserve: %v", err)
+	}
+	env, err := src.Save(ctx)
+	if err != nil {
+		t.Fatalf("Save: %v", err)
+	}
+	c, err := Import(ctx, env)
+	if err != nil {
+		t.Fatalf("Import: %v", err)
+	}
+
+	// Rejected while migrating.
+	if _, err := c.Reserve(ctx, defaultDiskConfig(), defaultMountConfig()); err == nil {
+		t.Fatal("expected migrating error before Resume")
+	}
+
+	// Resuming binds live interfaces and lifts the guard.
+	c.Resume(ctx, &mockVMOps{}, newMockGuestOps())
+	dc := disk.Config{HostPath: `C:\other.vhdx`, Type: disk.TypeVirtualDisk}
+	if _, err := c.Reserve(ctx, dc, defaultMountConfig()); err != nil {
+		t.Fatalf("Reserve after Resume: %v", err)
+	}
+}

internal/controller/device/scsi/disk/save.go

@@ -0,0 +1,84 @@
+//go:build windows && (lcow || wcow)
+
+package disk
+
+import (
+	"fmt"
+
+	"github.com/Microsoft/hcsshim/internal/controller/device/scsi/mount"
+	scsisave "github.com/Microsoft/hcsshim/internal/controller/device/scsi/save"
+)
+
+// Save returns a migration snapshot of the disk and its mounts. It fails unless
+// the disk is attached or reserved and every mount can be saved.
+func (d *Disk) Save() (*scsisave.DiskState, error) {
+	if d.state != StateAttached && d.state != StateReserved {
+		return nil, fmt.Errorf("scsi disk controller=%d lun=%d in state %s; want %s", d.controller, d.lun, d.state, StateAttached)
+	}
+
+	out := &scsisave.DiskState{
+		Config: &scsisave.DiskConfig{
+			HostPath: d.config.HostPath,
+			ReadOnly: d.config.ReadOnly,
+			Type:     string(d.config.Type),
+			EvdType:  d.config.EVDType,
+		},
+	}
+
+	if len(d.mounts) > 0 {
+		out.Mounts = make(map[uint64]*scsisave.MountState, len(d.mounts))
+
+		// Snapshot every mount; abort if any cannot be saved.
+		for partition, m := range d.mounts {
+			ms, err := m.Save()
+			if err != nil {
+				return nil, err
+			}
+			out.Mounts[partition] = ms
+		}
+	}
+	return out, nil
+}
+
+// Import reconstructs a disk and its mounts from a migration snapshot at the
+// given controller and lun. It returns nil if the snapshot is nil.
+func Import(state *scsisave.DiskState, controller, lun uint) *Disk {
+	if state == nil {
+		return nil
+	}
+
+	// Rebuild the host-side config from the snapshot, if present.
+	cfg := Config{}
+	if c := state.GetConfig(); c != nil {
+		cfg = Config{
+			HostPath: c.GetHostPath(),
+			ReadOnly: c.GetReadOnly(),
+			Type:     Type(c.GetType()),
+			EVDType:  c.GetEvdType(),
+		}
+	}
+
+	// An imported disk is assumed to be live on the SCSI bus.
+	d := &Disk{
+		controller: controller,
+		lun:        lun,
+		config:     cfg,
+		state:      StateAttached,
+		mounts:     make(map[uint64]*mount.Mount, len(state.GetMounts())),
+	}
+
+	// Reconstruct each partition mount, skipping any that fail to import.
+	for partition, ms := range state.GetMounts() {
+		m := mount.Import(ms, controller, lun, partition)
+		if m == nil {
+			continue
+		}
+		d.mounts[partition] = m
+	}
+	return d
+}
+
+// UpdateHostPath rewrites the host-side path of the disk image.
+func (d *Disk) UpdateHostPath(p string) {
+	d.config.HostPath = p
+}