Merge pull request #3088 from dotnet-bot/from-tfs

Merge changes from TFS

9 files changed, 649 insertions, 79 deletions

diff --git a/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.cs b/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.cs
index 82d942c75..5937bae7f 100644
--- a/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.cs
+++ b/src/System.Private.CoreLib/src/Internal/Runtime/Augments/RuntimeThread.cs
@@ -22,6 +22,9 @@ namespace Internal.Runtime.Augments
         [ThreadStatic]
         private static RuntimeThread t_currentThread;
 
+        private ExecutionContext _executionContext;
+        private SynchronizationContext _synchronizationContext;
+
         private volatile int _threadState;
         private ThreadPriority _priority;
         private ManagedThreadId _managedThreadId;
@@ -117,6 +120,18 @@ namespace Internal.Runtime.Augments
             }
         }
 
+        internal ExecutionContext ExecutionContext
+        {
+            get { return _executionContext; }
+            set { _executionContext = value; }
+        }
+
+        internal SynchronizationContext SynchronizationContext
+        {
+            get { return _synchronizationContext; }
+            set { _synchronizationContext = value; }
+        }
+
         public bool IsAlive
         {
             get
diff --git a/src/System.Private.CoreLib/src/Resources/Strings.resx b/src/System.Private.CoreLib/src/Resources/Strings.resx
index 9a4148b5a..cdae3ed92 100644
--- a/src/System.Private.CoreLib/src/Resources/Strings.resx
+++ b/src/System.Private.CoreLib/src/Resources/Strings.resx
@@ -1122,6 +1122,24 @@
   <data name="InvalidOperation_TooEarly" xml:space="preserve">
     <value>Internal Error: This operation cannot be invoked in an eager class constructor.</value>
   </data>
+  <data name="InvalidOperation_NullContext" xml:space="preserve">
+    <value>Cannot call Set on a null context</value>
+  </data>
+  <data name="InvalidOperation_CannotUseAFCOtherThread" xml:space="preserve">
+    <value>AsyncFlowControl object must be used on the thread where it was created.</value>
+  </data>
+  <data name="InvalidOperation_CannotRestoreUnsupressedFlow" xml:space="preserve">
+    <value>Cannot restore context flow when it is not suppressed.</value>
+  </data>
+  <data name="InvalidOperation_CannotSupressFlowMultipleTimes" xml:space="preserve">
+    <value>Context flow is already suppressed.</value>
+  </data>
+  <data name="InvalidOperation_CannotUseAFCMultiple" xml:space="preserve">
+    <value>AsyncFlowControl object can be used only once to call Undo().</value>
+  </data>
+  <data name="InvalidOperation_AsyncFlowCtrlCtxMismatch" xml:space="preserve">
+    <value>AsyncFlowControl objects can be used to restore flow only on a Context that had its flow suppressed.</value>
+  </data>
   <data name="InvalidProgram_Default" xml:space="preserve">
     <value>Common Language Runtime detected an invalid program.</value>
   </data>
@@ -1603,9 +1621,6 @@
   <data name="Argument_MinMaxValue" xml:space="preserve">
     <value>'{0}' cannot be greater than {1}.</value>
   </data>
-  <data name="InvalidOperation_NullContext" xml:space="preserve">
-    <value>Cannot call Set on a null context</value>
-  </data>
   <data name="ExecutionContext_ExceptionInAsyncLocalNotification" xml:space="preserve">
     <value>An exception was not handled in an AsyncLocal&lt;T&gt; notification callback.</value>
   </data>
diff --git a/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs b/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
index 896b89717..74b1a1b77 100644
--- a/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
+++ b/src/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncMethodBuilder.cs
@@ -24,6 +24,7 @@ using CausalityRelation = Internal.Runtime.Augments.CausalityRelation;
 using CausalitySource = Internal.Runtime.Augments.CausalitySource;
 using CausalityTraceLevel = Internal.Runtime.Augments.CausalityTraceLevel;
 using CausalitySynchronousWork = Internal.Runtime.Augments.CausalitySynchronousWork;
+using Thread = Internal.Runtime.Augments.RuntimeThread;
 
 namespace System.Runtime.CompilerServices
 {
@@ -729,10 +730,11 @@ namespace System.Runtime.CompilerServices
             where TStateMachine : IAsyncStateMachine
         {
             // Async state machines are required not to throw, so no need for try/finally here.
+            Thread currentThread = Thread.CurrentThread;
             ExecutionContextSwitcher ecs = default(ExecutionContextSwitcher);
-            ExecutionContext.EstablishCopyOnWriteScope(ref ecs);
+            ExecutionContext.EstablishCopyOnWriteScope(currentThread, ref ecs);
             stateMachine.MoveNext();
-            ecs.Undo();
+            ecs.Undo(currentThread);
         }
 
         //
diff --git a/src/System.Private.CoreLib/src/System/Threading/AsyncLocal.cs b/src/System.Private.CoreLib/src/System/Threading/AsyncLocal.cs
index 8ee2c8b13..59c8fb3c8 100644
--- a/src/System.Private.CoreLib/src/System/Threading/AsyncLocal.cs
+++ b/src/System.Private.CoreLib/src/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.
 
+using System.Collections.Generic;
 using System.Diagnostics;
 
 namespace System.Threading
@@ -35,7 +36,7 @@ namespace System.Threading
     //      NativeMethods.SetThreadCulture(args.CurrentValue.LCID);
     //   });
     //
-    public sealed class AsyncLocal<T> : IAsyncLocal, IEquatable<IAsyncLocal>
+    public sealed class AsyncLocal<T> : IAsyncLocal
     {
         private readonly Action<AsyncLocalValueChangedArgs<T>> m_valueChangedHandler;
 
@@ -75,17 +76,12 @@ namespace System.Threading
             T currentValue = currentValueObj == null ? default(T) : (T)currentValueObj;
             m_valueChangedHandler(new AsyncLocalValueChangedArgs<T>(previousValue, currentValue, contextChanged));
         }
-
-        bool IEquatable<IAsyncLocal>.Equals(IAsyncLocal other)
-        {
-            return this == other;
-        }
     }
 
     //
     // Interface to allow non-generic code in ExecutionContext to call into the generic AsyncLocal<T> type.
     //
-    internal interface IAsyncLocal : IEquatable<IAsyncLocal>
+    internal interface IAsyncLocal
     {
         void OnValueChanged(object previousValue, object currentValue, bool contextChanged);
     }
@@ -109,4 +105,380 @@ namespace System.Threading
             ThreadContextChanged = contextChanged;
         }
     }
