Consolidate pod-level state into podConfigStore

Author: purvavj

pkg/driver/driver.go

@@ -64,9 +64,6 @@ type inventoryDB interface {
 	IsIBOnlyDevice(deviceName string) bool
 	GetRDMADeviceName(deviceName string) (string, error)
 	GetDeviceConfig(deviceName string) (*apis.NetworkConfig, bool)
-	AddPodNetNs(podKey string, netNs string)
-	RemovePodNetNs(podKey string)
-	GetPodNetNs(podKey string) (netNs string)
 	RequestRescan()
 }
 

pkg/driver/driver_test.go

@@ -47,14 +47,12 @@ func (m *fakePluginHelper) RegistrationStatus() *registerapi.RegistrationStatus
 // mockNetDB is a mock implementation of the inventoryDB interface for testing.
 type fakeInventoryDB struct {
 	resources   chan []resourcev1.Device
-	podNetNs    map[string]string
 	rescanCalls atomic.Int32
 }
 
 func newFakeInventoryDB() *fakeInventoryDB {
 	return &fakeInventoryDB{
 		resources: make(chan []resourcev1.Device, 1),
-		podNetNs:  make(map[string]string),
 	}
 }
 
@@ -70,18 +68,6 @@ func (m *fakeInventoryDB) IsIBOnlyDevice(_ string) bool { return false }
 
 func (m *fakeInventoryDB) GetRDMADeviceName(_ string) (string, error) { return "", nil }
 
-func (m *fakeInventoryDB) AddPodNetNs(podKey string, netNs string) {
-	m.podNetNs[podKey] = netNs
-}
-
-func (m *fakeInventoryDB) RemovePodNetNs(podKey string) {
-	delete(m.podNetNs, podKey)
-}
-
-func (m *fakeInventoryDB) GetPodNetNs(podKey string) string {
-	return m.podNetNs[podKey]
-}
-
 func (m *fakeInventoryDB) GetDeviceConfig(deviceName string) (*apis.NetworkConfig, bool) {
 	return nil, false
 }

pkg/driver/nri_hooks.go

