From 76308b1612c7b989777dc0bbc755bebc963d19bc Mon Sep 17 00:00:00 2001
From: Gaetan Lepage <gaetan@glepage.com>
Date: Wed, 9 Oct 2024 09:00:19 +0200
Subject: [PATCH] Use threads instead of processes for parallel 'nix eval'

---
 nixpkgs_review/nix.py | 40 ++++++++++++++++++----------------------
 1 file changed, 18 insertions(+), 22 deletions(-)

diff --git a/nixpkgs_review/nix.py b/nixpkgs_review/nix.py
index 8d00531..b587fd3 100644
--- a/nixpkgs_review/nix.py
+++ b/nixpkgs_review/nix.py
@@ -1,11 +1,10 @@
+import concurrent.futures
 import json
-import multiprocessing as mp
 import os
 import shlex
 import shutil
 import subprocess
 from dataclasses import dataclass, field
-from functools import partial
 from pathlib import Path
 from sys import platform
 from tempfile import NamedTemporaryFile
@@ -271,32 +270,29 @@ def nix_eval(
             os.unlink(attr_json.name)
 
 
-def nix_eval_thread(
-    system: System,
-    attr_names: set[str],
-    allow: AllowedFeatures,
-    nix_path: str,
-) -> tuple[System, list[Attr]]:
-    return system, nix_eval(attr_names, system, allow, nix_path)
-
-
 def multi_system_eval(
     attr_names_per_system: dict[System, set[str]],
     allow: AllowedFeatures,
     nix_path: str,
     n_procs: int,
 ) -> dict[System, list[Attr]]:
-    nix_eval_partial = partial(
-        nix_eval_thread,
-        allow=allow,
-        nix_path=nix_path,
-    )
-
-    args: list[tuple[System, set[str]]] = list(attr_names_per_system.items())
-    with mp.Pool(n_procs) as pool:
-        results: list[tuple[System, list[Attr]]] = pool.starmap(nix_eval_partial, args)
-
-    return {system: attrs for system, attrs in results}
+    results: dict[System, list[Attr]] = {}
+    with concurrent.futures.ThreadPoolExecutor(max_workers=n_procs) as executor:
+        future_to_system = {
+            executor.submit(
+                nix_eval,
+                attrs=attrs,
+                system=system,
+                allow=allow,
+                nix_path=nix_path,
+            ): system
+            for system, attrs in attr_names_per_system.items()
+        }
+        for future in concurrent.futures.as_completed(future_to_system):
+            system = future_to_system[future]
+            results[system] = future.result()
+
+    return results
 
 
 def nix_build(