+
+    //
+    // Interface used to store an IAsyncLocal => object mapping in ExecutionContext.
+    // Implementations are specialized based on the number of elements in the immutable
+    // map in order to minimize memory consumption and look-up times.
+    //
+    internal interface IAsyncLocalValueMap
+    {
+        bool TryGetValue(IAsyncLocal key, out object value);
+        IAsyncLocalValueMap Set(IAsyncLocal key, object value);
+    }
+
+    //
+    // Utility functions for getting/creating instances of IAsyncLocalValueMap
+    //
+    internal static class AsyncLocalValueMap
+    {
+        public static IAsyncLocalValueMap Empty { get; } = new EmptyAsyncLocalValueMap();
+
+        // Instance without any key/value pairs.  Used as a singleton/
+        private sealed class EmptyAsyncLocalValueMap : IAsyncLocalValueMap
+        {
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
+            {
+                // If the value isn't null, then create a new one-element map to store
+                // the key/value pair.  If it is null, then we're still empty.
+                return value != null ?
+                    new OneElementAsyncLocalValueMap(key, value) :
+                    (IAsyncLocalValueMap)this;
+            }
+
+            public bool TryGetValue(IAsyncLocal key, out object value)
+            {
+                value = null;
+                return false;
+            }
+        }
+
+        // Instance with one key/value pair.
+        private sealed class OneElementAsyncLocalValueMap : IAsyncLocalValueMap
+        {
+            private readonly IAsyncLocal _key1;
+            private readonly object _value1;
+
+            public OneElementAsyncLocalValueMap(IAsyncLocal key, object value)
+            {
+                _key1 = key; _value1 = value;
+            }
+
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
+            {
+                if (value != null)
+                {
+                    // The value is non-null.  If the key matches one already contained in this map,
+                    // then create a new one-element map with the updated value, otherwise create
+                    // a two-element map with the additional key/value.
+                    return ReferenceEquals(key, _key1) ?
+                        new OneElementAsyncLocalValueMap(key, value) :
+                        (IAsyncLocalValueMap)new TwoElementAsyncLocalValueMap(_key1, _value1, key, value);
+                }
+                else
+                {
+                    // The value is null.  If the key exists in this map, remove it by downgrading to an empty map.
+                    // Otherwise, there's nothing to add or remove, so just return this map.
+                    return ReferenceEquals(key, _key1) ?
+                        Empty :
+                        (IAsyncLocalValueMap)this;
+                }
+            }
+
+            public bool TryGetValue(IAsyncLocal key, out object value)
+            {
+                if (ReferenceEquals(key, _key1))
+                {
+                    value = _value1;
+                    return true;
+                }
+                else
+                {
+                    value = null;
+                    return false;
+                }
+            }
+        }
+
+        // Instance with two key/value pairs.
+        private sealed class TwoElementAsyncLocalValueMap : IAsyncLocalValueMap
+        {
+            private readonly IAsyncLocal _key1, _key2;
+            private readonly object _value1, _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)
+            {
+                if (value != null)
+                {
+                    // The value is non-null.  If the key matches one already contained in this map,
+                    // then create a new two-element map with the updated value, otherwise create
+                    // a three-element map with the additional key/value.
+                    return
+                        ReferenceEquals(key, _key1) ? new TwoElementAsyncLocalValueMap(key, value, _key2, _value2) :
+                        ReferenceEquals(key, _key2) ? new TwoElementAsyncLocalValueMap(_key1, _value1, key, value) :
+                        (IAsyncLocalValueMap)new ThreeElementAsyncLocalValueMap(_key1, _value1, _key2, _value2, key, value);
+                }
+                else
+                {
+                    // The value is null.  If the key exists in this map, remove it by downgrading to a one-element map
+                    // without the key.  Otherwise, there's nothing to add or remove, so just return this map.
+                    return
+                        ReferenceEquals(key, _key1) ? new OneElementAsyncLocalValueMap(_key2, _value2) :
+                        ReferenceEquals(key, _key2) ? new OneElementAsyncLocalValueMap(_key1, _value1) :
+                        (IAsyncLocalValueMap)this;
+                }
+            }
+
+            public bool TryGetValue(IAsyncLocal key, out object value)
+            {
+                if (ReferenceEquals(key, _key1))
+                {
+                    value = _value1;
+                    return true;
+                }
+                else if (ReferenceEquals(key, _key2))
+                {
+                    value = _value2;
+                    return true;
+                }
+                else
+                {
+                    value = null;
+                    return false;
+                }
+            }
+        }
+
+        // Instance with three key/value pairs.
+        private sealed class ThreeElementAsyncLocalValueMap : IAsyncLocalValueMap
+        {
+            private readonly IAsyncLocal _key1, _key2, _key3;
+            private readonly object _value1, _value2, _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)
+            {
+                if (value != null)
+                {
+                    // The value is non-null.  If the key matches one already contained in this map,
+                    // then create a new three-element map with the updated value.
+                    if (ReferenceEquals(key, _key1)) return new ThreeElementAsyncLocalValueMap(key, value, _key2, _value2, _key3, _value3);
+                    if (ReferenceEquals(key, _key2)) return new ThreeElementAsyncLocalValueMap(_key1, _value1, key, value, _key3, _value3);
+                    if (ReferenceEquals(key, _key3)) return new ThreeElementAsyncLocalValueMap(_key1, _value1, _key2, _value2, key, value);
+
+                    // The key doesn't exist in this map, so upgrade to a multi map that contains
+                    // the additional key/value pair.
+                    var multi = new MultiElementAsyncLocalValueMap(4);
+                    multi.UnsafeStore(0, _key1, _value1);
+                    multi.UnsafeStore(1, _key2, _value2);
+                    multi.UnsafeStore(2, _key3, _value3);
+                    multi.UnsafeStore(3, key, value);
+                    return multi;
+                }
+                else
+                {
+                    // The value is null.  If the key exists in this map, remove it by downgrading to a two-element map
+                    // without the key.  Otherwise, there's nothing to add or remove, so just return this map.
+                    return
+                        ReferenceEquals(key, _key1) ? new TwoElementAsyncLocalValueMap(_key2, _value2, _key3, _value3) :
+                        ReferenceEquals(key, _key2) ? new TwoElementAsyncLocalValueMap(_key1, _value1, _key3, _value3) :
+                        ReferenceEquals(key, _key3) ? new TwoElementAsyncLocalValueMap(_key1, _value1, _key2, _value2) :
+                        (IAsyncLocalValueMap)this;
+                }
+            }
+
+            public bool TryGetValue(IAsyncLocal key, out object value)
+            {
+                if (ReferenceEquals(key, _key1))
+                {
+                    value = _value1;
+                    return true;
+                }
+                else if (ReferenceEquals(key, _key2))
+                {
+                    value = _value2;
+                    return true;
+                }
+                else if (ReferenceEquals(key, _key3))
+                {
+                    value = _value3;
+                    return true;
+                }
+                else
+                {
+                    value = null;
+                    return false;
+                }
+            }
+        }
+
+        // Instance with up to 16 key/value pairs.
+        private sealed class MultiElementAsyncLocalValueMap : IAsyncLocalValueMap
+        {
+            internal const int MaxMultiElements = 16;
+            private readonly KeyValuePair<IAsyncLocal, object>[] _keyValues;
+
+            internal MultiElementAsyncLocalValueMap(int count)
+            {
+                Debug.Assert(count <= MaxMultiElements);
+                _keyValues = new KeyValuePair<IAsyncLocal, object>[count];
+            }
+
+            internal void UnsafeStore(int index, IAsyncLocal key, object value)
+            {
+                Debug.Assert(index < _keyValues.Length);
+                _keyValues[index] = new KeyValuePair<IAsyncLocal, object>(key, value);
+            }
+
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
+            {
+                // Find the key in this map.
+                for (int i = 0; i < _keyValues.Length; i++)
+                {
+                    if (ReferenceEquals(key, _keyValues[i].Key))
+                    {
+                        // The key is in the map.  If the value isn't null, then create a new map of the same
+                        // size that has all of the same pairs, with this new key/value pair overwriting the old.
+                        if (value != null)
+                        {
+                            var multi = new MultiElementAsyncLocalValueMap(_keyValues.Length);
+                            Array.Copy(_keyValues, 0, multi._keyValues, 0, _keyValues.Length);
+                            multi._keyValues[i] = new KeyValuePair<IAsyncLocal, object>(key, value);
+                            return multi;
+                        }
+                        else if (_keyValues.Length == 4)
+                        {
+                            // The value is null, and we only have four elements, one of which we're removing,
+                            // so downgrade to a three-element map, without the matching element.
+                            return
+                                i == 0 ? new ThreeElementAsyncLocalValueMap(_keyValues[1].Key, _keyValues[1].Value, _keyValues[2].Key, _keyValues[2].Value, _keyValues[3].Key, _keyValues[3].Value) :
+                                i == 1 ? new ThreeElementAsyncLocalValueMap(_keyValues[0].Key, _keyValues[0].Value, _keyValues[2].Key, _keyValues[2].Value, _keyValues[3].Key, _keyValues[3].Value) :
+                                i == 2 ? new ThreeElementAsyncLocalValueMap(_keyValues[0].Key, _keyValues[0].Value, _keyValues[1].Key, _keyValues[1].Value, _keyValues[3].Key, _keyValues[3].Value) :
+                                (IAsyncLocalValueMap)new ThreeElementAsyncLocalValueMap(_keyValues[0].Key, _keyValues[0].Value, _keyValues[1].Key, _keyValues[1].Value, _keyValues[2].Key, _keyValues[2].Value);
+                        }
+                        else
+                        {
+                            // The value is null, and we have enough elements remaining to warrant a multi map.
+                            // Create a new one and copy all of the elements from this one, except the one to be removed.
+                            var multi = new MultiElementAsyncLocalValueMap(_keyValues.Length - 1);
+                            if (i != 0) Array.Copy(_keyValues, 0, multi._keyValues, 0, i);
+                            if (i != _keyValues.Length - 1) Array.Copy(_keyValues, i + 1, multi._keyValues, i, _keyValues.Length - i - 1);
+                            return multi;
+                        }
+                    }
+                }
+
+                // The key does not already exist in this map.
+
+                // If the value is null, then we can simply return this same map, as there's nothing to add or remove.
+                if (value == null)
+                {
+                    return this;
+                }
+
+                // We need to create a new map that has the additional key/value pair.
+                // If with the addition we can still fit in a multi map, create one.
+                if (_keyValues.Length < MaxMultiElements)
+                {
+                    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);
+                    return multi;
+                }
+
+                // Otherwise, upgrade to a many map.
+                var many = new ManyElementAsyncLocalValueMap(MaxMultiElements + 1);
+                foreach (KeyValuePair<IAsyncLocal, object> pair in _keyValues)
+                {
+                    many[pair.Key] = pair.Value;
+                }
+                many[key] = value;
+                return many;
+            }
+
+            public bool TryGetValue(IAsyncLocal key, out object value)
+            {
+                foreach (KeyValuePair<IAsyncLocal, object> pair in _keyValues)
+                {
+                    if (ReferenceEquals(key, pair.Key))
+                    {
+                        value = pair.Value;
+                        return true;
+                    }
+                }
+                value = null;
+                return false;
+            }
+        }
+
+        // Instance with any number of key/value pairs.
+        private sealed class ManyElementAsyncLocalValueMap : Dictionary<IAsyncLocal, object>, IAsyncLocalValueMap
+        {
+            public ManyElementAsyncLocalValueMap(int capacity) : base(capacity) { }
+
+            public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
+            {
+                int count = Count;
+                bool containsKey = ContainsKey(key);
+
+                // If the value being set exists, create a new many map, copy all of the elements from this one,
+                // and then store the new key/value pair into it.  This is the most common case.
+                if (value != null)
+                {
+                    var map = new ManyElementAsyncLocalValueMap(count + (containsKey ? 0 : 1));
+                    foreach (KeyValuePair<IAsyncLocal, object> pair in this)
+                    {
+                        map[pair.Key] = pair.Value;
+                    }
+                    map[key] = value;
+                    return map;
+                }
+
+                // Otherwise, the value is null, which means null is being stored into an AsyncLocal.Value.
+                // Since there's no observable difference at the API level between storing null and the key
+                // not existing at all, we can downgrade to a smaller map rather than storing null.
+
+                // If the key is contained in this map, we're going to create a new map that's one pair smaller.
+                if (containsKey)
+                {
+                    // If the new count would be within range of a multi map instead of a many map,
+                    // downgrade to the many map, which uses less memory and is faster to access.
+                    // Otherwise, just create a new many map that's missing this key.
+                    if (count == MultiElementAsyncLocalValueMap.MaxMultiElements + 1)
+                    {
+                        var multi = new MultiElementAsyncLocalValueMap(MultiElementAsyncLocalValueMap.MaxMultiElements);
+                        int index = 0;
+                        foreach (KeyValuePair<IAsyncLocal, object> pair in this)
+                        {
+                            if (!ReferenceEquals(key, pair.Key))
+                            {
+                                multi.UnsafeStore(index++, pair.Key, pair.Value);
+                            }
+                        }
+                        Debug.Assert(index == MultiElementAsyncLocalValueMap.MaxMultiElements);
+                        return multi;
+                    }
+                    else
+                    {
+                        var map = new ManyElementAsyncLocalValueMap(count - 1);
+                        foreach (KeyValuePair<IAsyncLocal, object> pair in this)
+                        {
+                            if (!ReferenceEquals(key, pair.Key))
+                            {
+                                map[pair.Key] = pair.Value;
+                            }
+                        }
+                        Debug.Assert(map.Count == count - 1);
+                        return map;
+                    }
+                }
+
+                // We were storing null, but the key wasn't in the map, so there's nothing to change.
+                // Just return this instance.
+                return this;
+            }
+        }
+    }
 }
