1 // Licensed to the .NET Foundation under one or more agreements.
2 // The .NET Foundation licenses this file to you under the MIT license.
3 // See the LICENSE file in the project root for more information.
5 using System.Collections.Generic;
6 using System.Diagnostics;
8 namespace System.Threading
11 // AsyncLocal<T> represents "ambient" data that is local to a given asynchronous control flow, such as an
12 // async method. For example, say you want to associate a culture with a given async flow:
14 // static AsyncLocal<Culture> s_currentCulture = new AsyncLocal<Culture>();
16 // static async Task SomeOperationAsync(Culture culture)
18 // s_currentCulture.Value = culture;
23 // static async Task FooAsync()
25 // PrintStringWithCulture(s_currentCulture.Value);
28 // AsyncLocal<T> also provides optional notifications when the value associated with the current thread
29 // changes, either because it was explicitly changed by setting the Value property, or implicitly changed
30 // when the thread encountered an "await" or other context transition. For example, we might want our
31 // current culture to be communicated to the OS as well:
33 // static AsyncLocal<Culture> s_currentCulture = new AsyncLocal<Culture>(
36 // NativeMethods.SetThreadCulture(args.CurrentValue.LCID);
39 public sealed class AsyncLocal<T> : IAsyncLocal
41 private readonly Action<AsyncLocalValueChangedArgs<T>> m_valueChangedHandler;
44 // Constructs an AsyncLocal<T> that does not receive change notifications.
51 // Constructs an AsyncLocal<T> with a delegate that is called whenever the current value changes
54 public AsyncLocal(Action<AsyncLocalValueChangedArgs<T>> valueChangedHandler)
56 m_valueChangedHandler = valueChangedHandler;
63 object obj = ExecutionContext.GetLocalValue(this);
64 return (obj == null) ? default(T) : (T)obj;
68 ExecutionContext.SetLocalValue(this, value, m_valueChangedHandler != null);
72 void IAsyncLocal.OnValueChanged(object previousValueObj, object currentValueObj, bool contextChanged)
74 Debug.Assert(m_valueChangedHandler != null);
75 T previousValue = previousValueObj == null ? default(T) : (T)previousValueObj;
76 T currentValue = currentValueObj == null ? default(T) : (T)currentValueObj;
77 m_valueChangedHandler(new AsyncLocalValueChangedArgs<T>(previousValue, currentValue, contextChanged));
82 // Interface to allow non-generic code in ExecutionContext to call into the generic AsyncLocal<T> type.
84 internal interface IAsyncLocal
86 void OnValueChanged(object previousValue, object currentValue, bool contextChanged);
89 public struct AsyncLocalValueChangedArgs<T>
91 public T PreviousValue { get; private set; }
92 public T CurrentValue { get; private set; }
95 // If the value changed because we changed to a different ExecutionContext, this is true. If it changed
96 // because someone set the Value property, this is false.
98 public bool ThreadContextChanged { get; private set; }
100 internal AsyncLocalValueChangedArgs(T previousValue, T currentValue, bool contextChanged)
103 PreviousValue = previousValue;
104 CurrentValue = currentValue;
105 ThreadContextChanged = contextChanged;
110 // Interface used to store an IAsyncLocal => object mapping in ExecutionContext.
111 // Implementations are specialized based on the number of elements in the immutable
112 // map in order to minimize memory consumption and look-up times.
114 internal interface IAsyncLocalValueMap
116 bool TryGetValue(IAsyncLocal key, out object value);
117 IAsyncLocalValueMap Set(IAsyncLocal key, object value);
121 // Utility functions for getting/creating instances of IAsyncLocalValueMap
123 internal static class AsyncLocalValueMap
125 public static IAsyncLocalValueMap Empty { get; } = new EmptyAsyncLocalValueMap();
127 // Instance without any key/value pairs. Used as a singleton/
128 private sealed class EmptyAsyncLocalValueMap : IAsyncLocalValueMap
130 public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
132 // If the value isn't null, then create a new one-element map to store
133 // the key/value pair. If it is null, then we're still empty.
134 return value != null ?
135 new OneElementAsyncLocalValueMap(key, value) :
136 (IAsyncLocalValueMap)this;
139 public bool TryGetValue(IAsyncLocal key, out object value)
146 // Instance with one key/value pair.
147 private sealed class OneElementAsyncLocalValueMap : IAsyncLocalValueMap
149 private readonly IAsyncLocal _key1;
150 private readonly object _value1;
152 public OneElementAsyncLocalValueMap(IAsyncLocal key, object value)
154 _key1 = key; _value1 = value;
157 public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
161 // The value is non-null. If the key matches one already contained in this map,
162 // then create a new one-element map with the updated value, otherwise create
163 // a two-element map with the additional key/value.
164 return ReferenceEquals(key, _key1) ?
165 new OneElementAsyncLocalValueMap(key, value) :
166 (IAsyncLocalValueMap)new TwoElementAsyncLocalValueMap(_key1, _value1, key, value);
170 // The value is null. If the key exists in this map, remove it by downgrading to an empty map.
171 // Otherwise, there's nothing to add or remove, so just return this map.
172 return ReferenceEquals(key, _key1) ?
174 (IAsyncLocalValueMap)this;
178 public bool TryGetValue(IAsyncLocal key, out object value)
180 if (ReferenceEquals(key, _key1))
193 // Instance with two key/value pairs.
194 private sealed class TwoElementAsyncLocalValueMap : IAsyncLocalValueMap
196 private readonly IAsyncLocal _key1, _key2;
197 private readonly object _value1, _value2;
199 public TwoElementAsyncLocalValueMap(IAsyncLocal key1, object value1, IAsyncLocal key2, object value2)
201 _key1 = key1; _value1 = value1;
202 _key2 = key2; _value2 = value2;
205 public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
209 // The value is non-null. If the key matches one already contained in this map,
210 // then create a new two-element map with the updated value, otherwise create
211 // a three-element map with the additional key/value.
213 ReferenceEquals(key, _key1) ? new TwoElementAsyncLocalValueMap(key, value, _key2, _value2) :
214 ReferenceEquals(key, _key2) ? new TwoElementAsyncLocalValueMap(_key1, _value1, key, value) :
215 (IAsyncLocalValueMap)new ThreeElementAsyncLocalValueMap(_key1, _value1, _key2, _value2, key, value);
219 // The value is null. If the key exists in this map, remove it by downgrading to a one-element map
220 // without the key. Otherwise, there's nothing to add or remove, so just return this map.
222 ReferenceEquals(key, _key1) ? new OneElementAsyncLocalValueMap(_key2, _value2) :
223 ReferenceEquals(key, _key2) ? new OneElementAsyncLocalValueMap(_key1, _value1) :
224 (IAsyncLocalValueMap)this;
228 public bool TryGetValue(IAsyncLocal key, out object value)
230 if (ReferenceEquals(key, _key1))
235 else if (ReferenceEquals(key, _key2))
248 // Instance with three key/value pairs.
249 private sealed class ThreeElementAsyncLocalValueMap : IAsyncLocalValueMap
251 private readonly IAsyncLocal _key1, _key2, _key3;
252 private readonly object _value1, _value2, _value3;
254 public ThreeElementAsyncLocalValueMap(IAsyncLocal key1, object value1, IAsyncLocal key2, object value2, IAsyncLocal key3, object value3)
256 _key1 = key1; _value1 = value1;
257 _key2 = key2; _value2 = value2;
258 _key3 = key3; _value3 = value3;
261 public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
265 // The value is non-null. If the key matches one already contained in this map,
266 // then create a new three-element map with the updated value.
267 if (ReferenceEquals(key, _key1)) return new ThreeElementAsyncLocalValueMap(key, value, _key2, _value2, _key3, _value3);
268 if (ReferenceEquals(key, _key2)) return new ThreeElementAsyncLocalValueMap(_key1, _value1, key, value, _key3, _value3);
269 if (ReferenceEquals(key, _key3)) return new ThreeElementAsyncLocalValueMap(_key1, _value1, _key2, _value2, key, value);
271 // The key doesn't exist in this map, so upgrade to a multi map that contains
272 // the additional key/value pair.
273 var multi = new MultiElementAsyncLocalValueMap(4);
274 multi.UnsafeStore(0, _key1, _value1);
275 multi.UnsafeStore(1, _key2, _value2);
276 multi.UnsafeStore(2, _key3, _value3);
277 multi.UnsafeStore(3, key, value);
282 // The value is null. If the key exists in this map, remove it by downgrading to a two-element map
283 // without the key. Otherwise, there's nothing to add or remove, so just return this map.
285 ReferenceEquals(key, _key1) ? new TwoElementAsyncLocalValueMap(_key2, _value2, _key3, _value3) :
286 ReferenceEquals(key, _key2) ? new TwoElementAsyncLocalValueMap(_key1, _value1, _key3, _value3) :
287 ReferenceEquals(key, _key3) ? new TwoElementAsyncLocalValueMap(_key1, _value1, _key2, _value2) :
288 (IAsyncLocalValueMap)this;
292 public bool TryGetValue(IAsyncLocal key, out object value)
294 if (ReferenceEquals(key, _key1))
299 else if (ReferenceEquals(key, _key2))
304 else if (ReferenceEquals(key, _key3))
317 // Instance with up to 16 key/value pairs.
318 private sealed class MultiElementAsyncLocalValueMap : IAsyncLocalValueMap
320 internal const int MaxMultiElements = 16;
321 private readonly KeyValuePair<IAsyncLocal, object>[] _keyValues;
323 internal MultiElementAsyncLocalValueMap(int count)
325 Debug.Assert(count <= MaxMultiElements);
326 _keyValues = new KeyValuePair<IAsyncLocal, object>[count];
329 internal void UnsafeStore(int index, IAsyncLocal key, object value)
331 Debug.Assert(index < _keyValues.Length);
332 _keyValues[index] = new KeyValuePair<IAsyncLocal, object>(key, value);
335 public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
337 // Find the key in this map.
338 for (int i = 0; i < _keyValues.Length; i++)
340 if (ReferenceEquals(key, _keyValues[i].Key))
342 // The key is in the map. If the value isn't null, then create a new map of the same
343 // size that has all of the same pairs, with this new key/value pair overwriting the old.
346 var multi = new MultiElementAsyncLocalValueMap(_keyValues.Length);
347 Array.Copy(_keyValues, 0, multi._keyValues, 0, _keyValues.Length);
348 multi._keyValues[i] = new KeyValuePair<IAsyncLocal, object>(key, value);
351 else if (_keyValues.Length == 4)
353 // The value is null, and we only have four elements, one of which we're removing,
354 // so downgrade to a three-element map, without the matching element.
356 i == 0 ? new ThreeElementAsyncLocalValueMap(_keyValues[1].Key, _keyValues[1].Value, _keyValues[2].Key, _keyValues[2].Value, _keyValues[3].Key, _keyValues[3].Value) :
357 i == 1 ? new ThreeElementAsyncLocalValueMap(_keyValues[0].Key, _keyValues[0].Value, _keyValues[2].Key, _keyValues[2].Value, _keyValues[3].Key, _keyValues[3].Value) :
358 i == 2 ? new ThreeElementAsyncLocalValueMap(_keyValues[0].Key, _keyValues[0].Value, _keyValues[1].Key, _keyValues[1].Value, _keyValues[3].Key, _keyValues[3].Value) :
359 (IAsyncLocalValueMap)new ThreeElementAsyncLocalValueMap(_keyValues[0].Key, _keyValues[0].Value, _keyValues[1].Key, _keyValues[1].Value, _keyValues[2].Key, _keyValues[2].Value);
363 // The value is null, and we have enough elements remaining to warrant a multi map.
364 // Create a new one and copy all of the elements from this one, except the one to be removed.
365 var multi = new MultiElementAsyncLocalValueMap(_keyValues.Length - 1);
366 if (i != 0) Array.Copy(_keyValues, 0, multi._keyValues, 0, i);
367 if (i != _keyValues.Length - 1) Array.Copy(_keyValues, i + 1, multi._keyValues, i, _keyValues.Length - i - 1);
373 // The key does not already exist in this map.
375 // If the value is null, then we can simply return this same map, as there's nothing to add or remove.
381 // We need to create a new map that has the additional key/value pair.
382 // If with the addition we can still fit in a multi map, create one.
383 if (_keyValues.Length < MaxMultiElements)
385 var multi = new MultiElementAsyncLocalValueMap(_keyValues.Length + 1);
386 Array.Copy(_keyValues, 0, multi._keyValues, 0, _keyValues.Length);
387 multi._keyValues[_keyValues.Length] = new KeyValuePair<IAsyncLocal, object>(key, value);
391 // Otherwise, upgrade to a many map.
392 var many = new ManyElementAsyncLocalValueMap(MaxMultiElements + 1);
393 foreach (KeyValuePair<IAsyncLocal, object> pair in _keyValues)
395 many[pair.Key] = pair.Value;
401 public bool TryGetValue(IAsyncLocal key, out object value)
403 foreach (KeyValuePair<IAsyncLocal, object> pair in _keyValues)
405 if (ReferenceEquals(key, pair.Key))
416 // Instance with any number of key/value pairs.
417 private sealed class ManyElementAsyncLocalValueMap : Dictionary<IAsyncLocal, object>, IAsyncLocalValueMap
419 public ManyElementAsyncLocalValueMap(int capacity) : base(capacity) { }
421 public IAsyncLocalValueMap Set(IAsyncLocal key, object value)
424 bool containsKey = ContainsKey(key);
426 // If the value being set exists, create a new many map, copy all of the elements from this one,
427 // and then store the new key/value pair into it. This is the most common case.
430 var map = new ManyElementAsyncLocalValueMap(count + (containsKey ? 0 : 1));
431 foreach (KeyValuePair<IAsyncLocal, object> pair in this)
433 map[pair.Key] = pair.Value;
439 // Otherwise, the value is null, which means null is being stored into an AsyncLocal.Value.
440 // Since there's no observable difference at the API level between storing null and the key
441 // not existing at all, we can downgrade to a smaller map rather than storing null.
443 // If the key is contained in this map, we're going to create a new map that's one pair smaller.
446 // If the new count would be within range of a multi map instead of a many map,
447 // downgrade to the many map, which uses less memory and is faster to access.
448 // Otherwise, just create a new many map that's missing this key.
449 if (count == MultiElementAsyncLocalValueMap.MaxMultiElements + 1)
451 var multi = new MultiElementAsyncLocalValueMap(MultiElementAsyncLocalValueMap.MaxMultiElements);
453 foreach (KeyValuePair<IAsyncLocal, object> pair in this)
455 if (!ReferenceEquals(key, pair.Key))
457 multi.UnsafeStore(index++, pair.Key, pair.Value);
460 Debug.Assert(index == MultiElementAsyncLocalValueMap.MaxMultiElements);
465 var map = new ManyElementAsyncLocalValueMap(count - 1);
466 foreach (KeyValuePair<IAsyncLocal, object> pair in this)
468 if (!ReferenceEquals(key, pair.Key))
470 map[pair.Key] = pair.Value;
473 Debug.Assert(map.Count == count - 1);
478 // We were storing null, but the key wasn't in the map, so there's nothing to change.
479 // Just return this instance.