@@ -42,24 +42,24 @@ func (np *NetworkDriver) Synchronize(_ context.Context, pods []*api.PodSandbox,
 	klog.Infof("Synchronized state with the runtime (%d pods, %d containers)...",
 		len(pods), len(containers))
 
-	livePods := set.New[types.UID]()
+	// livePodNetNs map tracks live pods by UID and their network namespace paths.
+	livePodNetNs := make(map[types.UID]string)
 	for _, pod := range pods {
 		klog.Infof("Synchronize Pod %s/%s UID %s", pod.Namespace, pod.Name, pod.Uid)
 		klog.V(2).Infof("pod %s/%s: namespace=%s ips=%v", pod.GetNamespace(), pod.GetName(), getNetworkNamespace(pod), pod.GetIps())
-		livePods.Insert(types.UID(pod.Uid))
-		// get the pod network namespace
-		ns := getNetworkNamespace(pod)
-		// host network pods are skipped
-		if ns != "" {
-			// store the Pod metadata in the db
-			np.netdb.AddPodNetNs(podKey(pod), ns)
-		}
+		livePodNetNs[types.UID(pod.Uid)] = getNetworkNamespace(pod)
 	}
 
-	// Prune persisted configs for pods that no longer exist in the runtime.
-	// This handles the case where pods were deleted while the driver was down.
+	// Process stored pods: update NetNS for live pods, and prune configurations
+	// for pods that no longer exist in the runtime.
 	for _, storedUID := range np.podConfigStore.ListPods() {
-		if !livePods.Has(storedUID) {
+		ns, isLive := livePodNetNs[storedUID]
+
+		if isLive {
+			if ns != "" {
+				np.podConfigStore.SetPodNetNs(storedUID, ns)
+			}
+		} else {
 			klog.Infof("Synchronize: pruning stale config for pod %s", storedUID)
 			np.podConfigStore.DeletePod(storedUID)
 		}
@@ -152,8 +152,8 @@ func (np *NetworkDriver) runPodSandbox(_ context.Context, pod *api.PodSandbox, p
 	if ns == "" {
 		return fmt.Errorf("RunPodSandbox pod %s/%s using host network can not claim host devices", pod.Namespace, pod.Name)
 	}
-	// store the Pod metadata in the db
-	np.netdb.AddPodNetNs(podKey(pod), ns)
+	// store the Pod network namespace in the pod config store
+	np.podConfigStore.SetPodNetNs(types.UID(pod.GetUid()), ns)
 
 	// Track all the status updates needed for the resource claims of the pod.
 	statusUpdates := map[types.NamespacedName]*resourceapply.ResourceClaimStatusApplyConfiguration{}
@@ -363,18 +363,23 @@ func (np *NetworkDriver) StopPodSandbox(ctx context.Context, pod *api.PodSandbox
 }
 
 func (np *NetworkDriver) stopPodSandbox(_ context.Context, pod *api.PodSandbox, podConfig PodConfig) error {
-	defer func() {
-		np.netdb.RemovePodNetNs(podKey(pod))
-	}()
 	// get the pod network namespace
 	ns := getNetworkNamespace(pod)
 	if ns == "" {
 		// some version of containerd does not send the network namespace information on this hook so
 		// we workaround it using the local copy we have in the db to associate interfaces with Pods via
 		// the network namespace id.
-		ns = np.netdb.GetPodNetNs(podKey(pod))
-		if ns == "" {
-			klog.Infof("StopPodSandbox pod %s/%s using host network ... skipping", pod.Namespace, pod.Name)
+		storedNs, ok := np.podConfigStore.GetPodNetNs(types.UID(pod.GetUid()))
+		if ok {
+			if storedNs != "" {
+				ns = storedNs
+			} else {
+				// Pod is not configured for DRAnet (host network or other driver)
+				klog.Infof("StopPodSandbox pod %s/%s using host network ... skipping", pod.Namespace, pod.Name)
+				return nil
+			}
+		} else {
+			klog.Warningf("StopPodSandbox: pod %s/%s (UID %s) not found in podConfigStore when fetching fallback NetNS", pod.Namespace, pod.Name, pod.Uid)
 			return nil
 		}
 	}
@@ -452,7 +457,6 @@ func (np *NetworkDriver) RemovePodSandbox(ctx context.Context, pod *api.PodSandb
 }
 
 func (np *NetworkDriver) removePodSandbox(_ context.Context, pod *api.PodSandbox) error {
-	np.netdb.RemovePodNetNs(podKey(pod))
 	return nil
 }
 

pkg/driver/nri_hooks_test.go

@@ -225,6 +225,69 @@ func TestSynchronizePrunesStaleConfigs(t *testing.T) {
 	}
 }
 
+func TestSynchronizeStoresNetNSOnlyForConfiguredPods(t *testing.T) {
+	store := mustNewPodConfigStore()
+
+	// Pod 1: Has device config (configured)
+	store.SetDeviceConfig("configured-pod", "eth0", DeviceConfig{}) //nolint:errcheck
+
+	// Pod 2: Does not have device config (unconfigured)
+
+	np := &NetworkDriver{
+		podConfigStore: store,
+		netdb:          inventory.New(),
+	}
+
+	pods := []*api.PodSandbox{
+		{
+			Uid:       "configured-pod",
+			Name:      "configured",
+			Namespace: "default",
+			Linux: &api.LinuxPodSandbox{
+				Namespaces: []*api.LinuxNamespace{
+					{Type: "network", Path: "/var/run/netns/configured"},
+				},
+			},
+		},
+		{
+			Uid:       "unconfigured-pod",
+			Name:      "unconfigured",
+			Namespace: "default",
+			Linux: &api.LinuxPodSandbox{
+				Namespaces: []*api.LinuxNamespace{
+					{Type: "network", Path: "/var/run/netns/unconfigured"},
+				},
+			},
+		},
+	}
+
+	_, err := np.Synchronize(context.Background(), pods, nil)
+	if err != nil {
+		t.Fatalf("Synchronize() error: %v", err)
+	}
+
+	// Case 1: Configured pod should have its NetNS stored
+	netns, found := store.GetPodNetNs("configured-pod")
+	if !found {
+		t.Error("configured-pod should have its NetNS stored")
+	}
+	if netns != "/var/run/netns/configured" {
+		t.Errorf("expected NetNS /var/run/netns/configured, got %q", netns)
+	}
+
+	// Case 2: Unconfigured pod should NOT have its NetNS stored
+	_, found = store.GetPodNetNs("unconfigured-pod")
+	if found {
+		t.Error("unconfigured-pod should NOT have its NetNS stored")
+	}
+
+	// Also verify it didn't create a skeleton config for unconfigured-pod
+	_, found = store.GetPodConfig("unconfigured-pod")
+	if found {
+		t.Error("unconfigured-pod should NOT have any PodConfig in the store")
+	}
+}
+
 func TestCreateContainerMetrics(t *testing.T) {
 	testCases := []struct {
 		name           string

pkg/driver/pod_device_config.go

@@ -35,6 +35,10 @@ type PodConfig struct {
 	// LastNRIActivity timestamp is updated whenever an NRI hook processes
 	// a container for this Pod. Used to track pod initialization progress.
 	LastNRIActivity time.Time
+
+	// NetNS is the path to the Pod's network namespace as observed by the
+	// container runtime.
+	NetNS string `json:"netns,omitempty"`
 }
 
 // DeviceConfig holds the set of configurations to be applied for a single
@@ -87,9 +91,11 @@ type LinuxDevice struct {
 // daemon restarts. Following the kubelet DRA checkpoint pattern
 // (pkg/kubelet/cm/dra/state), the in-memory PodConfigStore is the source
 // of truth and the Checkpointer is a write-through backend.
+// Note: Pod-level metadata like NetNS is not persisted (rebuilt on restart
+// via Synchronize() which queries the container runtime).
 type Checkpointer interface {
-	// GetOrCreate returns all persisted pod device configs, or an empty map
-	// if the checkpoint does not yet exist. Used at startup to restore state.
+	// GetOrCreate returns all persisted device configs. Used at startup to
+	// restore state.
 	GetOrCreate() (map[types.UID]map[string]DeviceConfig, error)
 	// Store persists the device config for a single pod/device pair.
 	Store(podUID types.UID, deviceName string, config DeviceConfig) error
@@ -110,7 +116,9 @@ type PodConfigStore struct {
 }
 
 // NewPodConfigStore creates a new PodConfigStore. If a Checkpointer is
-// provided, existing state is loaded from the checkpoint into memory.
+// provided, existing device configs are loaded from the checkpoint into memory.
+// Pod-level state is not persisted; NetNS is rebuilt through Synchronize() on
+// driver startup, while LastNRIActivity resets to its zero value.
 func NewPodConfigStore(checkpointer Checkpointer) (*PodConfigStore, error) {
 	s := &PodConfigStore{
 		configs:      make(map[types.UID]PodConfig),
@@ -253,9 +261,34 @@ func (s *PodConfigStore) GetPodConfig(podUID types.UID) (PodConfig, bool) {
 	return PodConfig{
 		DeviceConfigs:   configsCopy,
 		LastNRIActivity: podConfig.LastNRIActivity,
+		NetNS:           podConfig.NetNS,
 	}, true
 }
 
+// SetPodNetNs stores the Pod's network namespace path in the pod-level config.
+// This is in-memory only; pod NetNS is rebuilt from the container runtime on
+// driver restart via Synchronize().
+func (s *PodConfigStore) SetPodNetNs(podUID types.UID, netns string) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	podCfg, ok := s.configs[podUID]
+	if !ok {
+		klog.Warningf("SetPodNetNs: pod UID %s not found in store; skipping NetNS update", podUID)
+		return
+	}
+	podCfg.NetNS = netns
+	s.configs[podUID] = podCfg
+}
+
+// GetPodNetNs returns the stored network namespace for the given pod UID.
+func (s *PodConfigStore) GetPodNetNs(podUID types.UID) (string, bool) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+	podCfg, ok := s.configs[podUID]
+	return podCfg.NetNS, ok
+}
+
 // DeleteClaim removes all configurations associated with a given claim and
 // returns the list of Pod UIDs that were associated with it.
 // Like DeletePod, checkpoint failures do not prevent in-memory cleanup.

pkg/driver/pod_device_config_bolt_test.go

@@ -176,6 +176,7 @@ func TestPodConfigStore_Persistence(t *testing.T) {
 		t.Fatalf("NewPodConfigStore() error: %v", err)
 	}
 	store1.SetDeviceConfig("pod-1", "eth0", config)
+	store1.SetPodNetNs("pod-1", "/var/run/netns/test-ns")
 	store1.Close()
 
 	// Reopen and verify data was restored from checkpoint.
@@ -204,6 +205,11 @@ func TestPodConfigStore_Persistence(t *testing.T) {
 	if len(podConfig.DeviceConfigs) != 1 {
 		t.Errorf("Expected 1 device config after reopen, got %d", len(podConfig.DeviceConfigs))
 	}
+
+	// Verify NetNS is NOT restored (in-memory only)
+	if podConfig.NetNS != "" {
+		t.Errorf("Expected NetNS to be empty after reopen, but got %s", podConfig.NetNS)
+	}
 }
 
 // TestPodConfigStore_DeletePodCheckpoints verifies that DeletePod and
@@ -385,3 +391,4 @@ func TestPodConfigStore_NoCheckpointer(t *testing.T) {
 		t.Errorf("Close() error: %v", err)
 	}
 }
+

pkg/driver/pod_device_config_test.go

@@ -107,6 +107,49 @@ func TestPodConfigStore_SetAndGet(t *testing.T) {
 	}
 }
 
+// TestPodConfigStore_NetNs verifies that NetNS path can be stored and retrieved correctly in memory.
+func TestPodConfigStore_NetNs(t *testing.T) {
+	store := mustNewPodConfigStore()
+	podUID := types.UID("test-pod-uid-1")
+	netns := "/var/run/netns/test-ns"
+
+	// Test Get on non-existent item
+	_, found := store.GetPodNetNs(podUID)
+	if found {
+		t.Errorf("GetPodNetNs() found a netns before SetPodNetNs(), expected not found")
+	}
+
+	// Add a dummy device config so the pod exists in the store
+	store.SetDeviceConfig(podUID, "dummy-device", DeviceConfig{})
+
+	store.SetPodNetNs(podUID, netns)
+
+	retrievedNetNs, found := store.GetPodNetNs(podUID)
+	if !found {
+		t.Fatalf("GetPodNetNs() did not find netns after SetPodNetNs(), expected found")
+	}
+	if retrievedNetNs != netns {
+		t.Errorf("GetPodNetNs() retrieved %s, want %s", retrievedNetNs, netns)
+	}
+
+	// Test Get with different podUID
+	_, found = store.GetPodNetNs(types.UID("other-pod-uid"))
+	if found {
+		t.Errorf("GetPodNetNs() found netns for wrong podUID, expected not found")
+	}
+
+	// Test overwriting
+	newNetNs := "/var/run/netns/new-ns"
+	store.SetPodNetNs(podUID, newNetNs)
+	retrievedNetNs, found = store.GetPodNetNs(podUID)
+	if !found {
+		t.Fatalf("GetPodNetNs() did not find netns after overwrite, expected found")
+	}
+	if retrievedNetNs != newNetNs {
+		t.Errorf("GetPodNetNs() retrieved %s after overwrite, want %s", retrievedNetNs, newNetNs)
+	}
+}
+
 func TestPodConfigStore_DeletePod(t *testing.T) {
 	store := mustNewPodConfigStore()
 	podUID1 := types.UID("test-pod-uid-1")