Fix race condition in WaitForInstanceAsync causing RestartAsync_EndToEnd test to fail

The original code had a TOCTOU race condition where a completion notification
could be missed if the orchestration completed between checking completion
status and adding the waiter. The fix reorders the operations to add the
waiter first, then check for completion.

Co-authored-by: YunchuWang <[email protected]>

Author: Copilot

src/InProcessTestHost/Sidecar/InMemoryOrchestrationService.cs

@@ -739,6 +739,12 @@ public void ReleaseLock(string instanceId)
 
         public Task<OrchestrationState> WaitForInstanceAsync(string instanceId, CancellationToken cancellationToken)
         {
+            // First, add the waiter before checking completion to avoid a race condition.
+            // This ensures we don't miss a completion notification that happens between
+            // checking the status and adding the waiter.
+            var tcs = this.waiters.GetOrAdd(instanceId, _ => new TaskCompletionSource<OrchestrationState>());
+
+            // Now check if already completed - if so, complete the waiter immediately
             if (this.store.TryGetValue(instanceId, out SerializedInstanceState? state))
             {
                 lock (state)
@@ -750,16 +756,15 @@ public Task<OrchestrationState> WaitForInstanceAsync(string instanceId, Cancella
                             statusRecord.OrchestrationStatus == OrchestrationStatus.Failed ||
                             statusRecord.OrchestrationStatus == OrchestrationStatus.Terminated)
                         {
-                            // orchestration has already completed
-                            return Task.FromResult(statusRecord);
+                            // Orchestration has already completed - complete the waiter and return
+                            tcs.TrySetResult(statusRecord);
                         }
                     }
                 }
             }
 
             // Caller will be notified when the instance completes.
             // The ContinueWith is just to enable cancellation: https://stackoverflow.com/a/25652873/2069
-            var tcs = this.waiters.GetOrAdd(instanceId, _ => new TaskCompletionSource<OrchestrationState>());
             return tcs.Task.ContinueWith(t => t.GetAwaiter().GetResult(), cancellationToken);
         }