Distinguish cancellation from failure (#26988)

Closes #19526
dotnet · Dec 14, 2021 · 08e5ebd · 08e5ebd
1 parent 4a2e4eb
commit 08e5ebd
Show file tree

Hide file tree

Showing 54 changed files with 1,280 additions and 118 deletions.
diff --git a/...e.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs b/...e.Cosmos/Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs
@@ -112,6 +112,7 @@ private sealed class Enumerator : IEnumerator<T>
             private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
             private readonly bool _standAloneStateManager;
             private readonly IConcurrencyDetector _concurrencyDetector;
+            private readonly IExceptionDetector _exceptionDetector;
 
             private IEnumerator<JObject> _enumerator;
 
@@ -125,6 +126,7 @@ public Enumerator(QueryingEnumerable<T> queryingEnumerable)
                 _partitionKey = queryingEnumerable._partitionKey;
                 _queryLogger = queryingEnumerable._queryLogger;
                 _standAloneStateManager = queryingEnumerable._standAloneStateManager;
+                _exceptionDetector = _cosmosQueryContext.ExceptionDetector;
 
                 _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled
                     ? _cosmosQueryContext.ConcurrencyDetector
@@ -175,7 +177,14 @@ public bool MoveNext()
                 }
                 catch (Exception exception)
                 {
-                    _queryLogger.QueryIterationFailed(_contextType, exception);
+                    if (_exceptionDetector.IsCancellation(exception))
+                    {
+                        _queryLogger.QueryCanceled(_contextType);
+                    }
+                    else
+                    {
+                        _queryLogger.QueryIterationFailed(_contextType, exception);
+                    }
 
                     throw;
                 }
@@ -203,6 +212,7 @@ private sealed class AsyncEnumerator : IAsyncEnumerator<T>
             private readonly bool _standAloneStateManager;
             private readonly CancellationToken _cancellationToken;
             private readonly IConcurrencyDetector _concurrencyDetector;
+            private readonly IExceptionDetector _exceptionDetector;
 
             private IAsyncEnumerator<JObject> _enumerator;
 
@@ -216,6 +226,7 @@ public AsyncEnumerator(QueryingEnumerable<T> queryingEnumerable, CancellationTok
                 _partitionKey = queryingEnumerable._partitionKey;
                 _queryLogger = queryingEnumerable._queryLogger;
                 _standAloneStateManager = queryingEnumerable._standAloneStateManager;
+                _exceptionDetector = _cosmosQueryContext.ExceptionDetector;
                 _cancellationToken = cancellationToken;
 
                 _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled
@@ -264,7 +275,14 @@ public async ValueTask<bool> MoveNextAsync()
                 }
                 catch (Exception exception)
                 {
-                    _queryLogger.QueryIterationFailed(_contextType, exception);
+                    if (_exceptionDetector.IsCancellation(exception, _cancellationToken))
+                    {
+                        _queryLogger.QueryCanceled(_contextType);
+                    }
+                    else
+                    {
+                        _queryLogger.QueryIterationFailed(_contextType, exception);
+                    }
 
                     throw;
                 }

diff --git a/.../Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs b/.../Query/Internal/CosmosShapedQueryCompilingExpressionVisitor.ReadItemQueryingEnumerable.cs
@@ -169,6 +169,7 @@ private sealed class Enumerator : IEnumerator<T>, IAsyncEnumerator<T>
             private readonly IDiagnosticsLogger<DbLoggerCategory.Query> _queryLogger;
             private readonly bool _standAloneStateManager;
             private readonly IConcurrencyDetector _concurrencyDetector;
+            private readonly IExceptionDetector _exceptionDetector;
             private readonly ReadItemQueryingEnumerable<T> _readItemEnumerable;
             private readonly CancellationToken _cancellationToken;
 
@@ -183,6 +184,7 @@ public Enumerator(ReadItemQueryingEnumerable<T> readItemEnumerable, Cancellation
                 _contextType = readItemEnumerable._contextType;
                 _queryLogger = readItemEnumerable._queryLogger;
                 _standAloneStateManager = readItemEnumerable._standAloneStateManager;
+                _exceptionDetector = _cosmosQueryContext.ExceptionDetector;
                 _readItemEnumerable = readItemEnumerable;
                 _cancellationToken = cancellationToken;
 
@@ -235,7 +237,14 @@ public bool MoveNext()
                 }
                 catch (Exception exception)
                 {
-                    _queryLogger.QueryIterationFailed(_contextType, exception);
+                    if (_exceptionDetector.IsCancellation(exception))
+                    {
+                        _queryLogger.QueryCanceled(_contextType);
+                    }
+                    else
+                    {
+                        _queryLogger.QueryIterationFailed(_contextType, exception);
+                    }
 
                     throw;
                 }
@@ -282,7 +291,14 @@ public async ValueTask<bool> MoveNextAsync()
                 }
                 catch (Exception exception)
                 {
-                    _queryLogger.QueryIterationFailed(_contextType, exception);
+                    if (_exceptionDetector.IsCancellation(exception, _cancellationToken))
+                    {
+                        _queryLogger.QueryCanceled(_contextType);
+                    }
+                    else
+                    {
+                        _queryLogger.QueryIterationFailed(_contextType, exception);
+                    }
 
                     throw;
                 }