diff --git a/src/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs b/src/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs
index e309e8955..c99bc534b 100644
--- a/src/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs
+++ b/src/System.Private.CoreLib/src/System/Threading/ExecutionContext.cs
@@ -8,11 +8,15 @@
 **
 ** Purpose: Capture execution  context for a thread
 **
-** 
+**
 ===========================================================*/
 
-using System.Collections.Generic;
 using System.Diagnostics;
+using System.Diagnostics.Contracts;
+using System.Runtime.ExceptionServices;
+using System.Runtime.Serialization;
+
+using Thread = Internal.Runtime.Augments.RuntimeThread;
 
 namespace System.Threading
 {
@@ -23,88 +27,176 @@ namespace System.Threading
         internal ExecutionContext m_ec;
         internal SynchronizationContext m_sc;
 
-        internal void Undo()
+        internal void Undo(Thread currentThread)
         {
-            SynchronizationContext.SetSynchronizationContext(m_sc);
-            ExecutionContext.Restore(m_ec);
+            Debug.Assert(currentThread == Thread.CurrentThread);
+
+            // The common case is that these have not changed, so avoid the cost of a write if not needed.
+            if (currentThread.SynchronizationContext != m_sc)
+            {
+                currentThread.SynchronizationContext = m_sc;
+            }
+
+            if (currentThread.ExecutionContext != m_ec)
+            {
+                ExecutionContext.Restore(currentThread, m_ec);
+            }
         }
     }
 
