From 6ec94cdfaa2036fb8a78d53c65c22ae6d782e1ae Mon Sep 17 00:00:00 2001
From: John Shaffer <jshaffer2112@protonmail.com>
Date: Tue, 8 Mar 2022 13:50:28 -0800
Subject: [PATCH] Add :parallel-threads option to set thread pool size

---
 deps.edn                |  1 +
 src/kaocha/testable.clj | 49 +++++++++++------------------------------
 2 files changed, 14 insertions(+), 36 deletions(-)

diff --git a/deps.edn b/deps.edn
index b2fad127..f536ece7 100644
--- a/deps.edn
+++ b/deps.edn
@@ -5,6 +5,7 @@
   org.clojure/spec.alpha       {:mvn/version "0.2.194"}
   org.clojure/tools.cli        {:mvn/version "1.0.206"}
   lambdaisland/tools.namespace {:mvn/version "0.0-237"}
+  com.climate/claypoole        {:mvn/version "1.1.4"}
   lambdaisland/deep-diff       {:mvn/version "0.0-47"}
   org.tcrawley/dynapath        {:mvn/version "1.1.0"}
   slingshot/slingshot          {:mvn/version "0.12.2"}
diff --git a/src/kaocha/testable.clj b/src/kaocha/testable.clj
index d257fa8e..082183f0 100644
--- a/src/kaocha/testable.clj
+++ b/src/kaocha/testable.clj
@@ -4,6 +4,7 @@
             [clojure.pprint :as pprint]
             [clojure.spec.alpha :as s]
             [clojure.test :as t]
+            [com.climate.claypoole :as cp]
             [kaocha.classpath :as classpath]
             [kaocha.hierarchy :as hierarchy]
             [kaocha.history :as history]
@@ -58,16 +59,6 @@
     (retry-assert-spec type testable (dec n))) ;otherwise, retry
     ))
 
-(defn deref-recur [testables]
-  (cond (future? testables) (deref testables)
-        (vector? testables) (doall (mapv deref-recur testables))
-        (seq? testables) (deref-recur (into [] (doall testables)))
-        (contains? testables :kaocha.test-plan/tests)
-        (update testables :kaocha.test-plan/tests deref-recur)
-        (contains? testables :kaocha.result/tests)
-        (update testables :kaocha.result/tests deref-recur)
-        :else testables))
-
 (defn- load-type+validate
   "Try to load a testable type, and validate it both to be a valid generic testable, and a valid instance given the type.
 
@@ -287,32 +278,18 @@
 (defn run-testables-parallel
   "Run a collection of testables, returning a result collection."
   [testables test-plan]
-  (doall testables)
-  (let [load-error? (some ::load-error testables)
-        types (set (:parallel-children-exclude *config*))
-        futures (map #(do 
-                        (future 
-                          (binding [*config*  
-                                    (cond-> *config* 
-                                      (contains? types (:kaocha.testable/type %)) (dissoc :parallel)
-                                      true (update :levels (fn [x] (if (nil? x) 1 (inc x))))) ] 
-                            (run-testable % test-plan))))
-                     testables)]
-    (comment (loop [result [] ;(ArrayBlockingQueue. 1024)
-                    [test & testables] testables]
-               (if test
-                 (let [test (cond-> test
-                              (and load-error? (not (::load-error test)))
-                              (assoc ::skip true))
-                       r (run-testable test test-plan)]
-                   (if (or (and *fail-fast?* (result/failed? r)) (::skip-remaining? r))
-                     ;(reduce put-return result [[r] testables])
-                     (reduce into result [[r] testables])
-                     ;(recur (doto result (.put r)) testables)
-                     (recur (conj result r) testables)))
-                 result)))
-    (deref-recur futures)))
-
+  (let [num-threads (or (:parallel-threads *config*) (+ 2 (cp/ncpus)))
+        types (set (:parallel-children-exclude *config*))]
+    (cp/with-shutdown! [pool (cp/threadpool num-threads :name "kaocha-test-runner")]
+      (doall
+       (cp/pmap
+        pool
+        #(binding [*config*
+                   (cond-> *config*
+                     (contains? types (:kaocha.testable/type %)) (dissoc :parallel)
+                     true (update :levels (fn [x] (if (nil? x) 1 (inc x)))))]
+           (run-testable % test-plan))
+        testables)))))
 
 (defn run-testables 
   [testables test-plan]