Nullable: AsyncLocal (#23618)

* Nullable: AsyncLocal

* Address PR feedback

4 files changed, 55 insertions, 54 deletions

diff --git a/src/System.Private.CoreLib/shared/System/Threading/AsyncLocal.cs b/src/System.Private.CoreLib/shared/System/Threading/AsyncLocal.cs
index 12d6502432..463289ae55 100644
--- a/src/System.Private.CoreLib/shared/System/Threading/AsyncLocal.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/AsyncLocal.cs
@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+#nullable enable
 using System.Collections.Generic;
 using System.Diagnostics;
 
@@ -38,7 +39,7 @@ namespace System.Threading
     //
     public sealed class AsyncLocal<T> : IAsyncLocal
     {
-        private readonly Action<AsyncLocalValueChangedArgs<T>> m_valueChangedHandler;
+        private readonly Action<AsyncLocalValueChangedArgs<T>>? m_valueChangedHandler;
 
         //
         // Constructs an AsyncLocal<T> that does not receive change notifications.
@@ -51,7 +52,7 @@ namespace System.Threading
         // Constructs an AsyncLocal<T> with a delegate that is called whenever the current value changes
         // on any thread.
         //
-        public AsyncLocal(Action<AsyncLocalValueChangedArgs<T>> valueChangedHandler)
+        public AsyncLocal(Action<AsyncLocalValueChangedArgs<T>>? valueChangedHandler)
         {
             m_valueChangedHandler = valueChangedHandler;
         }
@@ -60,8 +61,8 @@ namespace System.Threading
         {
             get
             {
-                object obj = ExecutionContext.GetLocalValue(this);
-                return (obj == null) ? default : (T)obj;
+                object? obj = ExecutionContext.GetLocalValue(this);
+                return (obj == null) ? default! : (T)obj; // TODO-NULLABLE-GENERIC
             }
             set
             {
@@ -69,11 +70,11 @@ namespace System.Threading
             }
         }
 
-        void IAsyncLocal.OnValueChanged(object previousValueObj, object currentValueObj, bool contextChanged)
+        void IAsyncLocal.OnValueChanged(object? previousValueObj, object? currentValueObj, bool contextChanged)
         {
             Debug.Assert(m_valueChangedHandler != null);
-            T previousValue = previousValueObj == null ? default : (T)previousValueObj;
-            T currentValue = currentValueObj == null ? default : (T)currentValueObj;
+            T previousValue = previousValueObj == null ? default! : (T)previousValueObj; // TODO-NULLABLE-GENERIC
+            T currentValue = currentValueObj == null ? default! : (T)currentValueObj; // TODO-NULLABLE-GENERIC
             m_valueChangedHandler(new AsyncLocalValueChangedArgs<T>(previousValue, currentValue, contextChanged));
         }
     }
@@ -83,7 +84,7 @@ namespace System.Threading
     //
     internal interface IAsyncLocal
     {
-        void OnValueChanged(object previousValue, object currentValue, bool contextChanged);
+        void OnValueChanged(object? previousValue, object? currentValue, bool contextChanged);
     }
 
     public readonly struct AsyncLocalValueChangedArgs<T>
@@ -112,8 +113,8 @@ namespace System.Threading
     //
     internal interface IAsyncLocalValueMap
     {
-        bool TryGetValue(IAsyncLocal key, out object value);
-        IAsyncLocalValueMap Set(IAsyncLocal key, object value, bool treatNullValueAsNonexistent);
+        bool TryGetValue(IAsyncLocal key, out object? value);
+        IAsyncLocalValueMap Set(IAsyncLocal key, object? value, bool treatNullValueAsNonexistent);
     }
 
     //
@@ -131,7 +132,7 @@ namespace System.Threading
             return asyncLocalValueMap == Empty;
         }
 