-    public sealed class ExecutionContext
+    [Serializable]
+    public sealed class ExecutionContext : IDisposable, ISerializable
     {
-        public static readonly ExecutionContext Default = new ExecutionContext();
-
-        [ThreadStatic]
-        private static ExecutionContext t_currentMaybeNull;
+        internal static readonly ExecutionContext Default = new ExecutionContext();
 
-        private readonly LowLevelDictionaryWithIEnumerable<IAsyncLocal, object> m_localValues;
-        private readonly LowLevelListWithIList<IAsyncLocal> m_localChangeNotifications;
+        private readonly IAsyncLocalValueMap m_localValues;
+        private readonly IAsyncLocal[] m_localChangeNotifications;
+        private readonly bool m_isFlowSuppressed;
 
         private ExecutionContext()
         {
-            m_localValues = new LowLevelDictionaryWithIEnumerable<IAsyncLocal, object>();
-            m_localChangeNotifications = new LowLevelListWithIList<IAsyncLocal>();
+            m_localValues = AsyncLocalValueMap.Empty;
+            m_localChangeNotifications = Array.Empty<IAsyncLocal>();
+        }
+
+        private ExecutionContext(
+            IAsyncLocalValueMap localValues,
+            IAsyncLocal[] localChangeNotifications,
+            bool isFlowSuppressed)
+        {
+            m_localValues = localValues;
+            m_localChangeNotifications = localChangeNotifications;
+            m_isFlowSuppressed = isFlowSuppressed;
         }
 
-        private ExecutionContext(ExecutionContext other)
+        public void GetObjectData(SerializationInfo info, StreamingContext context)
         {
-            m_localValues = new LowLevelDictionaryWithIEnumerable<IAsyncLocal, object>();
-            foreach (KeyValuePair<IAsyncLocal, object> kvp in other.m_localValues)
+            if (info == null)
             {
-                m_localValues.Add(kvp.Key, kvp.Value);
+                throw new ArgumentNullException(nameof(info));
             }
+            Contract.EndContractBlock();
+        }
 
-            m_localChangeNotifications = new LowLevelListWithIList<IAsyncLocal>(other.m_localChangeNotifications);
+        private ExecutionContext(SerializationInfo info, StreamingContext context)
+        {
         }
 
         public static ExecutionContext Capture()
         {
-            return t_currentMaybeNull ?? ExecutionContext.Default;
+            ExecutionContext executionContext = Thread.CurrentThread.ExecutionContext;
+            return
+                executionContext == null ? Default :
+                executionContext.m_isFlowSuppressed ? null :
+                executionContext;
+        }
+
+        private ExecutionContext ShallowClone(bool isFlowSuppressed)
+        {
+            Debug.Assert(isFlowSuppressed != m_isFlowSuppressed);
+
+            if (!isFlowSuppressed &&
+                m_localValues == Default.m_localValues &&
+                m_localChangeNotifications == Default.m_localChangeNotifications)
+            {
+                return null; // implies the default context
+            }
+            return new ExecutionContext(m_localValues, m_localChangeNotifications, isFlowSuppressed);
+        }
+
+        public static AsyncFlowControl SuppressFlow()
+        {
+            Thread currentThread = Thread.CurrentThread;
+            ExecutionContext executionContext = currentThread.ExecutionContext ?? Default;
+            if (executionContext.m_isFlowSuppressed)
+            {
+                throw new InvalidOperationException(SR.InvalidOperation_CannotSupressFlowMultipleTimes);
+            }
+            Contract.EndContractBlock();
+
+            executionContext = executionContext.ShallowClone(isFlowSuppressed: true);
+            var asyncFlowControl = new AsyncFlowControl();
+            currentThread.ExecutionContext = executionContext;
+            asyncFlowControl.Initialize(currentThread);
+            return asyncFlowControl;
+        }
+
+        public static void RestoreFlow()
+        {
+            Thread currentThread = Thread.CurrentThread;
+            ExecutionContext executionContext = currentThread.ExecutionContext;
+            if (executionContext == null || !executionContext.m_isFlowSuppressed)
+            {
+                throw new InvalidOperationException(SR.InvalidOperation_CannotRestoreUnsupressedFlow);
+            }
+            Contract.EndContractBlock();
+
+            currentThread.ExecutionContext = executionContext.ShallowClone(isFlowSuppressed: false);
+        }
+
+        public static bool IsFlowSuppressed()
+        {
+            ExecutionContext executionContext = Thread.CurrentThread.ExecutionContext;
+            return executionContext != null && executionContext.m_isFlowSuppressed;
         }
 
+        //[HandleProcessCorruptedStateExceptions]
         public static void Run(ExecutionContext executionContext, ContextCallback callback, Object state)
         {
+            if (executionContext == null)
+                throw new InvalidOperationException(SR.InvalidOperation_NullContext);
+
+            Thread currentThread = Thread.CurrentThread;
             ExecutionContextSwitcher ecsw = default(ExecutionContextSwitcher);
             try
             {
-                EstablishCopyOnWriteScope(ref ecsw);
-
-                ExecutionContext.Restore(executionContext);
+                EstablishCopyOnWriteScope(currentThread, ref ecsw);
+                ExecutionContext.Restore(currentThread, executionContext);
                 callback(state);
             }
             catch
             {
                 // Note: we have a "catch" rather than a "finally" because we want
                 // to stop the first pass of EH here.  That way we can restore the previous
-                // context before any of our callers' EH filters run.  That means we need to 
+                // context before any of our callers' EH filters run.  That means we need to
                 // end the scope separately in the non-exceptional case below.
-                ecsw.Undo();
+                ecsw.Undo(currentThread);
                 throw;
             }
-            ecsw.Undo();
+            ecsw.Undo(currentThread);
         }
 
-        internal static void Restore(ExecutionContext executionContext)
+        internal static void Restore(Thread currentThread, ExecutionContext executionContext)
         {
-            if (executionContext == null)
-                throw new InvalidOperationException(SR.InvalidOperation_NullContext);
+            Debug.Assert(currentThread == Thread.CurrentThread);
+
+            ExecutionContext previous = currentThread.ExecutionContext ?? Default;
+            currentThread.ExecutionContext = executionContext;
 
-            ExecutionContext previous = t_currentMaybeNull ?? Default;
-            t_currentMaybeNull = executionContext;
+            // New EC could be null if that's what ECS.Undo saved off.
+            // For the purposes of dealing with context change, treat this as the default EC
+            executionContext = executionContext ?? Default;
 
             if (previous != executionContext)
+            {
                 OnContextChanged(previous, executionContext);
+            }
         }
 
-        internal static void EstablishCopyOnWriteScope(ref ExecutionContextSwitcher ecsw)
+        internal static void EstablishCopyOnWriteScope(Thread currentThread, ref ExecutionContextSwitcher ecsw)
         {
-            ecsw.m_ec = Capture();
-            ecsw.m_sc = SynchronizationContext.CurrentExplicit;
+            Debug.Assert(currentThread == Thread.CurrentThread);
+
+            ecsw.m_ec = currentThread.ExecutionContext;
+            ecsw.m_sc = currentThread.SynchronizationContext;
         }
 
+        //[HandleProcessCorruptedStateExceptions]
         private static void OnContextChanged(ExecutionContext previous, ExecutionContext current)
         {
-            previous = previous ?? Default;
+            Debug.Assert(previous != null);
+            Debug.Assert(current != null);
+            Debug.Assert(previous != current);
 
             foreach (IAsyncLocal local in previous.m_localChangeNotifications)
             {
@@ -138,14 +230,16 @@ namespace System.Threading
                 }
                 catch (Exception ex)
                 {
-                    Environment.FailFast(SR.ExecutionContext_ExceptionInAsyncLocalNotification, ex);
+                    Environment.FailFast(
+                        SR.ExecutionContext_ExceptionInAsyncLocalNotification,
+                        ex);
                 }
             }
         }
 
         internal static object GetLocalValue(IAsyncLocal local)
         {
-            ExecutionContext current = t_currentMaybeNull;
+            ExecutionContext current = Thread.CurrentThread.ExecutionContext;
             if (current == null)
                 return null;
 
@@ -156,7 +250,7 @@ namespace System.Threading
 
         internal static void SetLocalValue(IAsyncLocal local, object newValue, bool needChangeNotifications)
         {
-            ExecutionContext current = t_currentMaybeNull ?? ExecutionContext.Default;
+            ExecutionContext current = Thread.CurrentThread.ExecutionContext ?? ExecutionContext.Default;
 
             object previousValue;
             bool hadPreviousValue = current.m_localValues.TryGetValue(local, out previousValue);
@@ -164,41 +258,115 @@ namespace System.Threading
             if (previousValue == newValue)
                 return;
 
-            current = new ExecutionContext(current);
-            current.m_localValues[local] = newValue;
-
-            t_currentMaybeNull = current;
+            IAsyncLocalValueMap newValues = current.m_localValues.Set(local, newValue);
 
+            //
+            // Either copy the change notification array, or create a new one, depending on whether we need to add a new item.
+            //
+            IAsyncLocal[] newChangeNotifications = current.m_localChangeNotifications;
             if (needChangeNotifications)
             {
                 if (hadPreviousValue)
-                    Debug.Assert(current.m_localChangeNotifications.Contains(local));
+                {
+                    Debug.Assert(Array.IndexOf(newChangeNotifications, local) >= 0);
+                }
                 else
-                    current.m_localChangeNotifications.Add(local);
+                {
+                    int newNotificationIndex = newChangeNotifications.Length;
+                    Array.Resize(ref newChangeNotifications, newNotificationIndex + 1);
+                    newChangeNotifications[newNotificationIndex] = local;
+                }
+            }
+
+            Thread.CurrentThread.ExecutionContext =
+                new ExecutionContext(newValues, newChangeNotifications, current.m_isFlowSuppressed);
 
+            if (needChangeNotifications)
+            {
                 local.OnValueChanged(previousValue, newValue, false);
             }
         }
 
-        [Flags]
-        internal enum CaptureOptions
+        public ExecutionContext CreateCopy()
         {
-            None = 0x00,
-            IgnoreSyncCtx = 0x01,
-            OptimizeDefaultCase = 0x02,
+            return this; // since CoreCLR's ExecutionContext is immutable, we don't need to create copies.
         }
 
-        internal static ExecutionContext PreAllocatedDefault
+        public void Dispose()
         {
-            get
-            { return ExecutionContext.Default; }
+            // For CLR compat only
         }
+    }
 
