Parallelize tests with thread-safe version of with-redefs (#38620)

Gotta go fast and gotta stay safe.

We have 413 calls to with-redefs, so I'm not going to try get 'em all yet.

.clj-kondo/hooks/clojure/core.clj

@@ -43,6 +43,9 @@
      methodical.core/remove-aux-method-with-unique-key!
      next.jdbc/execute!
 
+     ;; Definitely thread safe
+     metabase.test.util.dynamic-redefs/patch-vars!
+
      ;; TODO: most of these symbols shouldn't be here, we should go through them and
      ;; find the functions/macros that use them and make sure their names end with !
      ;; best way to do this is try remove each of these and rely on kondo output to find places where it's used

test/metabase/test.clj

@@ -34,6 +34,7 @@
    [metabase.test.redefs :as test.redefs]
    [metabase.test.util :as tu]
    [metabase.test.util.async :as tu.async]
+   [metabase.test.util.dynamic-redefs :as tu.dr]
    [metabase.test.util.i18n :as i18n.tu]
    [metabase.test.util.log :as tu.log]
    [metabase.test.util.misc :as tu.misc]
@@ -314,3 +315,11 @@
 (alter-meta! #'with-temp update :doc str "\n\n  Note: by default, this will execute its body inside a transaction, making
   it thread safe. If it is wrapped in a call to [[metabase.test/test-helpers-set-global-values!]], it will affect the
   global state of the application database.")
+
+;; Cursive does not understand p/import-macro, so we just proxy this manually
+(defmacro with-dynamic-redefs
+  "A thread-safe version of with-redefs. It only support functions, and adds a fair amount of overhead.
+   It works by replacing each original definition with a proxy the first time it is redefined.
+   This proxy uses a dynamic mapping to check whether the function is currently redefined."
+  [bindings & body]
+  `(tu.dr/with-dynamic-redefs ~bindings ~@body))

test/metabase/test/util/dynamic_redefs.clj

+(ns metabase.test.util.dynamic-redefs
+  (:require [medley.core :as m])
+  (:import (clojure.lang Var)))
+
+(set! *warn-on-reflection* true)
+
+(def ^:dynamic *local-redefs*
+  "A thread-local mapping from vars to their most recently bound definition."
+  {})
+
+(defn- get-local-definition
+  "Get the version of this function that is in scope. It is the unpatched version if there is no override."
+  [a-var]
+  (get *local-redefs* a-var
+       (get (meta a-var) ::original)))
+
+(defn- var->proxy
+  "Build a proxy function to intercept the given var. The proxy checks the current scope for what to call."
+  [a-var]
+  (fn [& args]
+    (let [current-f (get-local-definition a-var)]
+      (apply current-f args))))
+
+(defn patch-vars!
+  "Rebind the given vars with proxies that wrap the original functions."
+  [vars]
+  (let [unpatched-vars (remove #(::patched? (meta %)) vars)]
+    (doseq [^Var a-var unpatched-vars]
+      (locking a-var
+        (when-not (::patched? (meta a-var))
+          (let [old-val (.getRawRoot a-var)
+                patch-meta #(assoc % ::original old-val ::patched? true)]
+            (.bindRoot a-var (with-meta (var->proxy a-var)
+                                        (patch-meta (meta (get *local-redefs* a-var)))))
+            (alter-meta! a-var patch-meta)))))))
+
+(defn- sym->var [sym] `(var ~sym))
+
+(defn- bindings->var->definition
+  "Given a with-redefs style binding, return a mapping from each corresponding var to its given replacement."
+  [binding]
+  (m/map-keys sym->var (into {} (partition-all 2) binding)))
+
+(defmacro with-dynamic-redefs
+  "A thread-safe version of with-redefs. It only supports functions, and adds a fair amount of overhead.
+   It works by replacing each original definition with a proxy the first time it is redefined.
+   This proxy uses a dynamic mapping to check whether the function is currently redefined."
+  [bindings & body]
+  (let [var->definition (bindings->var->definition bindings)]
+    `(do
+       (patch-vars! ~(vec (keys var->definition)))
+       (binding [*local-redefs* (merge *local-redefs* ~var->definition)]
+         ~@body))))

test/metabase/test/util_test.clj

@@ -7,7 +7,10 @@
    [metabase.test :as mt]
    [metabase.test.data :as data]
    [metabase.util :as u]
-   [toucan2.core :as t2]))
+   [toucan2.core :as t2])
+  (:import (java.util.concurrent CountDownLatch TimeUnit)))
+
+(set! *warn-on-reflection* true)
 
 (deftest with-temp-vals-in-db-test
   (testing "let's make sure this acutally works right!"
@@ -42,3 +45,44 @@
       (test-util-test-setting))
     (is (= ["A" "B" "C"]
            (test-util-test-setting)))))
+
+(defn- clump [x y] (str x y))
+
+(deftest ^:parallel with-dynamic-redefs-test
+  (testing "Three threads can independently redefine a regular var"
+    (let [n-threads 3
+          latch (CountDownLatch. (inc n-threads))
+          take-latch  #(do
+                         (.countDown latch)
+                         (when-not (.await latch 100 TimeUnit/MILLISECONDS)
+                           (throw (ex-info "Timeout waiting on all threads to pull their latch" {:latch latch}))))]
+
+      (testing "The original definition"
+        (is (= "original" (clump "o" "riginal"))))
+
+      (future
+        (testing "A thread that minds its own business"
+          (is (= "123" (clump 12 3)))
+          (take-latch)
+          (is (= "321" (clump 3 21)))))
+
+      (future
+        (testing "A thread that redefines it in reverse"
+          (mt/with-dynamic-redefs [clump #(str %2 %1)]
+            (is (= "ok" (clump "k" "o")))
+            (take-latch)
+            (is (= "ko" (clump "o" "k"))))))
+
+      (future
+        (testing "A thread that redefines it twice"
+          (mt/with-dynamic-redefs [clump #(str %2 %2)]
+            (is (= "zz" (clump "a" "z")))
+            (mt/with-dynamic-redefs [clump (fn [x _] (str x x))]
+              (is (= "aa" (clump "a" "z")))
+              (take-latch)
+              (is (= "mm" (clump "m" "l"))))
+            (is (= "bb" (clump "a" "b"))))))
+
+      (take-latch)
+      (testing "The original definition survives"
+        (is (= "original" (clump "orig" "inal")))))))

test/metabase/util/malli/defn_test.clj

@@ -4,6 +4,7 @@
    [clojure.test :refer :all]
    [malli.core :as mc]
    [malli.experimental :as mx]
+   [metabase.test :as mt]
    [metabase.util.malli :as mu]
    [metabase.util.malli.defn :as mu.defn]
    [metabase.util.malli.fn :as mu.fn]))
@@ -161,16 +162,16 @@
   (is (= 'Integer
          (-> #'add-ints meta :arglists first meta :tag))))
 
-(deftest defn-forms-are-not-emitted-for-skippable-ns-in-prod-test
+(deftest ^:parallel defn-forms-are-not-emitted-for-skippable-ns-in-prod-test
   (testing "omission in macroexpansion"
     (testing "returns a simple fn*"
-      (with-redefs [mu.fn/instrument-ns? (constantly false)]
+      (mt/with-dynamic-redefs [mu.fn/instrument-ns? (constantly false)]
         (let [expansion (macroexpand `(mu/defn ~'f :- :int [] "foo"))]
           (is (= '(def f
                     "Inputs: []\n  Return: :int" (clojure.core/fn [] "foo"))
                  expansion)))))
     (testing "returns an instrumented fn"
-      (with-redefs [mu.fn/instrument-ns? (constantly true)]
+      (mt/with-dynamic-redefs [mu.fn/instrument-ns? (constantly true)]
         (let [expansion (macroexpand `(mu/defn ~'f :- :int [] "foo"))]
           (is (= '(def f
                     "Inputs: []\n  Return: :int"

test/metabase/util/malli/fn_test.clj

@@ -252,14 +252,14 @@
                     (.resetMeta {:instrument/always true}))]
             (is (true? (mu.fn/instrument-ns? n)))))))))
 
-(deftest instrumentation-can-be-omitted
+(deftest ^:parallel instrumentation-can-be-omitted
   (testing "omission in macroexpansion"
     (testing "returns a simple fn*"
-      (with-redefs [mu.fn/instrument-ns? (constantly false)]
+      (mt/with-dynamic-redefs [mu.fn/instrument-ns? (constantly false)]
         (let [expansion (macroexpand `(mu.fn/fn :- :int [] "foo"))]
           (is (= expansion '(fn* ([] "foo")))))))
     (testing "returns an instrumented fn"
-      (with-redefs [mu.fn/instrument-ns? (constantly true)]
+      (mt/with-dynamic-redefs [mu.fn/instrument-ns? (constantly true)]
         (let [expansion (macroexpand `(mu.fn/fn :- :int [] "foo"))]
           (is (= (take 2 expansion)
                  '(let* [&f (clojure.core/fn [] "foo")]
@@ -271,7 +271,7 @@
            (catch Exception e
              (is (=? {:type ::mu.fn/invalid-output} (ex-data e)))))))
   (testing "when instrument-ns? returns false, unvalidated form is emitted"
-    (with-redefs [mu.fn/instrument-ns? (constantly false)]
+    (mt/with-dynamic-redefs [mu.fn/instrument-ns? (constantly false)]
       ;; we have to use eval here because `mu.fn/fn` is expanded at _read_ time and we want to change the
       ;; expansion via [[mu.fn/instrument-ns?]]. So that's why we call eval here. Could definitely use some
       ;; macroexpansion tests as well.