From 9891705d34239ece0bfffc2e5ba4913da519a493 Mon Sep 17 00:00:00 2001
From: Shay Rojansky <roji@roji.org>
Date: Wed, 8 Dec 2021 16:54:33 +0100
Subject: [PATCH] Preserve synchronization context in ExecutionStrategy

Fixes #26763
---
 src/EFCore/Storage/ExecutionStrategy.cs       |  6 +--
 .../SingleThreadSynchronizationContext.cs     | 43 +++++++++++++++++++
 .../Query/QueryBugsTest.cs                    | 11 +++--
 .../Storage/ExecutionStrategyTest.cs          | 30 +++++++++++++
 4 files changed, 83 insertions(+), 7 deletions(-)
 create mode 100644 test/EFCore.Specification.Tests/TestUtilities/SingleThreadSynchronizationContext.cs

diff --git a/src/EFCore/Storage/ExecutionStrategy.cs b/src/EFCore/Storage/ExecutionStrategy.cs
index 1259cc81599..37fb349fedc 100644
--- a/src/EFCore/Storage/ExecutionStrategy.cs
+++ b/src/EFCore/Storage/ExecutionStrategy.cs
@@ -332,7 +332,7 @@ private async Task<ExecutionResult<TResult>> ExecuteImplementationAsync<TState,
 
                 Current = this;
                 var result = await operation(Dependencies.CurrentContext.Context, state, cancellationToken)
-                    .ConfigureAwait(false);
+                    .ConfigureAwait(true);
                 Current = null;
                 return result;
             }
@@ -346,7 +346,7 @@ private async Task<ExecutionResult<TResult>> ExecuteImplementationAsync<TState,
                     && CallOnWrappedException(ex, ShouldVerifySuccessOn))
                 {
                     var result = await ExecuteImplementationAsync(verifySucceeded, null, state, cancellationToken)
-                        .ConfigureAwait(false);
+                        .ConfigureAwait(true);
                     if (result.IsSuccessful)
                     {
                         return result;
@@ -370,7 +370,7 @@ private async Task<ExecutionResult<TResult>> ExecuteImplementationAsync<TState,
 
                 OnRetry();
 
-                await Task.Delay(delay.Value, cancellationToken).ConfigureAwait(false);
+                await Task.Delay(delay.Value, cancellationToken).ConfigureAwait(true);
             }
         }
     }