-        internal bool IsPreAllocatedDefault
+    public struct AsyncFlowControl : IDisposable
+    {
+        private Thread _thread;
+
+        internal void Initialize(Thread currentThread)
         {
-            get { return this == ExecutionContext.Default; }
+            Debug.Assert(currentThread == Thread.CurrentThread);
+            _thread = currentThread;
+        }
+
+        public void Undo()
+        {
+            if (_thread == null)
+            {
+                throw new InvalidOperationException(SR.InvalidOperation_CannotUseAFCMultiple);
+            }
+            if (Thread.CurrentThread != _thread)
+            {
+                throw new InvalidOperationException(SR.InvalidOperation_CannotUseAFCOtherThread);
+            }
+
+            // An async flow control cannot be undone when a different execution context is applied. The desktop framework
+            // mutates the execution context when its state changes, and only changes the instance when an execution context
+            // is applied (for instance, through ExecutionContext.Run). The framework prevents a suppressed-flow execution
+            // context from being applied by returning null from ExecutionContext.Capture, so the only type of execution
+            // context that can be applied is one whose flow is not suppressed. After suppressing flow and changing an async
+            // local's value, the desktop framework verifies that a different execution context has not been applied by
+            // checking the execution context instance against the one saved from when flow was suppressed. In .NET Core,
+            // since the execution context instance will change after changing the async local's value, it verifies that a
+            // different execution context has not been applied, by instead ensuring that the current execution context's
+            // flow is suppressed.
+            if (!ExecutionContext.IsFlowSuppressed())
+            {
+                throw new InvalidOperationException(SR.InvalidOperation_AsyncFlowCtrlCtxMismatch);
+            }
+            Contract.EndContractBlock();
+
+            _thread = null;
+            ExecutionContext.RestoreFlow();
+        }
+
+        public void Dispose()
+        {
+            Undo();
+        }
+
+        public override bool Equals(object obj)
+        {
+            return obj is AsyncFlowControl && Equals((AsyncFlowControl)obj);
+        }
+
+        public bool Equals(AsyncFlowControl obj)
+        {
+            return _thread == obj._thread;
         }
-    }
-}
 
