From 3c68292963c049508fdee62259f19a6ecb0ee071 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/test_suite.clj | 10 --------
 src/kaocha/testable.clj   | 52 +++++++++++----------------------------
 3 files changed, 16 insertions(+), 47 deletions(-)

diff --git a/deps.edn b/deps.edn
index cde27a9e..917c53da 100644
--- a/deps.edn
+++ b/deps.edn
@@ -5,6 +5,7 @@
   org.clojure/spec.alpha       {:mvn/version "0.3.218"}
   org.clojure/tools.cli        {:mvn/version "1.0.206"}
   lambdaisland/tools.namespace {:mvn/version "0.1.247"}
+  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/test_suite.clj b/src/kaocha/test_suite.clj
index f4410c88..56ffe354 100644
--- a/src/kaocha/test_suite.clj
+++ b/src/kaocha/test_suite.clj
@@ -2,16 +2,6 @@
   (:require [clojure.test :as t]
             [kaocha.testable :as testable]))
 
-(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 run [testable test-plan]
   (t/do-report {:type :begin-test-suite})
   (let [results (testable/run-testables (:kaocha.test-plan/tests testable) test-plan) 
diff --git a/src/kaocha/testable.clj b/src/kaocha/testable.clj
index a21d9384..3c5fe0cb 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]
@@ -56,16 +57,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.
 
@@ -267,33 +258,20 @@
 (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)))
-
-(defn run-testables
+  (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]
   (if (:parallel *config*)
     (doall (run-testables-parallel testables test-plan))