nix_csi: improve code quality and observability

- Remove duplicate constants from identityservicer.py, import from constants.py
- Fix incorrect variable names in error messages in NodeUnpublishVolume
- Add error context logging on cache copy retry failures
- Cache store names in prepare_volume() to avoid duplicate extraction calls
- Parallelize gcroot installations using asyncio.gather
- Add extract_store_paths_set() convenience function
- Make CSI socket path configurable via CSI_SOCKET_PATH env var
- Add structured logging fields to Publish/Unpublish operations
- Enhance Probe method to verify /nix/store accessibility

Author: Lillecarl

python/nix_csi/cache.py

@@ -90,5 +90,11 @@ async def copy_to_cache(package_path: Path) -> None:
                 if nix_copy.returncode == 0:
                     logger.debug(f"Successfully copied to cache: {package_path}")
                     break
+                else:
+                    logger.debug(
+                        f"Copy attempt {attempt + 1}/6 failed for {package_path}: {nix_copy.combined}"
+                    )
             else:
-                logger.error(f"Failed to copy to cache after 6 attempts: {package_path}")
+                logger.error(
+                    f"Failed to copy to cache after 6 attempts: {package_path}"
+                )

python/nix_csi/cli.py

@@ -33,3 +33,6 @@
 
 # Simple string check is fine - value controlled by easykubenix (always "true" or "false")
 CACHE_ENABLED = os.environ.get("CACHE_ENABLED", "false") == "true"
+
+# CSI socket path for gRPC server
+CSI_SOCKET_PATH = os.environ.get("CSI_SOCKET_PATH", "/csi/csi.sock")

python/nix_csi/constants.py

@@ -1,11 +1,11 @@
+from pathlib import Path
+
 from csi import csi_grpc, csi_pb2
 from google.protobuf.wrappers_pb2 import BoolValue
 from grpclib import GRPCError
 from grpclib.const import Status
-from importlib import metadata
 
-CSI_PLUGIN_NAME = "nix.csi.store"
-CSI_VENDOR_VERSION = metadata.version("nix-csi")
+from .constants import CSI_PLUGIN_NAME, CSI_VENDOR_VERSION
 
 
 class IdentityServicer(csi_grpc.IdentityBase):
@@ -44,5 +44,10 @@ async def Probe(self, stream):
         request: csi_pb2.ProbeRequest | None = await stream.recv_message()
         if request is None:
             raise GRPCError(Status.INVALID_ARGUMENT, "Received None request in Probe")
-        reply = csi_pb2.ProbeResponse(ready=BoolValue(value=True))
+
+        # Verify /nix/store is accessible - critical for CSI driver operation
+        nix_store = Path("/nix/store")
+        ready = nix_store.is_dir()
+
+        reply = csi_pb2.ProbeResponse(ready=BoolValue(value=ready))
         await stream.send_message(reply)

python/nix_csi/identityservicer.py

@@ -16,11 +16,17 @@
 
 from .builders import build_builder_args, get_builder_uris
 from .cache import check_cache_connectivity, copy_to_cache, get_substituter_args
-from .constants import CSI_GCROOTS, CSI_VOLUMES, NIX_BUILD_TIMEOUT
+from .constants import CSI_GCROOTS, CSI_SOCKET_PATH, CSI_VOLUMES, NIX_BUILD_TIMEOUT
 from .identityservicer import IdentityServicer
 from .nix import build_flake_ref, build_nix_expr, build_store_path, get_current_system
-from .store import extract_store_paths
-from .volume import cleanup_failed_volume, is_mount, mount_volume, prepare_volume, unmount
+from .store import extract_store_paths_set
+from .volume import (
+    cleanup_failed_volume,
+    is_mount,
+    mount_volume,
+    prepare_volume,
+    unmount,
+)
 
 logger = logging.getLogger("nix-csi")
 
@@ -81,7 +87,7 @@ async def _build_pod_packages(
             )
 
         package_paths = []
-        pod_store_paths = set(extract_store_paths(pod.raw))
+        pod_store_paths = extract_store_paths_set(pod.raw)
 
         for package_path in pod_store_paths:
             logger.debug(f"{package_path=}")
@@ -155,7 +161,13 @@ async def NodePublishVolume(self, stream):
         if request is None:
             raise GRPCError(Status.INVALID_ARGUMENT, "NodePublishVolumeRequest is None")
 
