[runtime-async] Fix timestamp instrumentation

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

@@ -514,13 +514,13 @@ internal unsafe bool HandleSuspended(ref RuntimeAsyncAwaitState state)
                 return false;
             }
 
-            internal void InstrumentedHandleSuspended(AsyncInstrumentation.Flags flags, ref RuntimeAsyncAwaitState state, Continuation? newContinuation = null)
+            internal void InstrumentedHandleSuspended(AsyncInstrumentation.Flags flags, ref RuntimeAsyncAwaitState state)
             {
                 if (AsyncInstrumentation.IsEnabled.AsyncDebugger(flags))
                 {
                     Continuation? nextContinuation = state.SentinelContinuation!.Next;
 
-                    AsyncDebugger.HandleSuspended(nextContinuation, newContinuation);
+                    AsyncDebugger.HandleSuspended(nextContinuation);
 
                     if (!HandleSuspended(ref state))
                     {
@@ -696,7 +696,7 @@ private unsafe void InstrumentedDispatchContinuations(AsyncInstrumentation.Flags
                             newContinuation.Next = nextContinuation;
 
                             RuntimeAsyncInstrumentationHelpers.AwaitSuspendedRuntimeAsyncContext(ref asyncDispatcherInfo, flags, curContinuation, newContinuation, awaitState.SentinelContinuation!.Next);
-                            InstrumentedHandleSuspended(flags, ref awaitState, newContinuation);
+                            InstrumentedHandleSuspended(flags, ref awaitState);
 
                             awaitState.Pop();
                             refDispatcherInfo = asyncDispatcherInfo.Next;
@@ -754,7 +754,7 @@ private unsafe void InstrumentedDispatchContinuations(AsyncInstrumentation.Flags
 
                     if (QueueContinuationFollowUpActionIfNecessary(asyncDispatcherInfo.NextContinuation))
                     {
-                        RuntimeAsyncInstrumentationHelpers.QueueSuspendedRuntimeAsyncContext(ref asyncDispatcherInfo, flags, curContinuation, asyncDispatcherInfo.NextContinuation);
+                        RuntimeAsyncInstrumentationHelpers.QueueSuspendedRuntimeAsyncContext(ref asyncDispatcherInfo, flags, asyncDispatcherInfo.NextContinuation);
 
                         awaitState.Pop();
                         refDispatcherInfo = asyncDispatcherInfo.Next;
@@ -1211,7 +1211,7 @@ public static void ResumeRuntimeAsyncContext(Task task, ref AsyncDispatcherInfo
             }
 
             [MethodImpl(MethodImplOptions.AggressiveInlining)]
-            public static void QueueSuspendedRuntimeAsyncContext(ref AsyncDispatcherInfo info, AsyncInstrumentation.Flags flags, Continuation curContinuation, Continuation nextContinuation)
+            public static void QueueSuspendedRuntimeAsyncContext(ref AsyncDispatcherInfo info, AsyncInstrumentation.Flags flags, Continuation nextContinuation)
             {
                 if (AsyncInstrumentation.IsEnabled.SuspendAsyncContext(flags))
                 {
@@ -1222,7 +1222,7 @@ public static void QueueSuspendedRuntimeAsyncContext(ref AsyncDispatcherInfo inf
 
                     if (AsyncInstrumentation.IsEnabled.AsyncDebugger(flags))
                     {
-                        AsyncDebugger.SuspendAsyncContext(ref info, curContinuation);
+                        AsyncDebugger.SuspendAsyncContext();
                     }
                 }
             }
@@ -1370,19 +1370,14 @@ public static void ResumeAsyncContext(Task task)
                 TplEventSource.Log.TraceSynchronousWorkBegin(task.Id, CausalitySynchronousWork.Execution);
             }
 
-            public static void SuspendAsyncContext(ref AsyncDispatcherInfo info, Continuation curContinuation)
+            public static void SuspendAsyncContext(Continuation curContinuation, Continuation newContinuation)
             {
-                if (info.NextContinuation != null)
-                {
-                    Task.TryAddRuntimeAsyncContinuationChainTimestamps(info.NextContinuation, curContinuation);
-                }
-
+                Task.ReplaceOrAddRuntimeAsyncContinuationTimestamp(curContinuation, newContinuation);
                 TplEventSource.Log.TraceSynchronousWorkEnd(CausalitySynchronousWork.Execution);
             }
 
-            public static void SuspendAsyncContext(Continuation curContinuation, Continuation newContinuation)
+            public static void SuspendAsyncContext()
             {
-                Task.ReplaceOrAddRuntimeAsyncContinuationTimestamp(curContinuation, newContinuation);
                 TplEventSource.Log.TraceSynchronousWorkEnd(CausalitySynchronousWork.Execution);
             }
 
@@ -1424,18 +1419,11 @@ public static void CompleteAsyncMethod(Continuation curContinuation)
                 Task.RemoveRuntimeAsyncContinuationTimestamp(curContinuation);
             }
 
-            public static void HandleSuspended(Continuation? nextContinuation, Continuation? newContinuation)
+            public static void HandleSuspended(Continuation? nextContinuation)
             {
                 if (nextContinuation != null)
                 {
-                    if (newContinuation != null)
-                    {
-                        Task.TryAddRuntimeAsyncContinuationChainTimestamps(nextContinuation, newContinuation);
-                    }
-                    else
-                    {
-                        Task.TryAddRuntimeAsyncContinuationChainTimestamps(nextContinuation);
-                    }
+                    Task.TryAddRuntimeAsyncContinuationChainTimestamps(nextContinuation);
                 }
             }
 

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

@@ -264,21 +264,6 @@ internal static void TryAddRuntimeAsyncContinuationChainTimestamps(Continuation
             }
         }
 
-        internal static void TryAddRuntimeAsyncContinuationChainTimestamps(Continuation continuationChain, Continuation timestampSource)
-        {
-            var continuationTimestamps = GetOrCreateRuntimeAsyncContinuationTimestamps();
-            lock (continuationTimestamps)
-            {
-                long timestamp = continuationTimestamps.TryGetValue(timestampSource, out long timestampVal) ? timestampVal : Stopwatch.GetTimestamp();
-                Continuation? nc = continuationChain;
-                while (nc != null)
-                {
-                    continuationTimestamps.TryAdd(nc, timestamp);
-                    nc = nc.Next;
-                }
-            }
-        }
-
         internal static void RemoveRuntimeAsyncContinuationTimestamp(Continuation continuation)
         {
             var continuationTimestamps = s_runtimeAsyncContinuationTimestamps;

src/libraries/System.Runtime/tests/System.Threading.Tasks.Tests/System.Runtime.CompilerServices/RuntimeAsyncTests.cs

@@ -4,6 +4,7 @@
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics.Tracing;
+using System.Linq;
 using System.Reflection;
 using Microsoft.DotNet.RemoteExecutor;
 using Microsoft.DotNet.XUnitExtensions;
@@ -177,6 +178,48 @@ static async Task FuncThatInspectsContinuationTimestamps(TaskCompletionSource tc
             callback();
         }
 
+        [System.Runtime.CompilerServices.RuntimeAsyncMethodGeneration(true)]
+        static async Task FuncWithNestedDelays(List<long> observedTimestamps)
+        {
+            var continuationTimestamps = (Dictionary<object, long>)s_continuationTimestampsField.GetValue(null);
+
+            await Task.Delay(50);
+            lock (continuationTimestamps)
+            {
+                foreach (long ts in continuationTimestamps.Values)
+                    observedTimestamps.Add(ts);
+            }
+
+            await NestedDelay1(observedTimestamps);
+            await NestedDelay2(observedTimestamps);
+        }
+
+        [System.Runtime.CompilerServices.RuntimeAsyncMethodGeneration(true)]
+        static async Task NestedDelay1(List<long> observedTimestamps)
+        {
+            var continuationTimestamps = (Dictionary<object, long>)s_continuationTimestampsField.GetValue(null);
+
+            await Task.Delay(50);
+            lock (continuationTimestamps)
+            {
+                foreach (long ts in continuationTimestamps.Values)
+                    observedTimestamps.Add(ts);
+            }
+        }
+
+        [System.Runtime.CompilerServices.RuntimeAsyncMethodGeneration(true)]
+        static async Task NestedDelay2(List<long> observedTimestamps)
+        {
+            var continuationTimestamps = (Dictionary<object, long>)s_continuationTimestampsField.GetValue(null);
+
+            await Task.Delay(50);
+            lock (continuationTimestamps)
+            {
+                foreach (long ts in continuationTimestamps.Values)
+                    observedTimestamps.Add(ts);
+            }
+        }
+
         static void ValidateTimestampsCleared()
         {
             // some other tasks may be created by the runtime, so this is just using a reasonably small upper bound
@@ -532,6 +575,27 @@ public void RuntimeAsync_TplEvents()
             }).Dispose();
         }
 
+        [ConditionalFact(typeof(RuntimeAsyncTests), nameof(IsRemoteExecutorAndRuntimeAsyncSupported))]
+        public void RuntimeAsync_SuspensionTimestampsAreDistinct()
+        {
+            RemoteExecutor.Invoke(async () =>
+            {
+                AttachDebugger();
+
+                var observedTimestamps = new List<long>();
+                await FuncWithNestedDelays(observedTimestamps);
+
+                // Each suspension across nested async methods should produce a fresh, non-zero
+                // timestamp — not one inherited from a parent continuation.
+                Assert.True(observedTimestamps.Count >= 3, $"Expected at least 3 observed timestamps, got {observedTimestamps.Count}");
+                Assert.All(observedTimestamps, ts => Assert.True(ts > 0, "Expected non-zero timestamp"));
+                Assert.True(observedTimestamps.Distinct().Count() > 1, "Expected timestamps from different suspensions to not all be identical");
+
+                DetachDebugger();
+
+            }).Dispose();
+        }
+
         [ConditionalFact(typeof(RuntimeAsyncTests), nameof(IsRemoteExecutorAndRuntimeAsyncSupported))]
         public void RuntimeAsync_NoTplEventsWithoutDebugger()
         {