diff --git a/...Memory/Query/Internal/InMemoryShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs b/...Memory/Query/Internal/InMemoryShapedQueryCompilingExpressionVisitor.QueryingEnumerable.cs
@@ -63,6 +63,7 @@ private sealed class Enumerator : IEnumerator<T>, IAsyncEnumerator<T>
             private readonly bool _standAloneStateManager;
             private readonly CancellationToken _cancellationToken;
             private readonly IConcurrencyDetector? _concurrencyDetector;
+            private readonly IExceptionDetector _exceptionDetector;
 
             private IEnumerator<ValueBuffer>? _enumerator;
 
@@ -75,6 +76,7 @@ public Enumerator(QueryingEnumerable<T> queryingEnumerable, CancellationToken ca
                 _queryLogger = queryingEnumerable._queryLogger;
                 _standAloneStateManager = queryingEnumerable._standAloneStateManager;
                 _cancellationToken = cancellationToken;
+                _exceptionDetector = _queryContext.ExceptionDetector;
                 Current = default!;
 
                 _concurrencyDetector = queryingEnumerable._threadSafetyChecksEnabled
@@ -104,7 +106,14 @@ public bool MoveNext()
                 }
                 catch (Exception exception)
                 {
-                    _queryLogger.QueryIterationFailed(_contextType, exception);
+                    if (_exceptionDetector.IsCancellation(exception))
+                    {
+                        _queryLogger.QueryCanceled(_contextType);
+                    }
+                    else
+                    {
+                        _queryLogger.QueryIterationFailed(_contextType, exception);
+                    }
 
                     throw;
                 }
@@ -129,7 +138,14 @@ public ValueTask<bool> MoveNextAsync()
                 }
                 catch (Exception exception)
                 {
-                    _queryLogger.QueryIterationFailed(_contextType, exception);
+                    if (_exceptionDetector.IsCancellation(exception, _cancellationToken))
+                    {
+                        _queryLogger.QueryCanceled(_contextType);
+                    }
+                    else
+                    {
+                        _queryLogger.QueryIterationFailed(_contextType, exception);
+                    }
 
                     throw;
                 }

diff --git a/src/EFCore.InMemory/Storage/Internal/InMemoryDatabase.cs b/src/EFCore.InMemory/Storage/Internal/InMemoryDatabase.cs
@@ -64,7 +64,9 @@ public override int SaveChanges(IList<IUpdateEntry> entries)
     public override Task<int> SaveChangesAsync(
         IList<IUpdateEntry> entries,
         CancellationToken cancellationToken = default)
-        => Task.FromResult(_store.ExecuteTransaction(entries, _updateLogger));
+        => cancellationToken.IsCancellationRequested
+            ? Task.FromCanceled<int>(cancellationToken)
+            : Task.FromResult(_store.ExecuteTransaction(entries, _updateLogger));
 
     /// <summary>
     ///     This is an internal API that supports the Entity Framework Core infrastructure and not subject to

diff --git a/src/EFCore.Relational/Diagnostics/DbCommandInterceptor.cs b/src/EFCore.Relational/Diagnostics/DbCommandInterceptor.cs
@@ -361,6 +361,31 @@ public virtual ValueTask<int> NonQueryExecutedAsync(
         CancellationToken cancellationToken = default)
         => new(result);
 
+    /// <summary>
+    ///     Called when a command was canceled.
+    /// </summary>
+    /// <param name="command">The command.</param>
+    /// <param name="eventData">Contextual information about the command and execution.</param>
+    public virtual void CommandCanceled(
+        DbCommand command,
+        CommandEndEventData eventData)
+    {
+    }
+
+    /// <summary>
+    ///     Called when a command was canceled.
+    /// </summary>
+    /// <param name="command">The command.</param>
+    /// <param name="eventData">Contextual information about the command and execution.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken" /> to observe while waiting for the task to complete.</param>
+    /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
+    /// <exception cref="OperationCanceledException">If the <see cref="CancellationToken" /> is canceled.</exception>
+    public virtual Task CommandCanceledAsync(
+        DbCommand command,
+        CommandEndEventData eventData,
+        CancellationToken cancellationToken = default)
+        => Task.CompletedTask;
+
     /// <summary>
     ///     Called when execution of a command has failed with an exception.
     /// </summary>

diff --git a/src/EFCore.Relational/Diagnostics/IDbCommandInterceptor.cs b/src/EFCore.Relational/Diagnostics/IDbCommandInterceptor.cs
@@ -365,6 +365,28 @@ ValueTask<int> NonQueryExecutedAsync(
         int result,
         CancellationToken cancellationToken = default);
 
