search the podspec for storepaths and add them to the volume

Author: Lillecarl

kubenix/options.nix

@@ -151,11 +151,11 @@ in
           overlays = [
             (import ../pkgs)
             (
-              self: pkgs:
+              final: prev:
               let
                 callPackage =
-                  pp:
-                  pkgs.callPackage pp {
+                  packagePath:
+                  final.callPackage packagePath {
                     inherit (cfg) dinix;
                   };
               in

python/nix_csi/service.py

@@ -2,6 +2,7 @@
 import kr8s
 import logging
 import os
+import re
 import shutil
 import socket
 import tempfile
@@ -16,7 +17,9 @@
 from grpclib.const import Status
 from grpclib.server import Server
 from importlib import metadata
+from kr8s.asyncio.objects import Pod
 from pathlib import Path
+from typing import Any, Iterator
 
 logger = logging.getLogger("nix-csi")
 
@@ -48,6 +51,27 @@
 NAMESPACE = os.environ.get("KUBE_NAMESPACE", "nix-csi")
 BUILDERS_SERVICE = "nix-csi-builders"
 
+# Nix base32 excludes: e, o, t, u
+STORE_PATH_RE = re.compile(r"/nix/store/[0-9a-df-np-sv-z]{32}-[^\s/]+")
+
+
+def extract_store_paths(value: Any) -> Iterator[str]:
+    match value:
+        case str():
+            yield from Path(STORE_PATH_RE.findall(value))
+        case dict():
+            for v in value.values():
+                yield from Path(extract_store_paths(v))
+        case list() | tuple():
+            for item in value:
+                yield from Path(extract_store_paths(item))
+
+
+def extract_store_name(value: Any) -> Iterator[str]:
+    for path in extract_store_paths(value):
+        yield str(path).removeprefix("/nix/store/")
+
+
 async def get_current_system():
     """Get system string evaluated by nix"""
     return (
@@ -80,6 +104,7 @@ async def get_builder_uris():
         logger.warning(f"Failed to discover builder pods: {e}")
         return []
 
+
 class NodeServicer(csi_grpc.NodeBase):
     volumeLocks: defaultdict[str, Semaphore] = defaultdict(Semaphore)
 
@@ -102,8 +127,8 @@ async def NodePublishVolume(self, stream):
             # Using sentinel path instead of None to avoid Optional type and None checks everywhere.
             # The if/elif blocks below should set this to a real path; the exists() check at line ~157
             # will fail if none of the branches executed (sentinel path doesn't exist).
-            packagePath: Path = Path("/nonexistent/path/that/should/never/exist")
-            gcPath = CSI_GCROOTS / request.volume_id
+            primaryPackagePath: Path = Path("/nonexistent/path/that/should/never/exist")
+            gcRoot = CSI_GCROOTS / request.volume_id
 
             extraArgs = []
 
@@ -135,27 +160,57 @@ async def NodePublishVolume(self, stream):
                 except (GRPCError, OSError, asyncio.TimeoutError):
                     logger.warning("Cache connectivity check failed")
 
-            # Source selection order (intentional, documented in README):
-            # 1. storePath - if present, use directly
-            # 2. flakeRef - if storePath not present, build flake
-            # 3. nixExpr - if neither above present, evaluate expression
-            # Users can specify multiple; first non-None in priority order is used.
-            if storePath is not None:
-                async with self.volumeLocks[storePath]:
-                    logger.debug(f"{storePath=}")
-                    packagePath = Path(storePath)
-                    if not packagePath.exists():
-                        await try_console(
+            podName = request.volume_context.get("csi.storage.k8s.io/pod.name")
+            podNamespace = request.volume_context.get(
+                "csi.storage.k8s.io/pod.namespace"
+            )
+            podUid = request.volume_context.get("csi.storage.k8s.io/pod.uid")
+
+            packagePaths = []
+            if podName and podNamespace and podUid:
+                pod = await Pod.get(podName, podNamespace)
+                if pod.metadata.uid != podUid:
+                    raise GRPCError(Status.INTERNAL, "poduid doesn't match")
+
+                for packagePath in set(extract_store_paths(pod.raw)):
+                    async with self.volumeLocks[storePath]:
+                        name = extract_store_name(packagePath)
+                        result = await try_console(
                             "nix",
                             "build",
                             *extraArgs,
                             "--print-out-paths",
                             "--out-link",
-                            gcPath,
+                            gcRoot / name,
                             packagePath,
                             timeout=NIX_BUILD_TIMEOUT,
                         )
-            if flakeRef is not None and not packagePath.exists():
+                        packagePaths.append(Path(result.stdout.splitlines()[0]))
+
+            # Source selection order (intentional, documented in README):
+            # 1. storePath - if present, use directly
+            # 2. flakeRef - if storePath not present, build flake
+            # 3. nixExpr - if neither above present, evaluate expression
+            # Users can specify multiple; first non-None in priority order is used.
+            if storePath is not None:
+                async with self.volumeLocks[storePath]:
+                    logger.debug(f"{storePath=}")
+                    name = extract_store_name(storePath)
+                    result = await try_console(
+                        "nix",
+                        "build",
+                        *extraArgs,
+                        "--print-out-paths",
+                        "--out-link",
+                        gcRoot / storePath,
+                        storePath,
+                        timeout=NIX_BUILD_TIMEOUT,
+                    )
+                    packagePaths.append(Path(result.stdout.splitlines()[0]))
+                    if not primaryPackagePath.exists():
+                        primaryPackagePath = Path(result.stdout.splitlines()[0])
+
+            if flakeRef is not None:
                 async with self.volumeLocks[flakeRef]:
                     logger.debug(f"{flakeRef=}")
 
@@ -166,12 +221,15 @@ async def NodePublishVolume(self, stream):
                         *extraArgs,
                         "--print-out-paths",
                         "--out-link",
-                        gcPath,
+                        gcRoot / "flake",
                         flakeRef,
                         timeout=NIX_BUILD_TIMEOUT,
                     )
-                    packagePath = Path(result.stdout.splitlines()[0])
-            if nixExpr is not None and not packagePath.exists():
+                    packagePaths.append(Path(result.stdout.splitlines()[0]))
+                    if not primaryPackagePath.exists():
+                        primaryPackagePath = Path(result.stdout.splitlines()[0])
+
+            if nixExpr is not None:
                 async with self.volumeLocks[nixExpr]:
                     logger.debug(f"{nixExpr=}")
                     with tempfile.NamedTemporaryFile(mode="w", suffix=".nix") as tmp:
@@ -185,14 +243,16 @@ async def NodePublishVolume(self, stream):
                             *extraArgs,
                             "--print-out-paths",
                             "--out-link",
-                            gcPath,
+                            gcRoot / "expr",
                             "--file",
                             tmp.name,
                             timeout=NIX_BUILD_TIMEOUT,
                         )
-                        packagePath = Path(result.stdout.splitlines()[0])
+                        packagePaths.append(Path(result.stdout.splitlines()[0]))
+                        if not primaryPackagePath.exists():
+                            primaryPackagePath = Path(result.stdout.splitlines()[0])
 
-            if not packagePath.exists():
+            if not primaryPackagePath.exists():
                 logger.error("packagePath doesn't exist after building")
                 raise GRPCError(
                     Status.INVALID_ARGUMENT,
@@ -207,13 +267,13 @@ async def NodePublishVolume(self, stream):
             # Create NIX_STATE_DIR where database will be initialized
             NIX_STATE_DIR.mkdir(parents=True, exist_ok=True)
 
-            # Get closure
-            paths = (
+            # Get storepaths from all packages
+            storePaths = (
                 await try_captured(
                     "nix",
                     "path-info",
                     "--recursive",
-                    packagePath,
+                    *packagePaths,
                 )
             ).stdout.splitlines()
 
@@ -222,14 +282,16 @@ async def NodePublishVolume(self, stream):
                 # extra steps. (Hardlinking instead of dumbcopying)
 
                 # Install CSI gcroots
-                await try_captured(
-                    "nix",
-                    "build",
-                    "--out-link",
-                    gcPath,
-                    packagePath,
-                    timeout=NIX_BUILD_TIMEOUT,
-                )
+                for packagePath in packagePaths:
+                    name = extract_store_name(packagePath)
+                    await try_captured(
+                        "nix",
+                        "build",
+                        "--out-link",
+                        gcRoot / name,
+                        packagePath,
+                        timeout=NIX_BUILD_TIMEOUT,
+                    )
 
                 # Copy closure to substore, rsync saves a lot of implementation
                 # headache here. --archive keeps all attributes, --hard-links
@@ -242,7 +304,7 @@ async def NodePublishVolume(self, stream):
                         "--links",
                         "--hard-links",
                         "--mkpath",
-                        *paths,
+                        *storePaths,
                         volumeRoot / "nix/store",
                     )
 
@@ -251,21 +313,23 @@ async def NodePublishVolume(self, stream):
                 await try_captured(
                     "nix_init_db",
                     NIX_STATE_DIR,
-                    *paths,
+                    *storePaths,
                 )
 
-                # install gcroots in container using chroot store this is
+                # 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.
-                await try_captured(
-                    "nix",
-                    "build",
-                    "--store",
-                    volumeRoot,
-                    "--out-link",
-                    NIX_STATE_DIR / "gcroots/result",
-                    packagePath,
-                )
+                for packagePath in packagePaths:
+                    name = extract_store_name(packagePath)
+                    await try_captured(
+                        "nix",
+                        "build",
+                        "--store",
+                        volumeRoot,
+                        "--out-link",
+                        NIX_STATE_DIR / f"gcroots/result/{name}",
+                        packagePath,
+                    )
 
                 # install /nix/var/result in container using chroot store
                 await try_captured(
@@ -275,12 +339,12 @@ async def NodePublishVolume(self, stream):
                     volumeRoot,
                     "--out-link",
                     volumeRoot / "nix/var/result",
-                    packagePath,
+                    primaryPackagePath,
                 )
             except Exception:
                 logger.exception("Failed to build volume")
                 # Remove gcroots if we failed something else
-                gcPath.unlink(missing_ok=True)
+                shutil.rmtree(gcRoot, ignore_errors=True)
                 # Remove what we were working on
                 shutil.rmtree(volumeRoot, True)
                 raise
@@ -297,7 +361,7 @@ async def NodePublishVolume(self, stream):
                     "--bind",
                     "-o",
                     "ro",
-                    volumeRoot / "nix",
+                    volumeRoot,
                     targetPath,
                 ]
             else:
@@ -315,7 +379,7 @@ async def NodePublishVolume(self, stream):
                     "overlay",
                     "overlay",
                     "-o",
-                    f"rw,lowerdir={volumeRoot / 'nix'},upperdir={upperdir},workdir={workdir}",
+                    f"rw,lowerdir={volumeRoot},upperdir={upperdir},workdir={workdir}",
                     targetPath,
                 ]
 
@@ -324,7 +388,7 @@ async def NodePublishVolume(self, stream):
                 logger.debug(f"Mount target {targetPath} was already mounted")
             elif mount.returncode != 0:
                 # Clean up resources on mount failure
-                gcPath.unlink(missing_ok=True)
+                shutil.rmtree(gcRoot, ignore_errors=True)
                 shutil.rmtree(volumeRoot, ignore_errors=True)
                 raise GRPCError(
                     Status.INTERNAL,
@@ -334,7 +398,7 @@ async def NodePublishVolume(self, stream):
             reply = csi_pb2.NodePublishVolumeResponse()
             await stream.send_message(reply)
 
-            task = asyncio.create_task(copyToCache(packagePath))
+            task = asyncio.create_task(copyToCache(primaryPackagePath))
             task.add_done_callback(
                 lambda t: logger.error(f"copyToCache failed: {t.exception()}")
                 if t.exception()
@@ -385,11 +449,11 @@ async def NodeUnpublishVolume(self, stream):
                     )
 
             # Remove gcroots
-            gcPath = CSI_GCROOTS / request.volume_id
-            if gcPath.exists():
+            gcRoot = CSI_GCROOTS / request.volume_id
+            if gcRoot.exists():
                 try:
-                    gcPath.unlink()
-                    logger.debug(f"unlinked {gcPath=}")
+                    shutil.rmtree(gcRoot, ignore_errors=True)
+                    logger.debug(f"unlinked {gcRoot=}")
                 except Exception as ex:
                     raise GRPCError(
                         Status.INTERNAL, f"unlinking {targetPath=} failed", ex