-        public static IAsyncLocalValueMap Create(IAsyncLocal key, object value, bool treatNullValueAsNonexistent)
+        public static IAsyncLocalValueMap Create(IAsyncLocal key, object? value, bool treatNullValueAsNonexistent)
         {
             // If the value isn't null or a null value may not be treated as nonexistent, then create a new one-element map
             // to store the key/value pair.  Otherwise, use the empty map.
@@ -143,7 +144,7 @@ namespace System.Threading
         // Instance without any key/value pairs.  Used as a singleton/
         private sealed class EmptyAsyncLocalValueMap : IAsyncLocalValueMap
         {
-            public IAsyncLocalValueMap Set(IAsyncLocal key, object value, bool treatNullValueAsNonexistent)
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object? value, bool treatNullValueAsNonexistent)
             {
                 // If the value isn't null or a null value may not be treated as nonexistent, then create a new one-element map
                 // to store the key/value pair.  Otherwise, use the empty map.
@@ -152,7 +153,7 @@ namespace System.Threading
                     (IAsyncLocalValueMap)this;
             }
 
-            public bool TryGetValue(IAsyncLocal key, out object value)
+            public bool TryGetValue(IAsyncLocal key, out object? value)
             {
                 value = null;
                 return false;
@@ -163,14 +164,14 @@ namespace System.Threading
         private sealed class OneElementAsyncLocalValueMap : IAsyncLocalValueMap
         {
             private readonly IAsyncLocal _key1;
-            private readonly object _value1;
+            private readonly object? _value1;
 
-            public OneElementAsyncLocalValueMap(IAsyncLocal key, object value)
+            public OneElementAsyncLocalValueMap(IAsyncLocal key, object? value)
             {
                 _key1 = key; _value1 = value;
             }
 
-            public IAsyncLocalValueMap Set(IAsyncLocal key, object value, bool treatNullValueAsNonexistent)
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object? value, bool treatNullValueAsNonexistent)
             {
                 if (value != null || !treatNullValueAsNonexistent)
                 {
@@ -190,7 +191,7 @@ namespace System.Threading
                 }
             }
 
-            public bool TryGetValue(IAsyncLocal key, out object value)
+            public bool TryGetValue(IAsyncLocal key, out object? value)
             {
                 if (ReferenceEquals(key, _key1))
                 {
@@ -209,15 +210,15 @@ namespace System.Threading
         private sealed class TwoElementAsyncLocalValueMap : IAsyncLocalValueMap
         {
             private readonly IAsyncLocal _key1, _key2;
-            private readonly object _value1, _value2;
+            private readonly object? _value1, _value2;
 
-            public TwoElementAsyncLocalValueMap(IAsyncLocal key1, object value1, IAsyncLocal key2, object value2)
+            public TwoElementAsyncLocalValueMap(IAsyncLocal key1, object? value1, IAsyncLocal key2, object? value2)
             {
                 _key1 = key1; _value1 = value1;
                 _key2 = key2; _value2 = value2;
             }
 
-            public IAsyncLocalValueMap Set(IAsyncLocal key, object value, bool treatNullValueAsNonexistent)
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object? value, bool treatNullValueAsNonexistent)
             {
                 if (value != null || !treatNullValueAsNonexistent)
                 {
@@ -239,7 +240,7 @@ namespace System.Threading
                 }
             }
 
-            public bool TryGetValue(IAsyncLocal key, out object value)
+            public bool TryGetValue(IAsyncLocal key, out object? value)
             {
                 if (ReferenceEquals(key, _key1))
                 {
@@ -263,16 +264,16 @@ namespace System.Threading
         private sealed class ThreeElementAsyncLocalValueMap : IAsyncLocalValueMap
         {
             private readonly IAsyncLocal _key1, _key2, _key3;
-            private readonly object _value1, _value2, _value3;
+            private readonly object? _value1, _value2, _value3;
 
-            public ThreeElementAsyncLocalValueMap(IAsyncLocal key1, object value1, IAsyncLocal key2, object value2, IAsyncLocal key3, object value3)
+            public ThreeElementAsyncLocalValueMap(IAsyncLocal key1, object? value1, IAsyncLocal key2, object? value2, IAsyncLocal key3, object? value3)
             {
                 _key1 = key1; _value1 = value1;
                 _key2 = key2; _value2 = value2;
                 _key3 = key3; _value3 = value3;
             }
 
-            public IAsyncLocalValueMap Set(IAsyncLocal key, object value, bool treatNullValueAsNonexistent)
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object? value, bool treatNullValueAsNonexistent)
             {
                 if (value != null || !treatNullValueAsNonexistent)
                 {
@@ -303,7 +304,7 @@ namespace System.Threading
                 }
             }
 
-            public bool TryGetValue(IAsyncLocal key, out object value)
+            public bool TryGetValue(IAsyncLocal key, out object? value)
             {
                 if (ReferenceEquals(key, _key1))
                 {
@@ -332,21 +333,21 @@ namespace System.Threading
         private sealed class MultiElementAsyncLocalValueMap : IAsyncLocalValueMap
         {
             internal const int MaxMultiElements = 16;
-            private readonly KeyValuePair<IAsyncLocal, object>[] _keyValues;
+            private readonly KeyValuePair<IAsyncLocal, object?>[] _keyValues;
 
             internal MultiElementAsyncLocalValueMap(int count)
             {
                 Debug.Assert(count <= MaxMultiElements);
-                _keyValues = new KeyValuePair<IAsyncLocal, object>[count];
+                _keyValues = new KeyValuePair<IAsyncLocal, object?>[count];
             }
 
-            internal void UnsafeStore(int index, IAsyncLocal key, object value)
+            internal void UnsafeStore(int index, IAsyncLocal key, object? value)
             {
                 Debug.Assert(index < _keyValues.Length);
-                _keyValues[index] = new KeyValuePair<IAsyncLocal, object>(key, value);
+                _keyValues[index] = new KeyValuePair<IAsyncLocal, object?>(key, value);
             }
 
-            public IAsyncLocalValueMap Set(IAsyncLocal key, object value, bool treatNullValueAsNonexistent)
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object? value, bool treatNullValueAsNonexistent)
             {
                 // Find the key in this map.
                 for (int i = 0; i < _keyValues.Length; i++)
@@ -360,7 +361,7 @@ namespace System.Threading
                             // overwriting the old.
                             var multi = new MultiElementAsyncLocalValueMap(_keyValues.Length);
                             Array.Copy(_keyValues, 0, multi._keyValues, 0, _keyValues.Length);
-                            multi._keyValues[i] = new KeyValuePair<IAsyncLocal, object>(key, value);
+                            multi._keyValues[i] = new KeyValuePair<IAsyncLocal, object?>(key, value);
                             return multi;
                         }
                         else if (_keyValues.Length == 4)
@@ -399,13 +400,13 @@ namespace System.Threading
                 {
                     var multi = new MultiElementAsyncLocalValueMap(_keyValues.Length + 1);
                     Array.Copy(_keyValues, 0, multi._keyValues, 0, _keyValues.Length);
-                    multi._keyValues[_keyValues.Length] = new KeyValuePair<IAsyncLocal, object>(key, value);
+                    multi._keyValues[_keyValues.Length] = new KeyValuePair<IAsyncLocal, object?>(key, value);
                     return multi;
                 }
 
                 // Otherwise, upgrade to a many map.
                 var many = new ManyElementAsyncLocalValueMap(MaxMultiElements + 1);
-                foreach (KeyValuePair<IAsyncLocal, object> pair in _keyValues)
+                foreach (KeyValuePair<IAsyncLocal, object?> pair in _keyValues)
                 {
                     many[pair.Key] = pair.Value;
                 }
@@ -413,9 +414,9 @@ namespace System.Threading
                 return many;
             }
 
-            public bool TryGetValue(IAsyncLocal key, out object value)
+            public bool TryGetValue(IAsyncLocal key, out object? value)
             {
-                foreach (KeyValuePair<IAsyncLocal, object> pair in _keyValues)
+                foreach (KeyValuePair<IAsyncLocal, object?> pair in _keyValues)
                 {
                     if (ReferenceEquals(key, pair.Key))
                     {
@@ -429,11 +430,11 @@ namespace System.Threading
         }
 
         // Instance with any number of key/value pairs.
-        private sealed class ManyElementAsyncLocalValueMap : Dictionary<IAsyncLocal, object>, IAsyncLocalValueMap
+        private sealed class ManyElementAsyncLocalValueMap : Dictionary<IAsyncLocal, object?>, IAsyncLocalValueMap
         {
             public ManyElementAsyncLocalValueMap(int capacity) : base(capacity) { }
 
-            public IAsyncLocalValueMap Set(IAsyncLocal key, object value, bool treatNullValueAsNonexistent)
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object? value, bool treatNullValueAsNonexistent)
             {
                 int count = Count;
                 bool containsKey = ContainsKey(key);
@@ -443,7 +444,7 @@ namespace System.Threading
                 if (value != null || !treatNullValueAsNonexistent)
                 {
                     var map = new ManyElementAsyncLocalValueMap(count + (containsKey ? 0 : 1));
-                    foreach (KeyValuePair<IAsyncLocal, object> pair in this)
+                    foreach (KeyValuePair<IAsyncLocal, object?> pair in this)
                     {
                         map[pair.Key] = pair.Value;
                     }
@@ -477,7 +478,7 @@ namespace System.Threading
                     else
                     {
                         var map = new ManyElementAsyncLocalValueMap(count - 1);
-                        foreach (KeyValuePair<IAsyncLocal, object> pair in this)
+                        foreach (KeyValuePair<IAsyncLocal, object?> pair in this)
                         {
                             if (!ReferenceEquals(key, pair.Key))
                             {
diff --git a/src/System.Private.CoreLib/shared/System/Threading/ExecutionContext.cs b/src/System.Private.CoreLib/shared/System/Threading/ExecutionContext.cs
index 3e759d9b6a..6f8a4d75eb 100644
--- a/src/System.Private.CoreLib/shared/System/Threading/ExecutionContext.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/ExecutionContext.cs
@@ -399,8 +399,8 @@ namespace System.Threading
                     // Both contexts have change notifications, check previousExecutionCtx first
                     foreach (IAsyncLocal local in previousChangeNotifications)
                     {
-                        previousExecutionCtx.m_localValues.TryGetValue(local, out object previousValue);
-                        nextExecutionCtx.m_localValues.TryGetValue(local, out object currentValue);
+                        previousExecutionCtx.m_localValues.TryGetValue(local, out object? previousValue);
+                        nextExecutionCtx.m_localValues.TryGetValue(local, out object? currentValue);
 
                         if (previousValue != currentValue)
                         {
@@ -415,9 +415,9 @@ namespace System.Threading
                         {
                             // If the local has a value in the previous context, we already fired the event 
                             // for that local in the code above.
-                            if (!previousExecutionCtx.m_localValues.TryGetValue(local, out object previousValue))
+                            if (!previousExecutionCtx.m_localValues.TryGetValue(local, out object? previousValue))
                             {
-                                nextExecutionCtx.m_localValues.TryGetValue(local, out object currentValue);
+                                nextExecutionCtx.m_localValues.TryGetValue(local, out object? currentValue);
                                 if (previousValue != currentValue)
                                 {
                                     local.OnValueChanged(previousValue, currentValue, contextChanged: true);
@@ -433,7 +433,7 @@ namespace System.Threading
                     // No current values, so just check previous against null
                     foreach (IAsyncLocal local in previousChangeNotifications)
                     {
-                        previousExecutionCtx.m_localValues.TryGetValue(local, out object previousValue);
+                        previousExecutionCtx.m_localValues.TryGetValue(local, out object? previousValue);
                         if (previousValue != null)
                         {
                             local.OnValueChanged(previousValue, null, contextChanged: true);
@@ -447,7 +447,7 @@ namespace System.Threading
                     // No previous values, so just check current against null
                     foreach (IAsyncLocal local in nextChangeNotifications!) // TODO-NULLABLE: https://github.com/dotnet/roslyn/issues/34665
                     {
-                        nextExecutionCtx.m_localValues.TryGetValue(local, out object currentValue);
+                        nextExecutionCtx.m_localValues.TryGetValue(local, out object? currentValue);
                         if (currentValue != null)
                         {
                             local.OnValueChanged(null, currentValue, contextChanged: true);
@@ -479,11 +479,11 @@ namespace System.Threading
 
             Debug.Assert(!current.IsDefault);
             Debug.Assert(current.m_localValues != null, "Only the default context should have null, and we shouldn't be here on the default context");
-            current.m_localValues.TryGetValue(local, out object value);
+            current.m_localValues.TryGetValue(local, out object? value);
             return value;
         }
 
-        internal static void SetLocalValue(IAsyncLocal local, object newValue, bool needChangeNotifications)
+        internal static void SetLocalValue(IAsyncLocal local, object? newValue, bool needChangeNotifications)
         {
             ExecutionContext current = Thread.CurrentThread._executionContext;
 
diff --git a/src/System.Private.CoreLib/shared/System/Threading/LazyInitializer.cs b/src/System.Private.CoreLib/shared/System/Threading/LazyInitializer.cs
index ef03b7e97c..7f28a9fd53 100644
--- a/src/System.Private.CoreLib/shared/System/Threading/LazyInitializer.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/LazyInitializer.cs
@@ -61,7 +61,7 @@ namespace System.Threading
         {
             try
             {
-                Interlocked.CompareExchange(ref target, Activator.CreateInstance<T>(), null!); // TODO-NULLABLE: Need to be able to express ref T nullability
+                Interlocked.CompareExchange(ref target, Activator.CreateInstance<T>(), null!); // TODO-NULLABLE-GENERIC
             }
             catch (MissingMethodException)
             {
@@ -118,7 +118,7 @@ namespace System.Threading
                 throw new InvalidOperationException(SR.Lazy_StaticInit_InvalidOperation);
             }
 
-            Interlocked.CompareExchange(ref target, value, null!); // TODO-NULLABLE: Need to be able to express ref T nullability
+            Interlocked.CompareExchange(ref target, value, null!); // TODO-NULLABLE-GENERIC
             Debug.Assert(target != null);
             return target;
         }
diff --git a/src/System.Private.CoreLib/shared/System/Threading/ThreadLocal.cs b/src/System.Private.CoreLib/shared/System/Threading/ThreadLocal.cs
index b3b3eae064..8c19903ab7 100644
--- a/src/System.Private.CoreLib/shared/System/Threading/ThreadLocal.cs
+++ b/src/System.Private.CoreLib/shared/System/Threading/ThreadLocal.cs
@@ -202,7 +202,7 @@ namespace System.Threading
 
                     // And clear the references from the slot table to the linked slot and the value so that
                     // both can get garbage collected.
-                    slotArray[id].Value!._value = default!; // TODO-NULLABILITY: need a way to annotate Ts
+                    slotArray[id].Value!._value = default!; // TODO-NULLABLE-GENERIC
                     slotArray[id].Value = null;
                 }
             }
@@ -229,7 +229,7 @@ namespace System.Threading
         /// </remarks>
         public override string? ToString()
         {
-            return Value!.ToString(); // TODO-NULLABILITY: This is actually documented to null ref if Value is null!
+            return Value!.ToString(); // Throws NullReferenceException as if caller called ToString on the value itself
         }
 
         /// <summary>
@@ -320,7 +320,7 @@ namespace System.Threading
             T value;
             if (_valueFactory == null)
             {
-                value = default!; // TODO-NULLABILITY: Need a way to annotate T as nullable
+                value = default!; // TODO-NULLABLE-GENERIC
             }
             else
             {
@@ -520,7 +520,7 @@ namespace System.Threading
 
                 LinkedSlot? slot;
                 if (slotArray == null || id >= slotArray.Length || (slot = slotArray[id].Value) == null || !_initialized)
-                    return default!; // TODO-NULLABILITY: Need a way to annotate T
+                    return default!; // TODO-NULLABLE-GENERIC
                 return slot._value;
             }
         }
@@ -646,7 +646,7 @@ namespace System.Threading
             internal volatile LinkedSlotVolatile[]? _slotArray;
 
             // The value for this slot.
-            internal T _value = default!; // TODO-NULLABILITY: should not need to initialize this to default!
+            internal T _value = default!; // TODO-NULLABLE-GENERIC
         }
 
         /// <summary>