(PDB-5216) Add with-log-mdc to manage the MDC log context

Add a with-log-mdc macro that establishes new MDC contexts during the
execution of its body, and ensures the original values (if any) of
those contexts are always restored in the end.

Author: rbrw

src/puppetlabs/puppetdb/utils.clj

@@ -17,7 +17,8 @@
    [java.net MalformedURLException URISyntaxException URL]
    [java.nio ByteBuffer CharBuffer]
    [java.nio.charset Charset CoderResult StandardCharsets]
-   (java.util.concurrent ScheduledThreadPoolExecutor TimeUnit)))
+   (java.util.concurrent ScheduledThreadPoolExecutor TimeUnit)
+   (org.apache.log4j MDC)))
 
 (defmacro with-captured-throw [& body]
   `(try [(do ~@body)] (catch Throwable ex# ex#)))
@@ -34,6 +35,36 @@
     (apply print args)
     (flush)))
 
+(defmacro with-log-mdc
+  "Establishes the MDC contexts given by the alternating-kvs key value
+  pairs during the execution of the body, and ensures that the
+  original values (if any) are always restored before returning."
+  [alternating-kvs & body]
+  ;; For now, assume this isn't used in performance-critical code,
+  ;; i.e. within tight loops.
+  (when-not (even? (count alternating-kvs))
+    (throw (RuntimeException. "Odd number of MDC key value pairs")))
+  (when-not (every? string? (take-nth 2 alternating-kvs))
+    (throw (RuntimeException. "MDC keys are not all strings")))
+  (loop [[k v & alternating-kvs] alternating-kvs
+         expansion `(do ~@body)]
+    (if-not k
+      expansion
+      (recur alternating-kvs
+             ;; We know k is a string, so it's fine to repeat ~k
+             `(let [v# ~v]
+                (if (nil? v#)
+                  ~expansion
+                  (let [orig# (MDC/get ~k)]
+                    (try
+                      (MDC/put ~k v#)
+                      ~expansion
+                      (finally
+                        ;; After you put a nil value MDC/getContext crashes
+                        (if (nil? orig#)
+                          (MDC/remove ~k)
+                          (MDC/put ~k orig#)))))))))))
+
 (defn flush-and-exit
   "Attempts to flush *out* and *err*, reporting any failures to *err*,
   if possible, and then invokes (System/exit status)."

test/puppetlabs/puppetdb/utils_test.clj

@@ -2,7 +2,9 @@
   (:require [puppetlabs.puppetdb.utils :refer :all]
             [clojure.test :refer :all]
             [puppetlabs.puppetdb.testutils :as tu]
-            [clojure.walk :as walk]))
+            [clojure.walk :as walk])
+  (:import
+   (org.apache.log4j MDC)))
 
 (deftest test-println-err
   (is (= "foo\n"
@@ -17,6 +19,23 @@
   (is (= "foo bar"
          (tu/with-err-str (print-err "foo" "bar")))))
 
+(deftest with-log-mdc-behavior
+  (with-log-mdc ["foo" "bar"]
+    (is (= "bar" (MDC/get "foo"))))
+  (with-log-mdc ["foo" nil]
+    (is (= {} (MDC/getContext))))
+  (with-log-mdc ["foo" nil "bar" 1]
+    (is (= {"bar" "1"} (MDC/getContext))))
+  (with-log-mdc ["foo" 1 "bar" nil]
+    (is (= {"foo" "1"} (MDC/getContext))))
+  (with-log-mdc ["foo" "x" "bar" "y"]
+    (is (= {"foo" "x" "bar" "y"} (MDC/getContext))))
+  (with-log-mdc ["foo" "x" "bar" nil "baz" "z"]
+    (is (= {"foo" "x" "baz" "z"} (MDC/getContext)))
+    (with-log-mdc ["foo" "x" "bar" "y" "baz" "z"]
+      (is (= {"foo" "x" "bar" "y" "baz" "z"} (MDC/getContext))))
+    (is (= {"foo" "x" "baz" "z"} (MDC/getContext)))))
+
 (deftest test-assoc-when
   (is (= {:a 1 :b 2}
          (assoc-when {:a 1 :b 2} :b 100)))