From 0baf658c4bea02fd5310e1b3dd5d039d8e44e54a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?F=C3=A9lix=20Baylac=20Jacqu=C3=A9?=
 <felix@alternativebit.fr>
Date: Tue, 17 Jan 2023 12:51:03 +0100
Subject: [PATCH] wip

---
 src/nixglhost.py | 36 ++++++++++++++++++++++++------------
 1 file changed, 24 insertions(+), 12 deletions(-)

diff --git a/src/nixglhost.py b/src/nixglhost.py
index 4b5fd0a..086771d 100755
--- a/src/nixglhost.py
+++ b/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, cache_dir_root: str, rpath_dir_root: str
+) -> str:
     """Generate a cache directory for the LIBRARY_PATH host directory.
 
     This cache directory is mirroring the host directory containing
@@ -428,6 +431,7 @@ def cache_library_path(library_path: LibraryPath, cache_dir_root: str) -> str:
     # Paths
     cache_path_root: str = os.path.join(cache_dir_root, path_hash)
     lib_dir = os.path.join(cache_path_root, "lib")
+    rpath_lib_dir = os.path.join(rpath_dir_root, "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,23 @@ 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, (os.path.join(cache_dir, p)))
+                )
+            # 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: