Make Set methods of MDC and MDLC return IDisposable

src/NLog/MappedDiagnosticsContext.cs

@@ -52,6 +52,29 @@ public static class MappedDiagnosticsContext
 
         private static readonly IDictionary<string, object> EmptyDefaultDictionary = new SortHelpers.ReadOnlySingleBucketDictionary<string, object>();
 
+        /// <summary>
+        /// 
+        /// </summary>
+        private class ItemRemover : IDisposable
+        {
+            private readonly string _item;
+            private bool _disposed = false;
+
+            public ItemRemover(string item)
+            {
+                this._item = item;
+            }
+
+            public void Dispose()
+            {
+                if (!_disposed)
+                {
+                    _disposed = true;
+                    Remove(_item);
+                }
+            }
+        }
+
         /// <summary>
         /// Gets the thread-local dictionary
         /// </summary>
@@ -66,6 +89,30 @@ public static class MappedDiagnosticsContext
             return dictionary;
         }
 
+        /// <summary>
+        /// Sets the current thread MDC item to the specified value.
+        /// </summary>
+        /// <param name="item">Item name.</param>
+        /// <param name="value">Item value.</param>
+        /// <returns>An <see cref="IDisposable"/> that can be used to remove the item from the current thread MDC.</returns>
+        public static IDisposable SetScoped(string item, string value)
+        {
+            Set(item, value);
+            return new ItemRemover(item);
+        }
+
+        /// <summary>
+        /// Sets the current thread MDC item to the specified value.
+        /// </summary>
+        /// <param name="item">Item name.</param>
+        /// <param name="value">Item value.</param>
+        /// <returns>>An <see cref="IDisposable"/> that can be used to remove the item from the current thread MDC.</returns>
+        public static IDisposable SetScoped(string item, object value)
+        {
+            Set(item, value);
+            return new ItemRemover(item);
+        }
+
         /// <summary>
         /// Sets the current thread MDC item to the specified value.
         /// </summary>

src/NLog/MappedDiagnosticsLogicalContext.cs

@@ -37,6 +37,7 @@ namespace NLog
 {
     using System;
     using System.Collections.Generic;
+    using System.Threading;
     using Internal;
 
     /// <summary>
@@ -50,6 +51,29 @@ namespace NLog
     /// </remarks>
     public static class MappedDiagnosticsLogicalContext
     {
+        /// <summary>
+        /// 
+        /// </summary>
+        private class ItemRemover : IDisposable
+        {
+            private readonly string _item;
+            //boolean as int to allow the use of Interlocked.Exchange
+            private int _disposed = 0;
+
+            public ItemRemover(string item)
+            {
+                _item = item;
+            }
+
+            public void Dispose()
+            {
+                if (Interlocked.Exchange(ref _disposed, 1) == 0)
+                {
+                    Remove(_item); 
+                }
+            }
+        }
+
         /// <summary>
         /// Simulate ImmutableDictionary behavior (which is not yet part of all .NET frameworks).
         /// In future the real ImmutableDictionary could be used here to minimize memory usage and copying time.
@@ -110,6 +134,30 @@ public static object GetObject(string item)
             return value;
         }
 
+        /// <summary>
+        /// Sets the current logical context item to the specified value.
+        /// </summary>
+        /// <param name="item">Item name.</param>
+        /// <param name="value">Item value.</param>
+        /// <returns>>An <see cref="IDisposable"/> that can be used to remove the item from the current logical context.</returns>
+        public static IDisposable SetScoped(string item, string value)
+        {
+            Set(item, value);
+            return new ItemRemover(item);
+        }
+
+        /// <summary>
+        /// Sets the current logical context item to the specified value.
+        /// </summary>
+        /// <param name="item">Item name.</param>
+        /// <param name="value">Item value.</param>
+        /// <returns>>An <see cref="IDisposable"/> that can be used to remove the item from the current logical context.</returns>
+        public static IDisposable SetScoped(string item, object value)
+        {
+            Set(item, value);
+            return new ItemRemover(item);
+        }
+
         /// <summary>
         /// Sets the current logical context item to the specified value.
         /// </summary>

tests/NLog.UnitTests/Contexts/MappedDiagnosticsContextTests.cs

@@ -214,5 +214,39 @@ public void timer_cannot_inherit_mappedcontext()
             Assert.Null(getObject);
             Assert.Empty(getValue);
         }
+
+        [Fact]
+        public void disposable_removes_item()
+        {
+            const string itemNotRemovedKey = "itemNotRemovedKey";
+            const string itemRemovedKey = "itemRemovedKey";
+
+            MappedDiagnosticsContext.Clear();
+            MappedDiagnosticsContext.Set(itemNotRemovedKey, "itemNotRemoved");
+            using (MappedDiagnosticsContext.SetScoped(itemRemovedKey, "itemRemoved"))
+            {
+                Assert.Equal(MappedDiagnosticsContext.GetNames(), new[] { itemNotRemovedKey, itemRemovedKey });
+            }
+
+            Assert.Equal(MappedDiagnosticsContext.GetNames(), new[] { itemNotRemovedKey });
+        }
+
+        [Fact]
+        public void dispose_is_idempotent()
+        {
+            const string itemKey = "itemKey";
+
+            MappedDiagnosticsContext.Clear();            
+            IDisposable disposable = MappedDiagnosticsContext.SetScoped(itemKey, "item1");
+
+            disposable.Dispose();
+            Assert.False(MappedDiagnosticsContext.Contains(itemKey));
+
+            //This item shouldn't be removed since it is not the disposable one
+            MappedDiagnosticsContext.Set(itemKey, "item2");
+            disposable.Dispose();
+
+            Assert.True(MappedDiagnosticsContext.Contains(itemKey));
+        }
     }
 }

tests/NLog.UnitTests/Contexts/MappedDiagnosticsLogicalContextTests.cs

@@ -314,5 +314,39 @@ public void timer_cannot_inherit_mappedcontext()
             Assert.Null(getObject);
             Assert.Empty(getValue);
         }
+
+        [Fact]
+        public void disposable_removes_item()
+        {
+            const string itemNotRemovedKey = "itemNotRemovedKey";
+            const string itemRemovedKey = "itemRemovedKey";
+
+            MappedDiagnosticsLogicalContext.Clear();
+            MappedDiagnosticsLogicalContext.Set(itemNotRemovedKey, "itemNotRemoved");
+            using (MappedDiagnosticsLogicalContext.SetScoped(itemRemovedKey, "itemRemoved"))
+            {
+                Assert.Equal(MappedDiagnosticsLogicalContext.GetNames(), new[] { itemNotRemovedKey, itemRemovedKey });
+            }
+
+            Assert.Equal(MappedDiagnosticsLogicalContext.GetNames(), new[] { itemNotRemovedKey });
+        }
+
+        [Fact]
+        public void dispose_is_idempotent()
+        {
+            const string itemKey = "itemKey";
+
+            MappedDiagnosticsLogicalContext.Clear();
+            IDisposable disposable = MappedDiagnosticsLogicalContext.SetScoped(itemKey, "item1");
+
+            disposable.Dispose();
+            Assert.False(MappedDiagnosticsLogicalContext.Contains(itemKey));
+
+            //This item shouldn't be removed since it is not the disposable one
+            MappedDiagnosticsLogicalContext.Set(itemKey, "item2");
+            disposable.Dispose();
+
+            Assert.True(MappedDiagnosticsLogicalContext.Contains(itemKey));
+        }
     }
 }