+    /// <summary>
+    ///     Called when a command was canceled.
+    /// </summary>
+    /// <param name="command">The command.</param>
+    /// <param name="eventData">Contextual information about the command and execution.</param>
+    void CommandCanceled(
+        DbCommand command,
+        CommandEndEventData eventData);
+
+    /// <summary>
+    ///     Called when a command was canceled.
+    /// </summary>
+    /// <param name="command">The command.</param>
+    /// <param name="eventData">Contextual information about the command and execution.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken" /> to observe while waiting for the task to complete.</param>
+    /// <returns>A <see cref="Task" /> representing the asynchronous operation.</returns>
+    /// <exception cref="OperationCanceledException">If the <see cref="CancellationToken" /> is canceled.</exception>
+    Task CommandCanceledAsync(
+        DbCommand command,
+        CommandEndEventData eventData,
+        CancellationToken cancellationToken = default);
+
     /// <summary>
     ///     Called when execution of a command has failed with an exception.
     /// </summary>

diff --git a/src/EFCore.Relational/Diagnostics/IRelationalCommandDiagnosticsLogger.cs b/src/EFCore.Relational/Diagnostics/IRelationalCommandDiagnosticsLogger.cs
@@ -400,6 +400,56 @@ Task CommandErrorAsync(
         CommandSource commandSource,
         CancellationToken cancellationToken = default);
 
+    /// <summary>
+    ///     Logs for the <see cref="RelationalEventId.CommandCanceled" /> event.
+    /// </summary>
+    /// <param name="connection">The connection.</param>
+    /// <param name="command">The database command object.</param>
+    /// <param name="context">The <see cref="DbContext" /> currently being used, to null if not known.</param>
+    /// <param name="executeMethod">Represents the method that will be called to execute the command.</param>
+    /// <param name="commandId">The correlation ID associated with the given <see cref="DbCommand" />.</param>
+    /// <param name="connectionId">The correlation ID associated with the <see cref="DbConnection" /> being used.</param>
+    /// <param name="startTime">The time that execution began.</param>
+    /// <param name="duration">The amount of time that passed until the exception was raised.</param>
+    /// <param name="commandSource">Source of the command.</param>
+    void CommandCanceled(
+        IRelationalConnection connection,
+        DbCommand command,
+        DbContext? context,
+        DbCommandMethod executeMethod,
+        Guid commandId,
+        Guid connectionId,
+        DateTimeOffset startTime,
+        TimeSpan duration,
+        CommandSource commandSource);
+
+    /// <summary>
+    ///     Logs for the <see cref="RelationalEventId.CommandCanceled" /> event.
+    /// </summary>
+    /// <param name="connection">The connection.</param>
+    /// <param name="command">The database command object.</param>
+    /// <param name="context">The <see cref="DbContext" /> currently being used, to null if not known.</param>
+    /// <param name="executeMethod">Represents the method that will be called to execute the command.</param>
+    /// <param name="commandId">The correlation ID associated with the given <see cref="DbCommand" />.</param>
+    /// <param name="connectionId">The correlation ID associated with the <see cref="DbConnection" /> being used.</param>
+    /// <param name="startTime">The time that execution began.</param>
+    /// <param name="duration">The amount of time that passed until the exception was raised.</param>
+    /// <param name="commandSource">Source of the command.</param>
+    /// <param name="cancellationToken">A <see cref="CancellationToken" /> to observe while waiting for the task to complete.</param>
+    /// <returns>A <see cref="Task" /> representing the async operation.</returns>
+    /// <exception cref="OperationCanceledException">If the <see cref="CancellationToken" /> is canceled.</exception>
+    Task CommandCanceledAsync(
+        IRelationalConnection connection,
+        DbCommand command,
+        DbContext? context,
+        DbCommandMethod executeMethod,
+        Guid commandId,
+        Guid connectionId,
+        DateTimeOffset startTime,
+        TimeSpan duration,
+        CommandSource commandSource,
+        CancellationToken cancellationToken = default);
+
     /// <summary>
     ///     Logs for the <see cref="RelationalEventId.DataReaderDisposing" /> event.
     /// </summary>

diff --git a/src/EFCore.Relational/Diagnostics/Internal/DbCommandInterceptorAggregator.cs b/src/EFCore.Relational/Diagnostics/Internal/DbCommandInterceptorAggregator.cs
@@ -221,6 +221,26 @@ public async ValueTask<int> NonQueryExecutedAsync(
             return result;
         }
 
+        public void CommandCanceled(DbCommand command, CommandEndEventData eventData)
+        {
+            for (var i = 0; i < _interceptors.Length; i++)
+            {
+                _interceptors[i].CommandCanceled(command, eventData);
+            }
+        }
+
+        public async Task CommandCanceledAsync(
+            DbCommand command,
+            CommandEndEventData eventData,
+            CancellationToken cancellationToken = default)
+        {
+            for (var i = 0; i < _interceptors.Length; i++)
+            {
+                await _interceptors[i].CommandCanceledAsync(command, eventData, cancellationToken)
+                    .ConfigureAwait(false);
+            }
+        }
+
         public void CommandFailed(DbCommand command, CommandErrorEventData eventData)
         {
             for (var i = 0; i < _interceptors.Length; i++)