Lost Baggage: Fix TelemetryHttpModule losing Baggage when switching f…

…rom native to managed threads (#2314) * WIP working on reproducing and fixing lost baggage. * TelemetryHttpModule now restores Baggage. * Updated public api. * Unit tests. * Use SuppressFlow for more accurate test. * CHANGELOG update. * Attempting to fix unstable test. * Attempting to fix unstable test #2.
open-telemetry · Sep 15, 2021 · 2df1a62 · 2df1a62
1 parent 3cc0f7e
commit 2df1a62
Show file tree

Hide file tree

Showing 16 changed files with 355 additions and 153 deletions.
diff --git a/examples/AspNet/Controllers/WeatherForecastController.cs b/examples/AspNet/Controllers/WeatherForecastController.cs
@@ -20,9 +20,11 @@
 using System.Linq;
 using System.Net;
 using System.Net.Http;
+using System.Security.Cryptography;
 using System.Threading.Tasks;
 using System.Web.Http;
 using Examples.AspNet.Models;
+using OpenTelemetry;
 
 namespace Examples.AspNet.Controllers
 {
@@ -71,6 +73,57 @@ public async Task<IEnumerable<WeatherForecast>> Get(int customerId)
             return GetWeatherForecast();
         }
 
+        /// <summary>
+        /// For testing large async operation which causes IIS to jump threads and results in lost AsyncLocals.
+        /// </summary>
+        [Route("data")]
+        [HttpGet]
+        public async Task<string> GetData()
+        {
+            Baggage.SetBaggage("key1", "value1");
+
+            using var rng = RandomNumberGenerator.Create();
+
+            var requestData = new byte[1024 * 1024 * 100];
+            rng.GetBytes(requestData);
+
+            using var client = new HttpClient();
+
+            using var request = new HttpRequestMessage(HttpMethod.Post, this.Url.Content("~/data"));
+
+            request.Content = new ByteArrayContent(requestData);
+
+            using var response = await client.SendAsync(request).ConfigureAwait(false);
+
+            response.EnsureSuccessStatusCode();
+
+            var responseData = await response.Content.ReadAsByteArrayAsync().ConfigureAwait(false);
+
+            return responseData.SequenceEqual(responseData) ? "match" : "mismatch";
+        }
+
+        [Route("data")]
+        [HttpPost]
+        public async Task<HttpResponseMessage> PostData()
+        {
+            string value1 = Baggage.GetBaggage("key1");
+            if (string.IsNullOrEmpty(value1))
+            {
+                throw new InvalidOperationException("Key1 was not found on Baggage.");
+            }
+
+            var stream = await this.Request.Content.ReadAsStreamAsync().ConfigureAwait(false);
+
+            var result = new HttpResponseMessage(HttpStatusCode.OK)
+            {
+                Content = new StreamContent(stream),
+            };
+
+            result.Content.Headers.ContentType = this.Request.Content.Headers.ContentType;
+
+            return result;
+        }
+
         private static IEnumerable<WeatherForecast> GetWeatherForecast()
         {
             var rng = new Random();

diff --git a/examples/AspNet/Web.config b/examples/AspNet/Web.config
@@ -13,7 +13,9 @@
 	</appSettings>
 	<system.web>
 		<compilation debug="true" targetFramework="4.8"/>
-		<httpRuntime targetFramework="4.8"/>
+		<httpRuntime targetFramework="4.8"
+			maxRequestLength="2147483647"
+			executionTimeout="300" />
 	</system.web>
 	<system.webServer>
 		<handlers>
@@ -26,6 +28,11 @@
 			<add name="SuppressInstrumentationHttpModule" type="Examples.AspNet.SuppressInstrumentationHttpModule" preCondition="integratedMode,managedHandler"/>
 			<add name="TelemetryHttpModule" type="OpenTelemetry.Instrumentation.AspNet.TelemetryHttpModule, OpenTelemetry.Instrumentation.AspNet.TelemetryHttpModule" preCondition="integratedMode,managedHandler"/>
 		</modules>
+		<security>
+			<requestFiltering>
+				<requestLimits maxAllowedContentLength="4294967295" />
+			</requestFiltering>
+		</security>
 	</system.webServer>
 	<runtime>
 		<assemblyBinding xmlns="urn:schemas-microsoft-com:asm.v1">

diff --git a/src/OpenTelemetry.Api/.publicApi/net461/PublicAPI.Shipped.txt b/src/OpenTelemetry.Api/.publicApi/net461/PublicAPI.Shipped.txt
@@ -178,9 +178,7 @@ static OpenTelemetry.Context.Propagation.Propagators.DefaultTextMapPropagator.ge
 static OpenTelemetry.Context.RuntimeContext.ContextSlotType.get -> System.Type
 static OpenTelemetry.Context.RuntimeContext.ContextSlotType.set -> void
 static OpenTelemetry.Context.RuntimeContext.GetSlot<T>(string slotName) -> OpenTelemetry.Context.RuntimeContextSlot<T>
-static OpenTelemetry.Context.RuntimeContext.GetValue<T>(string name) -> T
 static OpenTelemetry.Context.RuntimeContext.RegisterSlot<T>(string slotName) -> OpenTelemetry.Context.RuntimeContextSlot<T>
-static OpenTelemetry.Context.RuntimeContext.SetValue<T>(string name, T value) -> void
 static OpenTelemetry.Trace.ActivityExtensions.GetStatus(this System.Diagnostics.Activity activity) -> OpenTelemetry.Trace.Status
 static OpenTelemetry.Trace.ActivityExtensions.RecordException(this System.Diagnostics.Activity activity, System.Exception ex) -> void
 static OpenTelemetry.Trace.ActivityExtensions.SetStatus(this System.Diagnostics.Activity activity, OpenTelemetry.Trace.Status status) -> void

diff --git a/src/OpenTelemetry.Api/.publicApi/net461/PublicAPI.Unshipped.txt b/src/OpenTelemetry.Api/.publicApi/net461/PublicAPI.Unshipped.txt
@@ -1,4 +1,17 @@
+OpenTelemetry.Context.AsyncLocalRuntimeContextSlot<T>.Value.get -> object
+OpenTelemetry.Context.AsyncLocalRuntimeContextSlot<T>.Value.set -> void
+OpenTelemetry.Context.IRuntimeContextSlotValueAccessor
+OpenTelemetry.Context.IRuntimeContextSlotValueAccessor.Value.get -> object
+OpenTelemetry.Context.IRuntimeContextSlotValueAccessor.Value.set -> void
+OpenTelemetry.Context.RemotingRuntimeContextSlot<T>.Value.get -> object
+OpenTelemetry.Context.RemotingRuntimeContextSlot<T>.Value.set -> void
+OpenTelemetry.Context.ThreadLocalRuntimeContextSlot<T>.Value.get -> object
+OpenTelemetry.Context.ThreadLocalRuntimeContextSlot<T>.Value.set -> void
 override OpenTelemetry.Context.AsyncLocalRuntimeContextSlot<T>.Get() -> T
 override OpenTelemetry.Context.AsyncLocalRuntimeContextSlot<T>.Set(T value) -> void
 OpenTelemetry.Context.AsyncLocalRuntimeContextSlot<T>
 OpenTelemetry.Context.AsyncLocalRuntimeContextSlot<T>.AsyncLocalRuntimeContextSlot(string name) -> void
+static OpenTelemetry.Context.RuntimeContext.GetValue(string slotName) -> object
+static OpenTelemetry.Context.RuntimeContext.GetValue<T>(string slotName) -> T
+static OpenTelemetry.Context.RuntimeContext.SetValue(string slotName, object value) -> void
+static OpenTelemetry.Context.RuntimeContext.SetValue<T>(string slotName, T value) -> void
diff --git a/src/OpenTelemetry.Api/.publicApi/netstandard2.0/PublicAPI.Shipped.txt b/src/OpenTelemetry.Api/.publicApi/netstandard2.0/PublicAPI.Shipped.txt
@@ -178,9 +178,7 @@ static OpenTelemetry.Context.Propagation.Propagators.DefaultTextMapPropagator.ge
 static OpenTelemetry.Context.RuntimeContext.ContextSlotType.get -> System.Type
 static OpenTelemetry.Context.RuntimeContext.ContextSlotType.set -> void
 static OpenTelemetry.Context.RuntimeContext.GetSlot<T>(string slotName) -> OpenTelemetry.Context.RuntimeContextSlot<T>
-static OpenTelemetry.Context.RuntimeContext.GetValue<T>(string name) -> T
 static OpenTelemetry.Context.RuntimeContext.RegisterSlot<T>(string slotName) -> OpenTelemetry.Context.RuntimeContextSlot<T>
-static OpenTelemetry.Context.RuntimeContext.SetValue<T>(string name, T value) -> void
 static OpenTelemetry.Trace.ActivityExtensions.GetStatus(this System.Diagnostics.Activity activity) -> OpenTelemetry.Trace.Status
 static OpenTelemetry.Trace.ActivityExtensions.RecordException(this System.Diagnostics.Activity activity, System.Exception ex) -> void
 static OpenTelemetry.Trace.ActivityExtensions.SetStatus(this System.Diagnostics.Activity activity, OpenTelemetry.Trace.Status status) -> void

diff --git a/src/OpenTelemetry.Api/.publicApi/netstandard2.0/PublicAPI.Unshipped.txt b/src/OpenTelemetry.Api/.publicApi/netstandard2.0/PublicAPI.Unshipped.txt
@@ -0,0 +1,11 @@
+OpenTelemetry.Context.AsyncLocalRuntimeContextSlot<T>.Value.get -> object
+OpenTelemetry.Context.AsyncLocalRuntimeContextSlot<T>.Value.set -> void
+OpenTelemetry.Context.IRuntimeContextSlotValueAccessor
+OpenTelemetry.Context.IRuntimeContextSlotValueAccessor.Value.get -> object
+OpenTelemetry.Context.IRuntimeContextSlotValueAccessor.Value.set -> void
+OpenTelemetry.Context.ThreadLocalRuntimeContextSlot<T>.Value.get -> object
+OpenTelemetry.Context.ThreadLocalRuntimeContextSlot<T>.Value.set -> void
+static OpenTelemetry.Context.RuntimeContext.GetValue(string slotName) -> object
+static OpenTelemetry.Context.RuntimeContext.GetValue<T>(string slotName) -> T
+static OpenTelemetry.Context.RuntimeContext.SetValue(string slotName, object value) -> void
+static OpenTelemetry.Context.RuntimeContext.SetValue<T>(string slotName, T value) -> void
diff --git a/src/OpenTelemetry.Api/Context/AsyncLocalRuntimeContextSlot.cs b/src/OpenTelemetry.Api/Context/AsyncLocalRuntimeContextSlot.cs
@@ -23,7 +23,7 @@ namespace OpenTelemetry.Context
     /// The async local implementation of context slot.
     /// </summary>
     /// <typeparam name="T">The type of the underlying value.</typeparam>
-    public class AsyncLocalRuntimeContextSlot<T> : RuntimeContextSlot<T>
+    public class AsyncLocalRuntimeContextSlot<T> : RuntimeContextSlot<T>, IRuntimeContextSlotValueAccessor
     {
         private readonly AsyncLocal<T> slot;
 
@@ -37,6 +37,13 @@ public AsyncLocalRuntimeContextSlot(string name)
             this.slot = new AsyncLocal<T>();
         }
 
+        /// <inheritdoc/>
+        public object Value
+        {
+            get => this.slot.Value;
+            set => this.slot.Value = (T)value;
+        }
+
         /// <inheritdoc/>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public override T Get()

diff --git a/src/OpenTelemetry.Api/Context/IRuntimeContextSlotValueAccessor.cs b/src/OpenTelemetry.Api/Context/IRuntimeContextSlotValueAccessor.cs
@@ -0,0 +1,29 @@
+// <copyright file="IRuntimeContextSlotValueAccessor.cs" company="OpenTelemetry Authors">
+// Copyright The OpenTelemetry Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+// </copyright>
+
+namespace OpenTelemetry.Context
+{
+    /// <summary>
+    /// Describes a type of <see cref="RuntimeContextSlot{T}"/> which can expose its value as an <see cref="object"/>.
+    /// </summary>
+    public interface IRuntimeContextSlotValueAccessor
+    {
+        /// <summary>
+        /// Gets or sets the value of the slot as an <see cref="object"/>.
+        /// </summary>
+        object Value { get; set; }
+    }
+}
diff --git a/src/OpenTelemetry.Api/Context/RemotingRuntimeContextSlot.cs b/src/OpenTelemetry.Api/Context/RemotingRuntimeContextSlot.cs
@@ -26,7 +26,7 @@ namespace OpenTelemetry.Context
     /// The .NET Remoting implementation of context slot.
     /// </summary>
     /// <typeparam name="T">The type of the underlying value.</typeparam>
-    public class RemotingRuntimeContextSlot<T> : RuntimeContextSlot<T>
+    public class RemotingRuntimeContextSlot<T> : RuntimeContextSlot<T>, IRuntimeContextSlotValueAccessor
     {
         // A special workaround to suppress context propagation cross AppDomains.
         //
@@ -50,24 +50,26 @@ public RemotingRuntimeContextSlot(string name)
         {
         }
 
+        /// <inheritdoc/>
+        public object Value
+        {
+            get => this.Get();
+            set => this.Set((T)value);
+        }
+
         /// <inheritdoc/>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public override T Get()
         {
-            var wrapper = CallContext.LogicalGetData(this.Name) as BitArray;
-
-            if (wrapper == null)
+            if (!(CallContext.LogicalGetData(this.Name) is BitArray wrapper))
             {
-                return default(T);
+                return default;
             }
 
             var value = WrapperField.GetValue(wrapper);
-            if (value is T)
-            {
-                return (T)value;
-            }
-
-            return default(T);
+            return value is T t
+                ? t
+                : default;
         }
 
         /// <inheritdoc/>

diff --git a/src/OpenTelemetry.Api/Context/RuntimeContext.cs b/src/OpenTelemetry.Api/Context/RuntimeContext.cs
@@ -73,8 +73,12 @@ public static RuntimeContextSlot<T> GetSlot<T>(string slotName)
                 throw new ArgumentException($"{nameof(slotName)} cannot be null or empty string.");
             }
 
-            Slots.TryGetValue(slotName, out var slot);
-            return slot as RuntimeContextSlot<T> ?? throw new ArgumentException($"The context slot {slotName} is not found.");
+            if (!Slots.TryGetValue(slotName, out var slot))
+            {
+                throw new ArgumentException($"The context slot {slotName} could not be found.");
+            }
+
+            return slot as RuntimeContextSlot<T> ?? throw new ArgumentException($"The context slot {slotName} cannot be cast as {typeof(RuntimeContextSlot<T>)}.");
         }
 
         /*
@@ -104,27 +108,76 @@ public static IDictionary<string, object> Snapshot()
         /// <summary>
         /// Sets the value to a registered slot.
         /// </summary>
-        /// <param name="name">The name of the context slot.</param>
+        /// <param name="slotName">The name of the context slot.</param>
         /// <param name="value">The value to be set.</param>
         /// <typeparam name="T">The type of the value.</typeparam>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public static void SetValue<T>(string name, T value)
+        public static void SetValue<T>(string slotName, T value)
         {
-            var slot = (RuntimeContextSlot<T>)Slots[name];
-            slot.Set(value);
+            GetSlot<T>(slotName).Set(value);
         }
 
         /// <summary>
         /// Gets the value from a registered slot.
         /// </summary>
-        /// <param name="name">The name of the context slot.</param>
+        /// <param name="slotName">The name of the context slot.</param>
         /// <typeparam name="T">The type of the value.</typeparam>
         /// <returns>The value retrieved from the context slot.</returns>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        public static T GetValue<T>(string name)
+        public static T GetValue<T>(string slotName)
+        {
+            return GetSlot<T>(slotName).Get();
+        }
+
+        /// <summary>
+        /// Sets the value to a registered slot.
+        /// </summary>
+        /// <param name="slotName">The name of the context slot.</param>
+        /// <param name="value">The value to be set.</param>
+        public static void SetValue(string slotName, object value)
         {
-            var slot = (RuntimeContextSlot<T>)Slots[name];
-            return slot.Get();
+            if (string.IsNullOrEmpty(slotName))
+            {
+                throw new ArgumentException($"{nameof(slotName)} cannot be null or empty string.");
+            }
+
+            if (!Slots.TryGetValue(slotName, out var slot))
+            {
+                throw new ArgumentException($"The context slot {slotName} could not be found.");
+            }
+
+            if (slot is IRuntimeContextSlotValueAccessor runtimeContextSlotValueAccessor)
+            {
+                runtimeContextSlotValueAccessor.Value = value;
+                return;
+            }
+
+            throw new NotSupportedException($"The context slot {slotName} value cannot be accessed as an object.");
+        }
+
+        /// <summary>
+        /// Gets the value from a registered slot.
+        /// </summary>
+        /// <param name="slotName">The name of the context slot.</param>
+        /// <returns>The value retrieved from the context slot.</returns>
+        public static object GetValue(string slotName)
+        {
+            if (string.IsNullOrEmpty(slotName))
+            {
+                throw new ArgumentException($"{nameof(slotName)} cannot be null or empty string.");
+            }
+
+            if (!Slots.TryGetValue(slotName, out var slot))
+            {
+                throw new ArgumentException($"The context slot {slotName} could not be found.");
+            }
+
+            if (slot is IRuntimeContextSlotValueAccessor runtimeContextSlotValueAccessor)
+            {
+                return runtimeContextSlotValueAccessor.Value;
+            }
+
+            throw new NotSupportedException($"The context slot {slotName} value cannot be accessed as an object.");
         }
 
         // For testing purpose

diff --git a/src/OpenTelemetry.Api/Context/ThreadLocalRuntimeContextSlot.cs b/src/OpenTelemetry.Api/Context/ThreadLocalRuntimeContextSlot.cs
@@ -23,7 +23,7 @@ namespace OpenTelemetry.Context
     /// The thread local (TLS) implementation of context slot.
     /// </summary>
     /// <typeparam name="T">The type of the underlying value.</typeparam>
-    public class ThreadLocalRuntimeContextSlot<T> : RuntimeContextSlot<T>
+    public class ThreadLocalRuntimeContextSlot<T> : RuntimeContextSlot<T>, IRuntimeContextSlotValueAccessor
     {
         private readonly ThreadLocal<T> slot;
         private bool disposedValue;
@@ -38,6 +38,13 @@ public ThreadLocalRuntimeContextSlot(string name)
             this.slot = new ThreadLocal<T>();
         }
 
+        /// <inheritdoc/>
+        public object Value
+        {
+            get => this.slot.Value;
+            set => this.slot.Value = (T)value;
+        }
+
         /// <inheritdoc/>
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
         public override T Get()
@@ -55,14 +62,11 @@ public override void Set(T value)
         /// <inheritdoc/>
         protected override void Dispose(bool disposing)
         {
-            base.Dispose(true);
-            if (!this.disposedValue)
-            {
-                if (disposing)
-                {
-                    this.slot.Dispose();
-                }
+            base.Dispose(disposing);
 
+            if (disposing && !this.disposedValue)
+            {
+                this.slot.Dispose();
                 this.disposedValue = true;
             }
         }