Managed fixes for async task tracking

src/coreclr/System.Private.CoreLib/src/System/Runtime/CompilerServices/AsyncHelpers.CoreCLR.cs

@@ -161,6 +161,15 @@ internal unsafe ref struct AsyncDispatcherInfo
  [FieldOffset(4)]
 #endif
  public Continuation? NextContinuation;
+#if TARGET_64BIT
+ [FieldOffset(16)]
+#else
+ [FieldOffset(8)]
+#endif
+ // The runtime async Task being dispatched.
+ // This is used by debuggers in the case of nested dispatcher info (multiple runtime-async Tasks on the same thread)
+ // to match an inflight Task to the corresponding Continuation chain.
+ public Task? CurrentTask;
 
  // Information about current task dispatching, to be used for async
  // stackwalking.
@@ -379,6 +388,18 @@ internal void HandleSuspended()
 
  SetContinuationState(headContinuation);
 
+ Continuation? nc = headContinuation;
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ long timestamp = Stopwatch.GetTimestamp();
+ while (nc != null)
+ {
+ // On suspension we set timestamp for all continuations that have not yet had it set.
+ Task.SetRuntimeAsyncContinuationTimestamp(nc, timestamp);
+ nc = nc.Next;
+ }
+ }
+
  try
  {
  if (critNotifier != null)
@@ -440,10 +461,29 @@ internal void HandleSuspended()
  }
  catch (Exception ex)
  {
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ Task.RemoveFromActiveTasks(this);
+ Task.RemoveRuntimeAsyncTaskTimestamp(this);
+ Continuation? nextCont = headContinuation;
+ while (nextCont != null)
+ {
+ Task.RemoveRuntimeAsyncContinuationTimestamp(nextCont);
+ nextCont = nextCont.Next;
+ }
+ }
+
  Task.ThrowAsync(ex, targetContext: null);
  }
  }
 
+#pragma warning disable CA1822 // Mark members as static
+ [MethodImpl(MethodImplOptions.NoOptimization)]
+ public void NotifyDebuggerOfRuntimeAsyncState()
+ {
+ }
+#pragma warning restore CA1822
+
  [StackTraceHidden]
  private unsafe void DispatchContinuations()
  {
@@ -453,31 +493,59 @@ private unsafe void DispatchContinuations()
  AsyncDispatcherInfo asyncDispatcherInfo;
  asyncDispatcherInfo.Next = AsyncDispatcherInfo.t_current;
  asyncDispatcherInfo.NextContinuation = MoveContinuationState();
+ asyncDispatcherInfo.CurrentTask = this;
  AsyncDispatcherInfo.t_current = &asyncDispatcherInfo;
+ bool isTplEnabled = TplEventSource.Log.IsEnabled();
+
+ if (isTplEnabled)
+ {
+ TplEventSource.Log.TraceSynchronousWorkBegin(this.Id, CausalitySynchronousWork.Execution);
+ }
 
  while (true)
  {
  Debug.Assert(asyncDispatcherInfo.NextContinuation != null);
+ Continuation curContinuation = asyncDispatcherInfo.NextContinuation;
  try
  {
- Continuation curContinuation = asyncDispatcherInfo.NextContinuation;
  Continuation? nextContinuation = curContinuation.Next;
  asyncDispatcherInfo.NextContinuation = nextContinuation;
 
  ref byte resultLoc = ref nextContinuation != null ? ref nextContinuation.GetResultStorageOrNull() : ref GetResultStorage();
+ long timestamp = 0;
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ timestamp = Task.GetRuntimeAsyncContinuationTimestamp(curContinuation, out long timestampVal) ? timestampVal : Stopwatch.GetTimestamp();
+ // we have dequeued curContinuation; update task tick info so that we can track its start time from a debugger
+ Task.UpdateRuntimeAsyncTaskTimestamp(this, timestamp);
+ }
  Continuation? newContinuation = curContinuation.ResumeInfo->Resume(curContinuation, ref resultLoc);
 
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ Task.RemoveRuntimeAsyncContinuationTimestamp(curContinuation);
+ }
+
  if (newContinuation != null)
  {
+ // we have a new Continuation that belongs to the same logical invocation as the previous; propagate debug info from previous continuation
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ Task.SetRuntimeAsyncContinuationTimestamp(newContinuation, timestamp);
+ }
  newContinuation.Next = nextContinuation;
  HandleSuspended();
  contexts.Pop();
  AsyncDispatcherInfo.t_current = asyncDispatcherInfo.Next;
-return;
+break;
  }
  }
  catch (Exception ex)
  {
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ Task.RemoveRuntimeAsyncContinuationTimestamp(curContinuation);
+ }
  Continuation? handlerContinuation = UnwindToPossibleHandler(asyncDispatcherInfo.NextContinuation, ex);
  if (handlerContinuation == null)
  {
@@ -490,12 +558,23 @@ private unsafe void DispatchContinuations()
 
  AsyncDispatcherInfo.t_current = asyncDispatcherInfo.Next;
 
+ if (isTplEnabled)
+ {
+ TplEventSource.Log.TraceOperationEnd(this.Id, ex is OperationCanceledException ? AsyncCausalityStatus.Canceled : AsyncCausalityStatus.Error);
+ }
+
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ Task.RemoveFromActiveTasks(this);
+ Task.RemoveRuntimeAsyncTaskTimestamp(this);
+ }
+
  if (!successfullySet)
  {
  ThrowHelper.ThrowInvalidOperationException(ExceptionResource.TaskT_TransitionToFinal_AlreadyCompleted);
  }
 
