Preserve synchronization context in ExecutionStrategy (#26953)

Fixes #26763

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);
  }
  }
  }

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();
+ }
+ }
+}

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

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(