refactor share state management to introduce StateInvalid for failed attempts

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

Author: rawahars

internal/controller/device/plan9/controller.go

@@ -23,7 +23,7 @@ import (
 //
 // 1. Obtain a reservation using Reserve().
 //
-// 2. Use the reservation in MapToGuest() to mount the share unto guest.
+// 2. Use the reservation in MapToGuest() to mount the share into the guest.
 //
 // 3. Call UnmapFromGuest() to release the reservation and all resources.
 //
@@ -171,8 +171,8 @@ func (c *Controller) MapToGuest(ctx context.Context, id guid.GUID) (string, erro
 
 	// Validate if the host path has an associated share.
 	// This should be reserved by the Reserve() call.
-	existingShare := c.sharesByHostPath[res.hostPath]
-	if existingShare == nil {
+	existingShare, ok := c.sharesByHostPath[res.hostPath]
+	if !ok {
 		return "", fmt.Errorf("share for host path %s not found", res.hostPath)
 	}
 
@@ -210,14 +210,9 @@ func (c *Controller) UnmapFromGuest(ctx context.Context, id guid.GUID) error {
 
 	// Validate that the share exists before proceeding with teardown.
 	// This should be reserved by the Reserve() call.
-	existingShare := c.sharesByHostPath[res.hostPath]
-	if existingShare == nil {
-		// Share is gone — a sibling reservation already cleaned it up.
-		// Example: A and B reserve the same path; A's AddToVM fails
-		// (share→StateRemoved) and A's UnmapFromGuest deletes the map entry.
-		// B has nothing left to clean up, so just drop the reservation.
-		delete(c.reservations, id)
-		return nil
+	existingShare, ok := c.sharesByHostPath[res.hostPath]
+	if !ok {
+		return fmt.Errorf("share for host path %s not found", res.hostPath)
 	}
 
 	log.G(ctx).WithField(logfields.HostPath, existingShare.HostPath()).Debug("unmapping Plan9 share from guest")

internal/controller/device/plan9/controller_test.go

@@ -274,7 +274,7 @@ func TestMapToGuest_GuestMountFails_RetryMapToGuest_Fails(t *testing.T) {
 
 	// Forward retry: MapToGuest again with the same reservation.
 	// AddToVM is idempotent (share already in StateAdded), but MountToGuest
-	// fails because the mount is in the terminal StateUnmounted.
+	// fails because the mount is in StateInvalid.
 	_, err = tc.c.MapToGuest(tc.ctx, id)
 	if err == nil {
 		t.Fatal("expected error on retry MapToGuest after terminal mount failure")
@@ -525,3 +525,134 @@ func TestFullLifecycle_ReuseAfterRelease(t *testing.T) {
 		t.Error("expected a new reservation ID after re-reserving a released path")
 	}
 }
+
+// TestUnmapFromGuest_AddToVMFails_MultipleReservations_AllDrain verifies that
+// when two callers reserve the same host path and AddToVM fails, the share
+// stays in the controller's map (in StateInvalid) until both callers call
+// UnmapFromGuest to drain their mount reservations. Only after the last ref
+// is drained does the share transition to StateRemoved and get cleaned up.
+func TestUnmapFromGuest_AddToVMFails_MultipleReservations_AllDrain(t *testing.T) {
+	t.Parallel()
+	tc := newTestController(t, false)
+
+	// Two callers reserve the same host path.
+	id1, _, _ := tc.c.Reserve(tc.ctx, share.Config{HostPath: "/host/path"}, mount.Config{})
+	id2, _, _ := tc.c.Reserve(tc.ctx, share.Config{HostPath: "/host/path"}, mount.Config{})
+
+	// First caller attempts MapToGuest — AddToVM fails.
+	tc.vmAdd.EXPECT().AddPlan9(gomock.Any(), gomock.Any()).Return(errVMAdd)
+	_, err := tc.c.MapToGuest(tc.ctx, id1)
+	if err == nil {
+		t.Fatal("expected error when VM add fails")
+	}
+
+	// First caller's UnmapFromGuest: decrements mount ref but share stays
+	// (mount ref > 0 → share stays in StateInvalid).
+	if err := tc.c.UnmapFromGuest(tc.ctx, id1); err != nil {
+		t.Fatalf("first UnmapFromGuest: %v", err)
+	}
+
+	// Share must still be in the map — second caller hasn't drained yet.
+	if len(tc.c.sharesByHostPath) != 1 {
+		t.Errorf("expected share to remain in map after first unmap, got %d shares", len(tc.c.sharesByHostPath))
+	}
+	if len(tc.c.reservations) != 1 {
+		t.Errorf("expected 1 remaining reservation, got %d", len(tc.c.reservations))
+	}
+
+	// Second caller's UnmapFromGuest: drains last mount ref → share transitions
+	// to StateRemoved → controller cleans up the share entry.
+	if err := tc.c.UnmapFromGuest(tc.ctx, id2); err != nil {
+		t.Fatalf("second UnmapFromGuest: %v", err)
+	}
+
+	// Everything should be cleaned up.
+	if len(tc.c.reservations) != 0 {
+		t.Errorf("expected 0 reservations after all unmaps, got %d", len(tc.c.reservations))
+	}
+	if len(tc.c.sharesByHostPath) != 0 {
+		t.Errorf("expected 0 shares after all unmaps, got %d", len(tc.c.sharesByHostPath))
+	}
+}
+
+// TestUnmapFromGuest_GuestMountFails_MultipleReservations_AllDrain verifies
+// that when two callers reserve the same host path and MapToGuest fails at the
+// guest mount stage (AddToVM succeeds, AddLCOWMappedDirectory fails), the share
+// and mount stay in the controller's maps until all callers have called
+// UnmapFromGuest to drain their mount reservations.
+func TestUnmapFromGuest_GuestMountFails_MultipleReservations_AllDrain(t *testing.T) {
+	t.Parallel()
+	tc := newTestController(t, false)
+
+	// Two callers reserve the same host path.
+	id1, _, _ := tc.c.Reserve(tc.ctx, share.Config{HostPath: "/host/path"}, mount.Config{})
+	id2, _, _ := tc.c.Reserve(tc.ctx, share.Config{HostPath: "/host/path"}, mount.Config{})
+
+	// First caller attempts MapToGuest — AddToVM succeeds, guest mount fails.
+	tc.vmAdd.EXPECT().AddPlan9(gomock.Any(), gomock.Any()).Return(nil)
+	tc.guestMount.EXPECT().AddLCOWMappedDirectory(gomock.Any(), gomock.Any()).Return(errMount)
+	_, err := tc.c.MapToGuest(tc.ctx, id1)
+	if err == nil {
+		t.Fatal("expected error when guest mount fails")
+	}
+
+	// First caller's UnmapFromGuest: decrements mount ref but share stays
+	// because second caller still holds a reservation.
+	// VM remove should NOT be called yet — second caller hasn't drained.
+	tc.vmRemove.EXPECT().RemovePlan9(gomock.Any(), gomock.Any()).Return(nil).MaxTimes(0)
+	if err := tc.c.UnmapFromGuest(tc.ctx, id1); err != nil {
+		t.Fatalf("first UnmapFromGuest: %v", err)
+	}
+
+	// Share must still be in the map — second caller hasn't drained yet.
+	if len(tc.c.sharesByHostPath) != 1 {
+		t.Errorf("expected share to remain in map after first unmap, got %d shares", len(tc.c.sharesByHostPath))
+	}
+	if len(tc.c.reservations) != 1 {
+		t.Errorf("expected 1 remaining reservation, got %d", len(tc.c.reservations))
+	}
+
+	// Second caller's UnmapFromGuest: drains last mount ref → share transitions
+	// to StateRemoved → controller cleans up the share entry.
+	tc.vmRemove.EXPECT().RemovePlan9(gomock.Any(), gomock.Any()).Return(nil)
+	if err := tc.c.UnmapFromGuest(tc.ctx, id2); err != nil {
+		t.Fatalf("second UnmapFromGuest: %v", err)
+	}
+
+	// Everything should be cleaned up.
+	if len(tc.c.reservations) != 0 {
+		t.Errorf("expected 0 reservations after all unmaps, got %d", len(tc.c.reservations))
+	}
+	if len(tc.c.sharesByHostPath) != 0 {
+		t.Errorf("expected 0 shares after all unmaps, got %d", len(tc.c.sharesByHostPath))
+	}
+}
+
+// TestUnmapFromGuest_AddToVMFails_SingleReservation_Drains verifies that when
+// a single caller reserves a host path and AddToVM fails, UnmapFromGuest
+// correctly drains the mount and transitions the share to StateRemoved.
+func TestUnmapFromGuest_AddToVMFails_SingleReservation_Drains(t *testing.T) {
+	t.Parallel()
+	tc := newTestController(t, false)
+
+	id, _, _ := tc.c.Reserve(tc.ctx, share.Config{HostPath: "/host/path"}, mount.Config{})
+
+	// MapToGuest fails at AddToVM.
+	tc.vmAdd.EXPECT().AddPlan9(gomock.Any(), gomock.Any()).Return(errVMAdd)
+	_, err := tc.c.MapToGuest(tc.ctx, id)
+	if err == nil {
+		t.Fatal("expected error when VM add fails")
+	}
+
+	// UnmapFromGuest drains the mount and removes the share.
+	if err := tc.c.UnmapFromGuest(tc.ctx, id); err != nil {
+		t.Fatalf("UnmapFromGuest: %v", err)
+	}
+
+	if len(tc.c.reservations) != 0 {
+		t.Errorf("expected 0 reservations, got %d", len(tc.c.reservations))
+	}
+	if len(tc.c.sharesByHostPath) != 0 {
+		t.Errorf("expected 0 shares, got %d", len(tc.c.sharesByHostPath))
+	}
+}

internal/controller/device/plan9/mount/doc.go

@@ -16,15 +16,15 @@
 // # Mount Lifecycle
 //
 //	┌──────────────────────┐
-//	│    StateReserved     │ ← mount failure → StateUnmounted (not retriable)
-//	└──────────┬───────────┘
-//	           │ guest mount succeeds
-//	           ▼
-//	┌──────────────────────┐
-//	│    StateMounted      │
-//	└──────────┬───────────┘
-//	           │ reference count → 0;
-//	           │ guest unmount succeeds
+//	│    StateReserved     │ ← mount failure → StateInvalid
+//	└──────────┬───────────┘                       │
+//	           │ guest mount succeeds              │ all refs drained
+//	           ▼                                   ▼
+//	┌──────────────────────┐              ┌──────────────────────┐
+//	│    StateMounted      │              │   StateUnmounted     │
+//	└──────────┬───────────┘              └──────────────────────┘
+//	           │ reference count → 0;       (terminal — entry
+//	           │ guest unmount succeeds      removed from share)
 //	           ▼
 //	┌──────────────────────┐
 //	│   StateUnmounted     │

internal/controller/device/plan9/mount/mount.go

@@ -100,8 +100,9 @@ func (m *Mount) MountToGuest(ctx context.Context, guest LinuxGuestPlan9Mounter)
 			Port:      vmutils.Plan9Port,
 			ReadOnly:  m.config.ReadOnly,
 		}); err != nil {
-			// Move to unmounted since no guest state was established from Reserved.
-			m.state = StateUnmounted
+			// Move to invalid since no guest state was established from Reserved,
+			// but outstanding reservations may still need to be drained.
+			m.state = StateInvalid
 			return "", fmt.Errorf("add LCOW mapped directory share=%s: %w", m.shareName, err)
 		}
 
@@ -117,10 +118,9 @@ func (m *Mount) MountToGuest(ctx context.Context, guest LinuxGuestPlan9Mounter)
 		// existing guest path directly.
 		return m.guestPath, nil
 
-	case StateUnmounted:
+	default:
 		return "", fmt.Errorf("cannot mount a share in state %s", m.state)
 	}
-	return "", nil
 }
 
 // UnmountFromGuest decrements the reference count and, when it reaches zero,
@@ -160,10 +160,18 @@ func (m *Mount) UnmountFromGuest(ctx context.Context, guest LinuxGuestPlan9Unmou
 		m.refCount--
 		return nil
 
+	case StateInvalid:
+		// The guest mount failed; no guest-side state to tear down.
+		// Decrement the ref count and transition to the terminal state
+		// once all reservations are drained.
+		m.refCount--
+		if m.refCount == 0 {
+			m.state = StateUnmounted
+		}
+		return nil
+
 	case StateUnmounted:
-		// Already in the terminal state — nothing to do. This can happen when
-		// MountToGuest failed (it transitions StateReserved → StateUnmounted),
-		// and the caller subsequently calls UnmapFromGuest to clean up.
+		// Already in the terminal state — nothing to do.
 	}
 	return nil
 }

internal/controller/device/plan9/mount/mount_test.go

@@ -118,10 +118,10 @@ func TestMountToGuest_HappyPath(t *testing.T) {
 	}
 }
 