-return;
+break;
  }
 
  handlerContinuation.SetException(ex);
@@ -504,6 +583,15 @@ private unsafe void DispatchContinuations()
 
  if (asyncDispatcherInfo.NextContinuation == null)
  {
+ if (isTplEnabled)
+ {
+ TplEventSource.Log.TraceOperationEnd(this.Id, AsyncCausalityStatus.Completed);
+ }
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ Task.RemoveFromActiveTasks(this);
+ Task.RemoveRuntimeAsyncTaskTimestamp(this);
+ }
  bool successfullySet = TrySetResult(m_result);
 
  contexts.Pop();
@@ -515,16 +603,20 @@ private unsafe void DispatchContinuations()
  ThrowHelper.ThrowInvalidOperationException(ExceptionResource.TaskT_TransitionToFinal_AlreadyCompleted);
  }
 
-return;
+break;
  }
 
  if (QueueContinuationFollowUpActionIfNecessary(asyncDispatcherInfo.NextContinuation))
  {
  contexts.Pop();
  AsyncDispatcherInfo.t_current = asyncDispatcherInfo.Next;
-return;
+break;
  }
  }
+ if (isTplEnabled)
+ {
+ TplEventSource.Log.TraceSynchronousWorkEnd(CausalitySynchronousWork.Execution);
+ }
  }
 
  private ref byte GetResultStorage() => ref Unsafe.As<T?, byte>(ref m_result);
@@ -544,7 +636,10 @@ private unsafe void DispatchContinuations()
  }
  if (continuation == null || (continuation.Flags & ContinuationFlags.HasException) != 0)
  return continuation;
-
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ Task.RemoveRuntimeAsyncContinuationTimestamp(continuation);
+ }
  continuation = continuation.Next;
  }
  }
@@ -632,13 +727,31 @@ private bool QueueContinuationFollowUpActionIfNecessary(Continuation continuatio
  private static Task<T?> FinalizeTaskReturningThunk<T>()
  {
  RuntimeAsyncTask<T?> result = new();
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ result.NotifyDebuggerOfRuntimeAsyncState();
+ Task.AddToActiveTasks(result);
+ }
+ if (TplEventSource.Log.IsEnabled())
+ {
+ TplEventSource.Log.TraceOperationBegin(result.Id, "System.Runtime.CompilerServices.AsyncHelpers+RuntimeAsyncTask", 0);
+ }
  result.HandleSuspended();
  return result;
  }
 
  private static Task FinalizeTaskReturningThunk()
  {
  RuntimeAsyncTask<VoidTaskResult> result = new();
+ if (Task.s_asyncDebuggingEnabled)
+ {
+ result.NotifyDebuggerOfRuntimeAsyncState();
+ Task.AddToActiveTasks(result);
+ }
+ if (TplEventSource.Log.IsEnabled())
+ {
+ TplEventSource.Log.TraceOperationBegin(result.Id, "System.Runtime.CompilerServices.AsyncHelpers+RuntimeAsyncTask", 0);
+ }
  result.HandleSuspended();
  return result;
  }

src/libraries/System.Private.CoreLib/src/System/Threading/Tasks/Task.cs

@@ -123,7 +123,6 @@ public class Task : IAsyncResult, IDisposable
  // The delegate to invoke for a delegate-backed Task.
  // This field also may be used by async state machines to cache an Action.
  internal Delegate? m_action;
-
  private protected object? m_stateObject; // A state object that can be optionally supplied, passed to action.
  internal TaskScheduler? m_taskScheduler; // The task scheduler this task runs under.
 
@@ -178,6 +177,13 @@ internal enum TaskStateFlags
  // task. This is to be used by the debugger ONLY. Task in this dictionary represent current active tasks.
  private static Dictionary<int, Task>? s_currentActiveTasks;
 
+#if !MONO
+ // Dictionary that relates a runtime-async Task's ID to the timestamp when the current inflight invocation started.
+ // Needed because Continuations that are inflight have already been dequeued from the chain.
+ private static Dictionary<int, long>? s_runtimeAsyncTaskTimestamps;
+ // Dictionary to store the timestamp when the logical invocation to which the Continuation belongs started.
+ private static Dictionary<object, long>? s_runtimeAsyncContinuationTimestamps;
+#endif
  // These methods are a way to access the dictionary both from this class and for other classes that also
  // activate dummy tasks. Specifically the AsyncTaskMethodBuilder and AsyncTaskMethodBuilder<>
  internal static bool AddToActiveTasks(Task task)
