From a4575581f17b63940173c098b3434afbd0fdc77f Mon Sep 17 00:00:00 2001
From: Utkarsh Umesan Pillai <66651184+utpilla@users.noreply.github.com>
Date: Tue, 6 Jun 2023 13:04:33 -0700
Subject: [PATCH] Fix AOT warnings for RuntimeContext (#4542)

---
 src/OpenTelemetry.Api/CHANGELOG.md            |  9 ++++
 .../Context/RuntimeContext.cs                 | 50 +++++++++++++++++--
 .../AotCompatibilityTests.cs                  |  2 +-
 3 files changed, 56 insertions(+), 5 deletions(-)

diff --git a/src/OpenTelemetry.Api/CHANGELOG.md b/src/OpenTelemetry.Api/CHANGELOG.md
index c87f1b4b3..7c4c03093 100644
--- a/src/OpenTelemetry.Api/CHANGELOG.md
+++ b/src/OpenTelemetry.Api/CHANGELOG.md
@@ -2,6 +2,15 @@
 
 ## Unreleased
 
+* **Breaking change** In order to make `RuntimeContext` compatible with
+  ahead-of-time compilation (AOT),
+  `RuntimeContext.ContextSlotType` can only be assigned one
+  of the following types: `AsyncLocalRuntimeContextSlot<>`,
+  `ThreadLocalRuntimeContextSlot<>`, and `RemotingRuntimeContextSlot<>`. A
+  `System.NotSupportedException` will be thrown if you try to assign any type
+  other than the three types mentioned.
+  ([#4542](https://github.com/open-telemetry/opentelemetry-dotnet/pull/4542))
+
 ## 1.5.0
 
 Released 2023-Jun-05
diff --git a/src/OpenTelemetry.Api/Context/RuntimeContext.cs b/src/OpenTelemetry.Api/Context/RuntimeContext.cs
index 6ff9699f9..a8f45c405 100644
--- a/src/OpenTelemetry.Api/Context/RuntimeContext.cs
+++ b/src/OpenTelemetry.Api/Context/RuntimeContext.cs
@@ -27,10 +27,38 @@ namespace OpenTelemetry.Context
     {
         private static readonly ConcurrentDictionary<string, object> Slots = new();
 
+        private static Type contextSlotType = typeof(AsyncLocalRuntimeContextSlot<>);
+
         /// <summary>
         /// Gets or sets the actual context carrier implementation.
         /// </summary>
-        public static Type ContextSlotType { get; set; } = typeof(AsyncLocalRuntimeContextSlot<>);
+        public static Type ContextSlotType
+        {
+            get => contextSlotType;
+            set
+            {
+                Guard.ThrowIfNull(value, nameof(value));
+
+                if (value == typeof(AsyncLocalRuntimeContextSlot<>))
+                {
+                    contextSlotType = value;
+                }
+                else if (value == typeof(ThreadLocalRuntimeContextSlot<>))
+                {
+                    contextSlotType = value;
+                }
+#if NETFRAMEWORK
+                else if (value == typeof(RemotingRuntimeContextSlot<>))
+                {
+                    contextSlotType = value;
+                }
+#endif
+                else
+                {
+                    throw new NotSupportedException($"{value} is not a supported type.");
+                }
+            }
+        }
 
         /// <summary>
         /// Register a named context slot.
@@ -41,6 +69,7 @@ namespace OpenTelemetry.Context
         public static RuntimeContextSlot<T> RegisterSlot<T>(string slotName)
         {
             Guard.ThrowIfNullOrEmpty(slotName);
+            RuntimeContextSlot<T> slot = null;
 
             lock (Slots)
             {
@@ -49,9 +78,22 @@ namespace OpenTelemetry.Context
                     throw new InvalidOperationException($"Context slot already registered: '{slotName}'");
                 }
 
-                var type = ContextSlotType.MakeGenericType(typeof(T));
-                var ctor = type.GetConstructor(new Type[] { typeof(string) });
-                var slot = (RuntimeContextSlot<T>)ctor.Invoke(new object[] { slotName });
+                if (ContextSlotType == typeof(AsyncLocalRuntimeContextSlot<>))
+                {
+                    slot = new AsyncLocalRuntimeContextSlot<T>(slotName);
+                }
+                else if (ContextSlotType == typeof(ThreadLocalRuntimeContextSlot<>))
+                {
+                    slot = new ThreadLocalRuntimeContextSlot<T>(slotName);
+                }
+
+#if NETFRAMEWORK
+                else if (ContextSlotType == typeof(RemotingRuntimeContextSlot<>))
+                {
+                    slot = new RemotingRuntimeContextSlot<T>(slotName);
+                }
+#endif
+
                 Slots[slotName] = slot;
                 return slot;
             }
diff --git a/test/OpenTelemetry.AotCompatibility.Tests/AotCompatibilityTests.cs b/test/OpenTelemetry.AotCompatibility.Tests/AotCompatibilityTests.cs
index 2b9cf77d9..4db5b9033 100644
--- a/test/OpenTelemetry.AotCompatibility.Tests/AotCompatibilityTests.cs
+++ b/test/OpenTelemetry.AotCompatibility.Tests/AotCompatibilityTests.cs
@@ -85,7 +85,7 @@ namespace OpenTelemetry.AotCompatibility.Tests
             Assert.True(process.ExitCode == 0, "Publishing the AotCompatibility app failed. See test output for more details.");
 
             var warnings = expectedOutput.ToString().Split('\n', '\r').Where(line => line.Contains("warning IL"));
-            Assert.Equal(43, warnings.Count());
+            Assert.Equal(40, warnings.Count());
         }
     }
 }