Support nested putCloseable invocations (issues/404)

Signed-off-by: Matthew Pearson <[email protected]>
qos-ch · Mar 15, 2024 · 1f438ae · 1f438ae
1 parent 448bca1
commit 1f438ae
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 3 deletions.
diff --git a/slf4j-api/src/main/java/org/slf4j/MDC.java b/slf4j-api/src/main/java/org/slf4j/MDC.java
@@ -75,13 +75,19 @@ public class MDC {
      */
     public static class MDCCloseable implements Closeable {
         private final String key;
+        private final String previousValue;
 
-        private MDCCloseable(String key) {
+        private MDCCloseable(String key, String previousValue) {
             this.key = key;
+            this.previousValue = previousValue;
         }
 
         public void close() {
-            MDC.remove(this.key);
+            if (previousValue != null) {
+                MDC.put(key, previousValue);
+            } else {
+                MDC.remove(key);
+            }
         }
     }
 
@@ -156,8 +162,9 @@ public static void put(String key, String val) throws IllegalArgumentException {
      *           in case the "key" parameter is null
      */
     public static MDCCloseable putCloseable(String key, String val) throws IllegalArgumentException {
+        String prev = get(key);
         put(key, val);
-        return new MDCCloseable(key);
+        return new MDCCloseable(key, prev);
     }
 
     /**

diff --git a/slf4j-reload4j/src/test/java/org/slf4j/reload4j/InvocationTest.java b/slf4j-reload4j/src/test/java/org/slf4j/reload4j/InvocationTest.java
@@ -206,4 +206,27 @@ public void testCallerInfoWithFluentAPI() {
         assertEquals(this.getClass().getName(), event.getLocationInformation().getClassName());
     }
 
+    @Test
+    public void testMDCCloseable() {
+        MDC.MDCCloseable closeable = MDC.putCloseable("k", "v");
+        assertEquals("v", MDC.get("k"));
+        closeable.close();
+        assertNull(MDC.get("k"));
+        MDC.clear();
+    }
+
+    @Test
+    public void testMDCCloseablePreExistingKey() {
+        MDC.put("k", "1");
+        MDC.MDCCloseable outerCloseable = MDC.putCloseable("k", "2");
+        assertEquals("2", MDC.get("k"));
+        MDC.MDCCloseable innerCloseable = MDC.putCloseable("k", "3");
+        assertEquals("3", MDC.get("k"));
+        innerCloseable.close();
+        assertEquals("2", MDC.get("k"));
+        outerCloseable.close();
+        assertEquals("1", MDC.get("k"));
+        MDC.clear();
+        assertNull(MDC.get("k"));
+    }
 }