@@ -211,6 +217,73 @@ internal static void RemoveFromActiveTasks(Task task)
  }
  }
 
+#if !MONO
+ internal static void SetRuntimeAsyncContinuationTimestamp(Continuation continuation, long timestamp)
+ {
+ Dictionary<object, long> continuationTimestamps =
+ Volatile.Read(ref s_runtimeAsyncContinuationTimestamps) ??
+ Interlocked.CompareExchange(ref s_runtimeAsyncContinuationTimestamps, new Dictionary<object, long>(ReferenceEqualityComparer.Instance), null) ??
+ s_runtimeAsyncContinuationTimestamps;
+
+ lock (continuationTimestamps)
+ {
+ continuationTimestamps.TryAdd(continuation, timestamp);
+ }
+ }
+
+ internal static bool GetRuntimeAsyncContinuationTimestamp(Continuation continuation, out long timestamp)
+ {
+ Dictionary<object, long>? continuationTimestamps = s_runtimeAsyncContinuationTimestamps;
+ if (continuationTimestamps is null)
+ {
+ timestamp = 0;
+ return false;
+ }
+
+ lock (continuationTimestamps)
+ {
+ return continuationTimestamps.TryGetValue(continuation, out timestamp);
+ }
+ }
+
+ internal static void RemoveRuntimeAsyncContinuationTimestamp(Continuation continuation)
+ {
+ Dictionary<object, long>? continuationTimestamps = s_runtimeAsyncContinuationTimestamps;
+ if (continuationTimestamps is null)
+ return;
+
+ lock (continuationTimestamps)
+ {
+ continuationTimestamps.Remove(continuation);
+ }
+ }
+
+ internal static void UpdateRuntimeAsyncTaskTimestamp(Task task, long inflightTimestamp)
+ {
+ Dictionary<int, long> runtimeAsyncTaskTimestamps =
+ Volatile.Read(ref s_runtimeAsyncTaskTimestamps) ??
+ Interlocked.CompareExchange(ref s_runtimeAsyncTaskTimestamps, new Dictionary<int, long>(), null) ??
+ s_runtimeAsyncTaskTimestamps;
+
+ lock (runtimeAsyncTaskTimestamps)
+ {
+ runtimeAsyncTaskTimestamps[task.Id] = inflightTimestamp;
+ }
+ }
+
+ internal static void RemoveRuntimeAsyncTaskTimestamp(Task task)
+ {
+ Dictionary<int, long>? runtimeAsyncTaskTimestamps = s_runtimeAsyncTaskTimestamps;
+ if (runtimeAsyncTaskTimestamps is null)
+ return;
+
+ lock (runtimeAsyncTaskTimestamps)
+ {
+ runtimeAsyncTaskTimestamps.Remove(task.Id);
+ }
+ }
+#endif
+
  // We moved a number of Task properties into this class. The idea is that in most cases, these properties never
  // need to be accessed during the life cycle of a Task, so we don't want to instantiate them every time. Once
  // one of these properties needs to be written, we will instantiate a ContingentProperties object and set
@@ -296,6 +369,7 @@ internal void UnregisterCancellationCallback()
  // Constructs the task as already completed
  internal Task(bool canceled, TaskCreationOptions creationOptions, CancellationToken ct)
  {
+ _ = s_asyncDebuggingEnabled; // Debugger depends on Task cctor being run when any Task gets created
  int optionFlags = (int)creationOptions;
  if (canceled)
  {
@@ -315,6 +389,7 @@ internal Task(bool canceled, TaskCreationOptions creationOptions, CancellationTo
  /// <summary>Constructor for use with promise-style tasks that aren't configurable.</summary>
  internal Task()
  {
+ _ = s_asyncDebuggingEnabled; // Debugger depends on Task cctor being run when any Task gets created
  m_stateFlags = (int)TaskStateFlags.WaitingForActivation | (int)InternalTaskOptions.PromiseTask;
  }
 
@@ -323,6 +398,7 @@ internal Task()
  // (action,TCO). It should always be true.
  internal Task(object? state, TaskCreationOptions creationOptions, bool promiseStyle)
  {
+ _ = s_asyncDebuggingEnabled; // Debugger depends on Task cctor being run when any Task gets created
  Debug.Assert(promiseStyle, "Promise CTOR: promiseStyle was false");
 
  // Check the creationOptions. We allow the AttachedToParent option to be specified for promise tasks.
@@ -503,6 +579,7 @@ public Task(Action<object?> action, object? state, CancellationToken cancellatio
  internal Task(Delegate action, object? state, Task? parent, CancellationToken cancellationToken,
  TaskCreationOptions creationOptions, InternalTaskOptions internalOptions, TaskScheduler? scheduler)
  {
+ _ = s_asyncDebuggingEnabled; // Debugger depends on Task cctor being run when any Task gets created
  if (action == null)
  {
  ThrowHelper.ThrowArgumentNullException(ExceptionArgument.action);