-// TestMountToGuest_GuestFails_TransitionsToUnmounted verifies that when the
+// TestMountToGuest_GuestFails_TransitionsToInvalid verifies that when the
 // guest AddLCOWMappedDirectory call fails, the mount transitions directly from
-// StateReserved to StateUnmounted (the terminal state).
-func TestMountToGuest_GuestFails_TransitionsToUnmounted(t *testing.T) {
+// StateReserved to StateInvalid so outstanding reservations can be drained.
+func TestMountToGuest_GuestFails_TransitionsToInvalid(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	guest := mocks.NewMockLinuxGuestPlan9Mounter(ctrl)
 
@@ -133,8 +133,8 @@ func TestMountToGuest_GuestFails_TransitionsToUnmounted(t *testing.T) {
 	if err == nil {
 		t.Fatal("expected error")
 	}
-	if m.State() != StateUnmounted {
-		t.Errorf("expected StateUnmounted after mount failure, got %v", m.State())
+	if m.State() != StateInvalid {
+		t.Errorf("expected StateInvalid after mount failure, got %v", m.State())
 	}
 }
 
@@ -165,20 +165,20 @@ func TestMountToGuest_AlreadyMounted_Idempotent(t *testing.T) {
 	}
 }
 
-// TestMountToGuest_OnUnmounted_Errors verifies that calling MountToGuest on a
-// mount in the terminal StateUnmounted returns an error.
-func TestMountToGuest_OnUnmounted_Errors(t *testing.T) {
+// TestMountToGuest_OnInvalid_Errors verifies that calling MountToGuest on a
+// mount in StateInvalid (from a prior mount failure) returns an error.
+func TestMountToGuest_OnInvalid_Errors(t *testing.T) {
 	ctrl := gomock.NewController(t)
 	guest := mocks.NewMockLinuxGuestPlan9Mounter(ctrl)
 
 	m := NewReserved("share0", Config{})
 	guest.EXPECT().AddLCOWMappedDirectory(gomock.Any(), gomock.Any()).Return(errGuestMount)
 	_, _ = m.MountToGuest(context.Background(), guest)
 
-	// Try to mount again from StateUnmounted.
+	// Try to mount again from StateInvalid.
 	_, err := m.MountToGuest(context.Background(), guest)
 	if err == nil {
-		t.Fatal("expected error when mounting from StateUnmounted")
+		t.Fatal("expected error when mounting from StateInvalid")
 	}
 }
 

internal/controller/device/plan9/mount/state.go

@@ -9,8 +9,11 @@ package mount
 //
 //	StateReserved → StateMounted → StateUnmounted
 //
-// Mount failure from StateReserved transitions directly to the terminal
-// StateUnmounted state; the entry is then removed by the parent [share.Share].
+// Mount failure from StateReserved transitions to [StateInvalid].
+// In [StateInvalid] the guest mount was never established, but outstanding
+// reservations may still exist. The mount remains in [StateInvalid] until
+// all reservations have been drained via [Mount.UnmountFromGuest], at which
+// point it transitions to [StateUnmounted].
 // An unmount failure from StateMounted leaves the mount in StateMounted so
 // the caller can retry.
 //
@@ -19,12 +22,14 @@ package mount
 //	Current State   │ Trigger                                    │ Next State
 //	────────────────┼────────────────────────────────────────────┼──────────────────────
 //	StateReserved   │ guest mount succeeds                       │ StateMounted
-//	StateReserved   │ guest mount fails                          │ StateUnmounted
+//	StateReserved   │ guest mount fails                          │ StateInvalid
 //	StateReserved   │ UnmountFromGuest (refCount > 1)            │ StateReserved (ref--)
 //	StateReserved   │ UnmountFromGuest (refCount == 1)           │ StateUnmounted
 //	StateMounted    │ UnmountFromGuest (refCount > 1)            │ StateMounted (ref--)
 //	StateMounted    │ UnmountFromGuest (refCount == 1) succeeds  │ StateUnmounted
 //	StateMounted    │ UnmountFromGuest (refCount == 1) fails     │ StateMounted
+//	StateInvalid    │ UnmountFromGuest (refCount > 1)            │ StateInvalid (ref--)
+//	StateInvalid    │ UnmountFromGuest (refCount == 1)           │ StateUnmounted
 //	StateUnmounted  │ UnmountFromGuest                           │ StateUnmounted (no-op)
 //	StateUnmounted  │ (terminal — entry removed from share)      │ —
 type State int
@@ -38,6 +43,13 @@ const (
 	// the guest. The guest path is valid from this state onward.
 	StateMounted
 
+	// StateInvalid means the guest mount operation failed. No guest-side
+	// state was established, but outstanding reservations may still need
+	// to be drained via [Mount.UnmountFromGuest]. Once all reservations
+	// are released (refCount reaches zero), the mount transitions to
+	// [StateUnmounted].
+	StateInvalid
+
 	// StateUnmounted means the guest has unmounted the share. This is a
 	// terminal state; the entry is removed from the parent share.
 	StateUnmounted
@@ -50,6 +62,8 @@ func (s State) String() string {
 		return "Reserved"
 	case StateMounted:
 		return "Mounted"
+	case StateInvalid:
+		return "Invalid"
 	case StateUnmounted:
 		return "Unmounted"
 	default:

internal/controller/device/plan9/share/doc.go

@@ -15,7 +15,7 @@
 // # Share Lifecycle
 //
 //	┌──────────────────────┐
-//	│    StateReserved     │ ← add failure → StateRemoved (not retriable)
+//	│    StateReserved     │ ← add failure → StateInvalid
 //	└──────────┬───────────┘
 //	           │ share added to VM
 //	           ▼
@@ -30,4 +30,15 @@
 //	│    StateRemoved      │
 //	└──────────────────────┘
 //	  (terminal — entry removed from controller)
+//
+//	┌──────────────────────┐
+//	│    StateInvalid      │ ← add failed; share never on VM
+//	└──────────┬───────────┘
+//	           │ all mount reservations drained
+//	           │ via UnmountFromGuest + RemoveFromVM
+//	           ▼
+//	┌──────────────────────┐
+//	│    StateRemoved      │
+//	└──────────────────────┘
+//	  (terminal — entry removed from controller)
 package share