wip

src/nixglhost.py

@@ -6,9 +6,10 @@ import json
 import os
 import re
 import shutil
-import subprocess
 import stat
+import subprocess
 import sys
+import tempfile
 import time
 from glob import glob
 from typing import List, Literal, Dict, Tuple, TypedDict, TextIO, Optional
@@ -412,7 +413,9 @@ def scan_dsos_from_dir(path: str) -> Optional[LibraryPath]:
         return None
 
 
-def cache_library_path(library_path: LibraryPath, cache_dir_root: str) -> str:
+def cache_library_path(
+    library_path: LibraryPath, temp_cache_dir_root: str, final_cache_dir_root: str
+) -> str:
     """Generate a cache directory for the LIBRARY_PATH host directory.
 
     This cache directory is mirroring the host directory containing
@@ -426,8 +429,9 @@ def cache_library_path(library_path: LibraryPath, cache_dir_root: str) -> str:
     h.update(library_path.path.encode("utf8"))
     path_hash: str = h.hexdigest()
     # Paths
-    cache_path_root: str = os.path.join(cache_dir_root, path_hash)
+    cache_path_root: str = os.path.join(temp_cache_dir_root, path_hash)
     lib_dir = os.path.join(cache_path_root, "lib")
+    rpath_lib_dir = os.path.join(final_cache_dir_root, path_hash, "lib")
     cuda_dir = os.path.join(cache_path_root, "cuda")
     egl_dir = os.path.join(cache_path_root, "egl")
     glx_dir = os.path.join(cache_path_root, "glx")
@@ -440,7 +444,7 @@ def cache_library_path(library_path: LibraryPath, cache_dir_root: str) -> str:
     ]:
         os.makedirs(d, exist_ok=True)
         if len(dsos) > 0:
-            copy_and_patch_libs(dsos=dsos, dest_dir=d, rpath=lib_dir)
+            copy_and_patch_libs(dsos=dsos, dest_dir=d, rpath=rpath_lib_dir)
         else:
             log_info(f"Did not find any DSO to put in {d}, skipping copy and patching.")
     return path_hash
@@ -534,15 +538,21 @@ def nvidia_main(
         cache_content, cache_file_path
     ) or not os.path.isfile(cached_ld_library_path):
         log_info("The cache is not up to date, regenerating it")
-        shutil.rmtree(cache_dir)
-        cache_paths: List[str] = []
-        for p in cache_content.paths:
-            log_info(f"Caching {p}")
-            cache_paths.append(cache_library_path(p, cache_dir))
-        cache_absolute_paths = [os.path.join(cache_dir, p) for p in cache_paths]
-        nix_gl_ld_library_path = generate_cache_metadata(
-            cache_dir, cache_content, cache_absolute_paths
-        )
+        with tempfile.TemporaryDirectory() as tmp_cache:
+            new_cache = os.path.join(tmp_cache, "nix-gl-host")
+            os.makedirs(new_cache)
+            cache_paths: List[str] = []
+            for p in cache_content.paths:
+                log_info(f"Caching {p}")
+                cache_paths.append(cache_library_path(p, tmp_cache, cache_dir))
+            # Pointing the LD_LIBRARY_PATH to the final destination
+            # instead of the tmp dir.
+            cache_absolute_paths = [os.path.join(cache_dir, p) for p in cache_paths]
+            nix_gl_ld_library_path = generate_cache_metadata(
+                cache_dir, cache_content, cache_absolute_paths
+            )
+            log_info(f"Mv {new_cache} to {cache_dir}")
+            shutil.move(new_cache, cache_dir)
     else:
         log_info("The cache is up to date, re-using it.")
         with open(cached_ld_library_path, "r", encoding="utf8") as f: