Preserve synchronization context in ExecutionStrategy

Fixes dotnet#26763
roji · Dec 9, 2021 · 9467b4d · 9467b4d
1 parent 9b2d0dc
commit 9467b4d
Show file tree

Hide file tree

Showing 4 changed files with 81 additions and 43 deletions.
diff --git 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
@@ -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
@@ -9041,20 +9041,22 @@ 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>();
 
         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);
@@ -9105,42 +9106,6 @@ public int Id
         public event Action Event;
     }
 
-    private class SingleThreadSynchronizationContext22841 : SynchronizationContext, IDisposable
-    {
-        private readonly CancellationTokenSource _cancellationTokenSource;
-        private readonly BlockingCollection<(SendOrPostCallback callback, object state)> _tasks = new();
-        internal Thread Thread { get; }
-
-        internal SingleThreadSynchronizationContext22841()
-        {
-            _cancellationTokenSource = new CancellationTokenSource();
-            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()
-        {
-            try
-            {
-                while (true)
-                {
-                    var (callback, state) = _tasks.Take();
-                    callback(state);
-                }
-            }
-            catch (InvalidOperationException)
-            {
-                _tasks.Dispose();
-            }
-        }
-    }
-
     #endregion Issue22841
 
     #region Issue12482

diff --git 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(