fix rare race condition in nvme disconnect subsystem

Author: Utkarshjh

core/node/detach.go

@@ -660,15 +660,6 @@ func (c *Core) detachNVMeVolume(
 	defer nvmeSelfHealingLock.RUnlock()
 	nvmeNodeOperationWaitingCount.Add(-1)
 
-	lockContext := "detachNVMeVolume.RemovePublishedNVMeSession"
-	if !attemptLock(ctx, lockContext, nvmeSelfHealingSessionLock, sharedLocksNodeLockTimeout) {
-		locks.Unlock(ctx, lockContext, nvmeSelfHealingSessionLock)
-		return errors.MaxWaitExceededError("request waited too long for the lock")
-	}
-	disconnect := c.nvme.RemovePublishedNVMeSession(&publishedNVMeSessions, publishInfo.NVMeSubsystemNQN,
-		publishInfo.NVMeNamespaceUUID)
-	locks.Unlock(ctx, lockContext, nvmeSelfHealingSessionLock)
-
 	nvmeSubsys := c.nvme.NewNVMeSubsystem(ctx, publishInfo.NVMeSubsystemNQN)
 	// Get the device using 'nvme-cli' commands. Flush the device IOs.
 	// Proceed further with detach flow, if device is not found.
@@ -749,8 +740,17 @@ func (c *Core) detachNVMeVolume(
 		locks.Unlock(ctx, "detachNVMeVolume.FlushRetryDelete", nvmeFlushRetryMapLock)
 	}
 
+	lockContext := "detachNVMeVolume.RemovePublishedNVMeSession"
+	if !attemptLock(ctx, lockContext, nvmeSelfHealingSessionLock, sharedLocksNodeLockTimeout) {
+		locks.Unlock(ctx, lockContext, nvmeSelfHealingSessionLock)
+		return errors.MaxWaitExceededError("request waited too long for the lock")
+	}
+	c.nvme.RemovePublishedNVMeSession(&publishedNVMeSessions, publishInfo.NVMeSubsystemNQN,
+		publishInfo.NVMeNamespaceUUID)
+	locks.Unlock(ctx, lockContext, nvmeSelfHealingSessionLock)
+
 	// Disconnect the subsystem if needed (handled under lock to prevent race conditions).
-	if err = c.disconnectNVMeSubsystemIfNeeded(ctx, nvmeSubsys, publishInfo, disconnect); err != nil {
+	if err = c.disconnectNVMeSubsystemIfNeeded(ctx, nvmeSubsys, publishInfo); err != nil {
 		Logc(ctx).WithError(err).Warn("Error during subsystem disconnect check.")
 		// Continue with cleanup even if disconnect fails.
 	}

core/node/detach_test.go

@@ -751,13 +751,22 @@ type fakeDetachNVMeSubsystem struct {
 	getDeviceErr    error
 	disconnectErr   error
 	disconnectCalls int
+	hostNsCount     int
+	hostNsCountErr  error
 }
 
 func (f *fakeDetachNVMeSubsystem) Disconnect(context.Context) error {
 	f.disconnectCalls++
 	return f.disconnectErr
 }
 
+// GetNamespaceCount feeds disconnectNVMeSubsystemIfNeeded's host-level check, which only skips the
+// disconnect when the count exceeds one; the zero value therefore leaves the disconnect decision to
+// the published session count that each test below sets up.
+func (f *fakeDetachNVMeSubsystem) GetNamespaceCount(context.Context) (int, error) {
+	return f.hostNsCount, f.hostNsCountErr
+}
+
 // GetNVMeDevice/GetNVMeDeviceAt always report "no device found" (nil, getDeviceErr): every
 // detachNVMeVolume test below therefore exercises the nvmeDev==nil branches, which is sufficient
 // to cover subsystem lookup, LUKS teardown, and disconnect-decision logic without needing a real,
@@ -796,7 +805,8 @@ func TestDetachNVMeVolume_GetNVMeDeviceNonNotFoundError(t *testing.T) {
 	publishInfo := samplePublishInfo(NVMe)
 	subsystem := &fakeDetachNVMeSubsystem{getDeviceErr: errors.New("nvme-cli failed")}
 
-	mocks.NVMe.EXPECT().RemovePublishedNVMeSession(gomock.Any(), gomock.Any(), gomock.Any()).Return(false)
+	// No RemovePublishedNVMeSession expectation: the session is now removed only after the device is
+	// torn down, so bailing out here leaves the published session in place for a later retry.
 	mocks.NVMe.EXPECT().NewNVMeSubsystem(gomock.Any(), gomock.Any()).Return(subsystem)
 
 	err := core.detachNVMeVolume(context.Background(), "test-volume", publishInfo, false)
@@ -846,7 +856,8 @@ func TestDetachNVMeVolume_LUKSDevicePathLookupError(t *testing.T) {
 	publishInfo.LUKSEncryption = "true"
 	subsystem := &fakeDetachNVMeSubsystem{}
 
-	mocks.NVMe.EXPECT().RemovePublishedNVMeSession(gomock.Any(), gomock.Any(), gomock.Any()).Return(false)
+	// No RemovePublishedNVMeSession expectation: this failure precedes the session removal, which now
+	// happens only after LUKS teardown completes.
 	mocks.NVMe.EXPECT().NewNVMeSubsystem(gomock.Any(), gomock.Any()).Return(subsystem)
 	mocks.Devices.EXPECT().GetLUKSDevicePathForDevicePath(gomock.Any(), publishInfo.DevicePath).
 		Return("", errors.New("lookup failed"))

core/node/utils.go

@@ -185,7 +185,7 @@ func (c *Core) readAllTrackingFiles(ctx context.Context) []models.VolumePublishI
 // This lock serializes GetNamespaceCount() and Disconnect() operations to ensure accurate namespace counting
 // and prevent race conditions where multiple threads might see the same count simultaneously.
 func (c *Core) disconnectNVMeSubsystemIfNeeded(
-	ctx context.Context, nvmeSubsys nvme.NVMeSubsystemInterface, publishInfo *models.VolumePublishInfo, disconnect bool,
+	ctx context.Context, nvmeSubsys nvme.NVMeSubsystemInterface, publishInfo *models.VolumePublishInfo,
 ) error {
 	lockContext := "disconnectNVMeSubsystemIfNeeded"
 	if !attemptLock(ctx, lockContext, nvmeSubsystemDisconnectLock, sharedLocksNodeLockTimeout) {
@@ -194,20 +194,44 @@ func (c *Core) disconnectNVMeSubsystemIfNeeded(
 	}
 	defer locks.Unlock(ctx, lockContext, nvmeSubsystemDisconnectLock)
 
+	// publishedNVMeSessions is mutated (Add/Remove) under nvmeSelfHealingSessionLock so this read must
+	// take that same lock to avoid a concurrent map read/write with NodeStage, NodeUnstage, or self-healing.
+	sessionLockContext := "disconnectNVMeSubsystemIfNeeded.SessionRead"
+	if !attemptLock(ctx, sessionLockContext, nvmeSelfHealingSessionLock, sharedLocksNodeLockTimeout) {
+		locks.Unlock(ctx, sessionLockContext, nvmeSelfHealingSessionLock)
+		return errors.MaxWaitExceededError("request waited too long for the lock")
+	}
 	numNs := publishedNVMeSessions.GetNamespaceCountForSession(publishInfo.NVMeSubsystemNQN)
+	locks.Unlock(ctx, sessionLockContext, nvmeSelfHealingSessionLock)
 	Logc(ctx).WithFields(LogFields{
 		"subsystem":      publishInfo.NVMeSubsystemNQN,
 		"namespaceCount": numNs,
-		"disconnectFlag": disconnect,
 	}).Info("Checking if subsystem should be disconnected.")
 
-	if numNs == 0 || (c.nvmeSelfHealingInterval > 0 && disconnect) {
-		if err := nvmeSubsys.Disconnect(ctx); err != nil {
-			Logc(ctx).WithField(
-				"subsystem", publishInfo.NVMeSubsystemNQN,
-			).WithError(err).Debug("Error disconnecting subsystem.")
-			return err
-		}
+	// Another pod still has a published session; we must not disconnect.
+	if numNs > 0 {
+		return nil
+	}
+
+	// In-memory sessions show none left, but a concurrent NodeStage may have already attached a namespace
+	// for a new pod without recording its session yet (recorded only after format/mount). Checking the
+	// host's ground-truth namespace count; if it's >1, another namespace is active, we don't disconnect.
+	if hostNsCount, err := nvmeSubsys.GetNamespaceCount(ctx); err != nil {
+		Logc(ctx).WithField("subsystem", publishInfo.NVMeSubsystemNQN).WithError(err).Debug(
+			"Could not determine host namespace count; proceeding with disconnect based on published sessions.")
+	} else if hostNsCount > 1 {
+		Logc(ctx).WithFields(LogFields{
+			"subsystem":      publishInfo.NVMeSubsystemNQN,
+			"hostNamespaces": hostNsCount,
+		}).Info("Subsystem still has namespace devices attached on host; skipping disconnect.")
+		return nil
+	}
+
+	if err := nvmeSubsys.Disconnect(ctx); err != nil {
+		Logc(ctx).WithField(
+			"subsystem", publishInfo.NVMeSubsystemNQN,
+		).WithError(err).Debug("Error disconnecting subsystem.")
+		return err
 	}
 	return nil
 }

core/node/utils_test.go

@@ -351,20 +351,26 @@ func TestReadAllTrackingFiles_NoTrackingDirReturnsEmpty(t *testing.T) {
 
 // fakeNVMeSubsystem is a minimal hand-rolled fake for nvme.NVMeSubsystemInterface: no gomock
 // mock exists for this interface under mocks/mock_utils/nvme, only for NVMeInterface. Any method
-// besides Disconnect is intentionally left unimplemented (nil-embedded) since
+// besides Disconnect and GetNamespaceCount is intentionally left unimplemented (nil-embedded) since
 // disconnectNVMeSubsystemIfNeeded never calls them; using any of them would panic, which is the
 // desired failure mode if the code under test changes to call something unexpected.
 type fakeNVMeSubsystem struct {
 	nvme.NVMeSubsystemInterface
 	disconnectErr   error
 	disconnectCalls int
+	hostNsCount     int
+	hostNsCountErr  error
 }
 
 func (f *fakeNVMeSubsystem) Disconnect(_ context.Context) error {
 	f.disconnectCalls++
 	return f.disconnectErr
 }
 
+func (f *fakeNVMeSubsystem) GetNamespaceCount(_ context.Context) (int, error) {
+	return f.hostNsCount, f.hostNsCountErr
+}
+
 // withCleanPublishedNVMeSessions snapshots and restores the package-level publishedNVMeSessions
 // global so this test's seeding doesn't bleed into other tests in the package.
 func withCleanPublishedNVMeSessions(t *testing.T) {
@@ -373,59 +379,56 @@ func withCleanPublishedNVMeSessions(t *testing.T) {
 	t.Cleanup(func() { publishedNVMeSessions = original })
 }
 
-func TestDisconnectNVMeSubsystemIfNeeded_NoNamespaces_DisconnectsRegardlessOfFlag(t *testing.T) {
-	withCleanPublishedNVMeSessions(t)
-	core, _ := newTestCore(t)
-	pi := samplePublishInfo(NVMe)
-	fakeSubsys := &fakeNVMeSubsystem{}
-
-	err := core.disconnectNVMeSubsystemIfNeeded(context.Background(), fakeSubsys, pi, false)
-
-	require.NoError(t, err)
-	assert.Equal(t, 1, fakeSubsys.disconnectCalls)
-}
-
-func TestDisconnectNVMeSubsystemIfNeeded_NamespacesPresent_DisconnectFlagFalse_NoDisconnect(t *testing.T) {
+func TestDisconnectNVMeSubsystemIfNeeded_PublishedNamespacePresent_NoDisconnect(t *testing.T) {
 	withCleanPublishedNVMeSessions(t)
 	core, _ := newTestCore(t, WithNVMeSelfHealingInterval(5*time.Second))
 	pi := samplePublishInfo(NVMe)
 	publishedNVMeSessions.AddNVMeSession(nvme.NVMeSubsystem{NQN: pi.NVMeSubsystemNQN}, nil)
 	publishedNVMeSessions.AddNamespaceToSession(pi.NVMeSubsystemNQN, "ns-1")
-	fakeSubsys := &fakeNVMeSubsystem{}
+	fakeSubsys := &fakeNVMeSubsystem{hostNsCount: 1}
 
-	err := core.disconnectNVMeSubsystemIfNeeded(context.Background(), fakeSubsys, pi, false)
+	err := core.disconnectNVMeSubsystemIfNeeded(context.Background(), fakeSubsys, pi)
 
 	require.NoError(t, err)
 	assert.Equal(t, 0, fakeSubsys.disconnectCalls)
 }
 
-func TestDisconnectNVMeSubsystemIfNeeded_NamespacesPresent_SelfHealingDisabled_NoDisconnect(t *testing.T) {
-	withCleanPublishedNVMeSessions(t)
-	// nvmeSelfHealingInterval defaults to zero (disabled) when not set via WithNVMeSelfHealingInterval.
-	core, _ := newTestCore(t)
-	pi := samplePublishInfo(NVMe)
-	publishedNVMeSessions.AddNVMeSession(nvme.NVMeSubsystem{NQN: pi.NVMeSubsystemNQN}, nil)
-	publishedNVMeSessions.AddNamespaceToSession(pi.NVMeSubsystemNQN, "ns-1")
-	fakeSubsys := &fakeNVMeSubsystem{}
-
-	err := core.disconnectNVMeSubsystemIfNeeded(context.Background(), fakeSubsys, pi, true)
-
-	require.NoError(t, err)
-	assert.Equal(t, 0, fakeSubsys.disconnectCalls, "self-healing disabled must gate the disconnect hint even if disconnect=true")
-}
+// Once no published sessions remain, the host namespace count is the tie-breaker: a count above one
+// means a concurrent NodeStage already attached a namespace it hasn't recorded a session for yet, so
+// disconnecting would pull that device out from under the new pod. Any other outcome, including an
+// unreadable count, falls through to the disconnect.
+func TestDisconnectNVMeSubsystemIfNeeded_NoPublishedNamespaces_HostCountDecides(t *testing.T) {
+	tests := map[string]struct {
+		hostNsCount         int
+		hostNsCountErr      error
+		wantDisconnectCalls int
+	}{
+		"another namespace still attached on host": {hostNsCount: 2, wantDisconnectCalls: 0},
+		"several namespaces still attached":        {hostNsCount: 5, wantDisconnectCalls: 0},
+		"only our own namespace attached":          {hostNsCount: 1, wantDisconnectCalls: 1},
+		"no namespaces attached":                   {hostNsCount: 0, wantDisconnectCalls: 1},
+		"host count unreadable": {
+			hostNsCountErr:      errors.New("failed to read namespace count"),
+			wantDisconnectCalls: 1,
+		},
+	}
 
-func TestDisconnectNVMeSubsystemIfNeeded_NamespacesPresent_SelfHealingEnabledAndDisconnect_Disconnects(t *testing.T) {
-	withCleanPublishedNVMeSessions(t)
-	core, _ := newTestCore(t, WithNVMeSelfHealingInterval(5*time.Second))
-	pi := samplePublishInfo(NVMe)
-	publishedNVMeSessions.AddNVMeSession(nvme.NVMeSubsystem{NQN: pi.NVMeSubsystemNQN}, nil)
-	publishedNVMeSessions.AddNamespaceToSession(pi.NVMeSubsystemNQN, "ns-1")
-	fakeSubsys := &fakeNVMeSubsystem{}
+	for name, test := range tests {
+		t.Run(name, func(t *testing.T) {
+			withCleanPublishedNVMeSessions(t)
+			core, _ := newTestCore(t)
+			pi := samplePublishInfo(NVMe)
+			fakeSubsys := &fakeNVMeSubsystem{
+				hostNsCount:    test.hostNsCount,
+				hostNsCountErr: test.hostNsCountErr,
+			}
 
-	err := core.disconnectNVMeSubsystemIfNeeded(context.Background(), fakeSubsys, pi, true)
+			err := core.disconnectNVMeSubsystemIfNeeded(context.Background(), fakeSubsys, pi)
 
-	require.NoError(t, err)
-	assert.Equal(t, 1, fakeSubsys.disconnectCalls)
+			require.NoError(t, err)
+			assert.Equal(t, test.wantDisconnectCalls, fakeSubsys.disconnectCalls)
+		})
+	}
 }
 
 func TestDisconnectNVMeSubsystemIfNeeded_DisconnectErrorPropagates(t *testing.T) {
@@ -434,7 +437,7 @@ func TestDisconnectNVMeSubsystemIfNeeded_DisconnectErrorPropagates(t *testing.T)
 	pi := samplePublishInfo(NVMe)
 	fakeSubsys := &fakeNVMeSubsystem{disconnectErr: errors.New("disconnect failed")}
 
-	err := core.disconnectNVMeSubsystemIfNeeded(context.Background(), fakeSubsys, pi, false)
+	err := core.disconnectNVMeSubsystemIfNeeded(context.Background(), fakeSubsys, pi)
 
 	require.Error(t, err)
 	assert.Contains(t, err.Error(), "disconnect failed")