-        logger.info(f"Publish {request.target_path}")
+        logger.info(
+            "Publishing volume",
+            extra={
+                "volume_id": request.volume_id,
+                "target_path": request.target_path,
+            },
+        )
 
         if not request.volume_context.get("csi.storage.k8s.io/ephemeral"):
             raise GRPCError(
@@ -230,7 +242,13 @@ async def NodeUnpublishVolume(self, stream):
         if request is None:
             raise ValueError("NodeUnpublishVolumeRequest is None")
 
-        logger.info(f"Unpublish {request.target_path}")
+        logger.info(
+            "Unpublishing volume",
+            extra={
+                "volume_id": request.volume_id,
+                "target_path": request.target_path,
+            },
+        )
 
         async with self.volumeLocks[request.volume_id]:
             target_path = Path(request.target_path)
@@ -262,9 +280,7 @@ async def NodeUnpublishVolume(self, stream):
                     shutil.rmtree(gc_root, ignore_errors=True)
                     logger.debug(f"unlinked {gc_root=}")
                 except Exception as ex:
-                    raise GRPCError(
-                        Status.INTERNAL, f"unlinking {target_path=} failed", ex
-                    )
+                    raise GRPCError(Status.INTERNAL, f"unlinking {gc_root=} failed", ex)
 
             # Remove hardlink farm
             volume_path = CSI_VOLUMES / request.volume_id
@@ -274,7 +290,7 @@ async def NodeUnpublishVolume(self, stream):
                     logger.debug(f"removed {volume_path=}")
                 except Exception as ex:
                     raise GRPCError(
-                        Status.INTERNAL, f"recursive removing {target_path=} failed", ex
+                        Status.INTERNAL, f"removing {volume_path=} failed", ex
                     )
 
             await stream.send_message(csi_pb2.NodeUnpublishVolumeResponse())
@@ -313,7 +329,7 @@ async def NodeUnstageVolume(self, stream):
 
 
 async def serve():
-    sock_path = "/csi/csi.sock"
+    sock_path = CSI_SOCKET_PATH
     Path(sock_path).unlink(missing_ok=True)
 
     identityServicer = IdentityServicer()

python/nix_csi/service.py

@@ -19,5 +19,10 @@ def extract_store_paths(value: Any) -> Iterator[Path]:
                 yield from extract_store_paths(item)
 
 
+def extract_store_paths_set(value: Any) -> set[Path]:
+    """Convenience wrapper that returns a deduplicated set of store paths."""
+    return set(extract_store_paths(value))
+
+
 def extract_store_name(path: Path | str) -> str:
     return str(path).removeprefix("/nix/store/")

python/nix_csi/store.py

@@ -5,7 +5,12 @@
 from grpclib import GRPCError
 from grpclib.const import Status
 
-from .constants import CSI_GCROOTS, CSI_VOLUMES, MOUNT_ALREADY_MOUNTED, NIX_BUILD_TIMEOUT
+from .constants import (
+    CSI_GCROOTS,
+    CSI_VOLUMES,
+    MOUNT_ALREADY_MOUNTED,
+    NIX_BUILD_TIMEOUT,
+)
 from .hardlinks import deref_hardlink_tree, hardlink_closure
 from .nix import get_closure_paths, init_database, install_gcroot, install_result_link
 from .store import extract_store_name
@@ -32,20 +37,26 @@ async def prepare_volume(
     # Get storepaths from all packages
     store_paths = await get_closure_paths(package_paths)
 
+    # Pre-compute store names to avoid duplicate calls
+    package_names = {p: extract_store_name(p) for p in package_paths}
+
     # This block is essentially nix copy into a chroot store with
     # extra steps. (Hardlinking instead of dumbcopying)
 
-    # Install CSI gcroots
-    for package_path in package_paths:
-        name = extract_store_name(package_path)
-        await try_captured(
-            "nix",
-            "build",
-            "--out-link",
-            gc_root / name,
-            package_path,
-            timeout=NIX_BUILD_TIMEOUT,
-        )
+    # Install CSI gcroots (parallel since packages are independent)
+    await asyncio.gather(
+        *[
+            try_captured(
+                "nix",
+                "build",
+                "--out-link",
+                gc_root / package_names[package_path],
+                package_path,
+                timeout=NIX_BUILD_TIMEOUT,
+            )
+            for package_path in package_paths
+        ]
+    )
 
     # Copy closure to substore
     hardlink_closure([Path(p) for p in store_paths], volume_root / "nix/store")
@@ -56,9 +67,14 @@ async def prepare_volume(
     # Install gcroots in container using chroot store. This is
     # required because the auto roots created for /nix/var/result
     # will point to Narnia while this one points into store.
-    for package_path in package_paths:
-        name = extract_store_name(package_path)
-        await install_gcroot(volume_root, package_path, name, NIX_STATE_DIR)
+    await asyncio.gather(
+        *[
+            install_gcroot(
+                volume_root, package_path, package_names[package_path], NIX_STATE_DIR
+            )
+            for package_path in package_paths
+        ]
+    )
 
     # Install /nix/var/result in container using chroot store
     if primary_package is not None:
@@ -134,6 +150,4 @@ async def unmount(path: Path) -> None:
     """Unmount a path. Raises GRPCError on failure if still mounted."""
     result = await run_captured("umount", "--verbose", path)
     if result.returncode != 0 and await is_mount(path):
-        raise GRPCError(
-            Status.INTERNAL, "unmount failed", f"{result.combined=}"
-        )
+        raise GRPCError(Status.INTERNAL, "unmount failed", f"{result.combined=}")