From 1f438ae32c06e46cc7f6c4acddfed154c7e2eb7c Mon Sep 17 00:00:00 2001
From: Matthew Pearson <matthewpearson@gmail.com>
Date: Fri, 15 Mar 2024 14:38:04 -0700
Subject: [PATCH] Support nested putCloseable invocations (issues/404)

Signed-off-by: Matthew Pearson <matthewpearson@gmail.com>
---
 slf4j-api/src/main/java/org/slf4j/MDC.java    | 13 ++++++++---
 .../org/slf4j/reload4j/InvocationTest.java    | 23 +++++++++++++++++++
 2 files changed, 33 insertions(+), 3 deletions(-)

diff --git a/slf4j-api/src/main/java/org/slf4j/MDC.java b/slf4j-api/src/main/java/org/slf4j/MDC.java
index 6b0feded0..50014585a 100644
--- 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
index 12d6db190..0b4254dba 100644
--- 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"));
+    }
 }