+        public override int GetHashCode()
+        {
+            return _thread?.GetHashCode() ?? 0;
+        }
 
+        public static bool operator ==(AsyncFlowControl a, AsyncFlowControl b)
+        {
+            return a.Equals(b);
+        }
+
+        public static bool operator !=(AsyncFlowControl a, AsyncFlowControl b)
+        {
+            return !(a == b);
+        }
+    }
+}
diff --git a/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.Dummy.cs b/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.Dummy.cs
index 1f1efed8d..968f13875 100644
--- a/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.Dummy.cs
+++ b/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.Dummy.cs
@@ -14,7 +14,7 @@ namespace System.Threading
         {
             get
             {
-                return s_current;
+                return CurrentExplicit;
             }
         }
     }
diff --git a/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.WinRT.cs b/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.WinRT.cs
index bb5c84851..5972a6e00 100644
--- a/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.WinRT.cs
+++ b/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.WinRT.cs
@@ -2,11 +2,10 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using Internal.Runtime.Augments;
 using System.Diagnostics.Contracts;
 using System.Runtime.CompilerServices;
 
-using Internal.Runtime.Augments;
-
 namespace System.Threading
 {
     //
@@ -19,7 +18,7 @@ namespace System.Threading
         {
             get
             {
-                return s_current ?? GetWinRTContext();
+                return CurrentExplicit ?? GetWinRTContext();
             }
         }
 
diff --git a/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.cs b/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.cs
index 843381362..5fbfb4a65 100644
--- a/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.cs
+++ b/src/System.Private.CoreLib/src/System/Threading/SynchronizationContext.cs
@@ -11,6 +11,8 @@
 ** 
 ===========================================================*/
 
+using Internal.Runtime.Augments;
+
 namespace System.Threading
 {
     public partial class SynchronizationContext
@@ -43,20 +45,17 @@ namespace System.Threading
         {
         }
 
-        [ThreadStatic]
-        private static SynchronizationContext s_current;
-
         // Set the SynchronizationContext on the current thread
         public static void SetSynchronizationContext(SynchronizationContext syncContext)
         {
-            s_current = syncContext;
+            RuntimeThread.CurrentThread.SynchronizationContext = syncContext;
         }
 
         internal static SynchronizationContext CurrentExplicit
         {
             get
             {
-                return s_current;
+                return RuntimeThread.CurrentThread.SynchronizationContext;
             }
         }
 
diff --git a/src/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs b/src/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs
index b55f670c8..baa9e4f4c 100644
--- a/src/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs
+++ b/src/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs
@@ -1701,11 +1701,11 @@ namespace System.Threading.Tasks
             {
                 var props = m_contingentProperties;
                 if (props != null && props.m_capturedContext != null) return props.m_capturedContext;
-                else return ExecutionContext.PreAllocatedDefault;
+                else return ExecutionContext.Default;
             }
             set
             {
-                if (!value.IsPreAllocatedDefault) // not the default context, then inflate the contingent properties and set it
+                if (value != ExecutionContext.Default) // not the default context, then inflate the contingent properties and set it
                 {
                     EnsureContingentPropertiesInitialized(needsProtection: false).m_capturedContext = value;
                 }