diff --git a/test/EFCore.Specification.Tests/TestUtilities/SingleThreadSynchronizationContext.cs b/test/EFCore.Specification.Tests/TestUtilities/SingleThreadSynchronizationContext.cs
new file mode 100644
index 00000000000..aa9d0b3a4de
--- /dev/null
+++ b/test/EFCore.Specification.Tests/TestUtilities/SingleThreadSynchronizationContext.cs
@@ -0,0 +1,43 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+
+using System.Collections.Concurrent;
+
+namespace Microsoft.EntityFrameworkCore.TestUtilities;
+
+public class SingleThreadSynchronizationContext : SynchronizationContext, IDisposable
+{
+    private readonly BlockingCollection<(SendOrPostCallback callback, object state)> _tasks = new();
+
+    public Thread Thread { get; }
+
+    public SingleThreadSynchronizationContext()
+    {
+        Thread = new Thread(WorkLoop);
+        Thread.Start();
+    }
+
+    public override void Post(SendOrPostCallback callback, object state)
+        => _tasks.Add((callback, state));
+
+    public void Dispose()
+        => _tasks.CompleteAdding();
+
+    private void WorkLoop()
+    {
+        SetSynchronizationContext(this);
+
+        try
+        {
+            while (true)
+            {
+                var (callback, state) = _tasks.Take();
+                callback(state);
+            }
+        }
+        catch (InvalidOperationException)
+        {
+            _tasks.Dispose();
+        }
+    }
+}
diff --git a/test/EFCore.SqlServer.FunctionalTests/Query/QueryBugsTest.cs b/test/EFCore.SqlServer.FunctionalTests/Query/QueryBugsTest.cs
index f31637bcae6..6b482e252e7 100644
--- a/test/EFCore.SqlServer.FunctionalTests/Query/QueryBugsTest.cs
+++ b/test/EFCore.SqlServer.FunctionalTests/Query/QueryBugsTest.cs
@@ -9041,7 +9041,7 @@ public SqlExpression Translate(
 
     #region Issue22841
 
-    [ConditionalFact(Skip = "Flaky, #26763")]
+    [ConditionalFact]
     public async Task SaveChangesAsync_accepts_changes_with_ConfigureAwait_true_22841()
     {
         var contextFactory = await InitializeAsync<MyContext22841>();
@@ -9049,12 +9049,14 @@ public async Task SaveChangesAsync_accepts_changes_with_ConfigureAwait_true_2284
         using var context = contextFactory.CreateContext();
         var observableThing = new ObservableThing22841();
 
+        using var trackingSynchronizationContext = new SingleThreadSynchronizationContext();
         var origSynchronizationContext = SynchronizationContext.Current;
-        var trackingSynchronizationContext = new SingleThreadSynchronizationContext22841();
         SynchronizationContext.SetSynchronizationContext(trackingSynchronizationContext);
 
         bool? isMySyncContext = null;
-        Action callback = () => isMySyncContext = Thread.CurrentThread == trackingSynchronizationContext.Thread;
+        Action callback = () => isMySyncContext =
+            SynchronizationContext.Current ==  trackingSynchronizationContext
+            && Thread.CurrentThread == trackingSynchronizationContext.Thread;
         observableThing.Event += callback;
 
         try
@@ -9066,7 +9068,6 @@ public async Task SaveChangesAsync_accepts_changes_with_ConfigureAwait_true_2284
         {
             observableThing.Event -= callback;
             SynchronizationContext.SetSynchronizationContext(origSynchronizationContext);
-            trackingSynchronizationContext.Dispose();
         }
 
         Assert.True(isMySyncContext);
@@ -9126,6 +9127,8 @@ public void Dispose()
 
         private void WorkLoop()
         {
+            SetSynchronizationContext(this);
+
             try
             {
                 while (true)
diff --git a/test/EFCore.Tests/Storage/ExecutionStrategyTest.cs b/test/EFCore.Tests/Storage/ExecutionStrategyTest.cs
index 5e1747f82f2..ebbe06156e6 100644
--- a/test/EFCore.Tests/Storage/ExecutionStrategyTest.cs
+++ b/test/EFCore.Tests/Storage/ExecutionStrategyTest.cs
@@ -1,6 +1,7 @@
 // Licensed to the .NET Foundation under one or more agreements.
 // The .NET Foundation licenses this file to you under the MIT license.
 
+using System.Collections.Concurrent;
 using System.Transactions;
 
 // ReSharper disable AccessToModifiedClosure
@@ -629,6 +630,35 @@ public void ShouldRetryOn_does_not_get_null_on_DbUpdateConcurrencyException()
         Assert.Equal(2, executionCount);
     }
 
+    [ConditionalFact]
+    public async Task ExecuteAsync_preserves_synchronization_context_across_retries()
+    {
+        var mockExecutionStrategy = new TestExecutionStrategy(Context, shouldRetryOn: e => e is DbUpdateConcurrencyException);
+
+        var origSyncContext = SynchronizationContext.Current;
+        using var syncContext = new SingleThreadSynchronizationContext();
+        SynchronizationContext.SetSynchronizationContext(syncContext);
+
+        try
+        {
+            var executionCount = 0;
+
+            await mockExecutionStrategy.ExecuteAsync(async _ =>
+            {
+                Assert.Same(syncContext, SynchronizationContext.Current);
+                await Task.Yield();
+                if (executionCount++ < 1)
+                {
+                    throw new DbUpdateConcurrencyException("");
+                }
+            }, cancellationToken: default);
+        }
+        finally
+        {
+            SynchronizationContext.SetSynchronizationContext(origSyncContext);
+        }
+    }
+
     protected DbContext CreateContext()
         => InMemoryTestHelpers.Instance.CreateContext(
             InMemoryTestHelpers.Instance.CreateServiceProvider(