CurtHagenlocher commented on code in PR #4895: URL: https://github.com/apache/arrow-adbc/pull/4895#discussion_r4278969727
########## csharp/src/Apache.Arrow.Adbc/AdbcAsync.cs: ########## @@ -0,0 +1,144 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +using System; +using System.Runtime.ExceptionServices; +using System.Threading; +using System.Threading.Tasks; + +namespace Apache.Arrow.Adbc +{ + /// <summary> + /// Helpers for the default implementations of the asynchronous methods on + /// the ADBC base classes. Those defaults run the synchronous method on the + /// thread pool and translate a <see cref="CancellationToken"/> into a call + /// to the object's <c>Cancel</c> method. + /// </summary> + internal static class AdbcAsync + { + private static readonly Action<object?> s_cancel = TryCancel; + + public static Task<T> Run<T>(Func<T> operation, Action? cancel, CancellationToken cancellationToken) + { + if (cancellationToken.IsCancellationRequested) + { + return Task.FromCanceled<T>(cancellationToken); + } + + return Task.Run(() => + { + CancelState state = new CancelState(cancel); + Exception failure; + using (Register(state, cancellationToken)) + { + try + { + return operation(); + } + catch (Exception e) when (!(e is OperationCanceledException)) + { + failure = e; + } + } + + // Disposing the registration waits for an in-progress Cancel to finish + throw Translate(failure, state, cancellationToken); + }, cancellationToken); + } + + public static Task Run(Action operation, Action? cancel, CancellationToken cancellationToken) + { + return Run<bool>(() => { operation(); return true; }, cancel, cancellationToken); + } + + /// <summary> + /// Attaches cancellation to a task produced by a driver override that + /// doesn't itself accept a <see cref="CancellationToken"/>. + /// </summary> + public static async Task<T> WithCancellation<T>(Task<T> task, Action? cancel, CancellationToken cancellationToken) + { + CancelState state = new CancelState(cancel); + Exception failure; + using (Register(state, cancellationToken)) + { + try + { + return await task.ConfigureAwait(false); + } + catch (Exception e) when (!(e is OperationCanceledException)) + { + failure = e; + } + } + + throw Translate(failure, state, cancellationToken); + } + + private static CancellationTokenRegistration Register(CancelState state, CancellationToken cancellationToken) + { + if (state.Cancel == null || !cancellationToken.CanBeCanceled) + { + return default; + } + return cancellationToken.Register(s_cancel, state); + } + + /// <summary> + /// Reports a failure as cancellation only if cancellation plausibly caused it: the + /// token was cancelled and either the driver accepted the call to Cancel or the + /// failure itself says the operation was cancelled. Otherwise, a driver that can't + /// cancel would have its own errors hidden behind <see cref="OperationCanceledException"/>. + /// </summary> + private static Exception Translate(Exception failure, CancelState state, CancellationToken cancellationToken) + { + bool cancelled = cancellationToken.IsCancellationRequested && + (state.Delivered || (failure is AdbcException adbc && adbc.Status == AdbcStatusCode.Cancelled)); + if (cancelled) + { + return new OperationCanceledException(failure.Message, failure, cancellationToken); + } + + ExceptionDispatchInfo.Capture(failure).Throw(); + return failure; // not reached + } + + private static void TryCancel(object? state) + { + CancelState cancelState = (CancelState)state!; + try + { + cancelState.Cancel!(); + cancelState.Delivered = true; Review Comment: Good catch, and confirmed. The race doesn't reproduce from the xUnit test thread, because its SynchronizationContext prevents the continuation from running inline. Cancelling from a thread-pool thread reproduces it: the caller got `TimeoutException` instead of `OperationCanceledException`. Fixed in ed6512562. When the legacy task isn't already complete, `WithCancellation` now resumes through a continuation scheduled on the thread pool, so it never runs inside `Cancel`. Disposing the registration then runs on a different thread from the callback, so it waits for `Cancel` to return and `Delivered` is accurate. A throwing `Cancel` is still treated as not delivered. The regression test is `CancelThatCompletesALegacyTaskInlineIsReportedAsCancellation`. ########## csharp/src/Apache.Arrow.Adbc/AdbcStatement.cs: ########## @@ -83,9 +89,49 @@ public virtual void BindStream(IArrowArrayStream stream) /// A <see cref="ValueTuple"/> where the first item is the number /// of records and the second is the <see cref="IArrowArrayStream"/>. /// </returns> - public virtual async ValueTask<QueryResult> ExecuteQueryAsync() + /// <remarks> + /// <para>New code should call <see cref="ExecuteQueryAsync(CancellationToken)"/>.</para> + /// <para>The default implementation runs <see cref="ExecuteQuery"/> on the thread pool. + /// It doesn't call <see cref="ExecuteQueryAsync(CancellationToken)"/>, because that + /// overload's default implementation calls this one. A driver that overrides + /// <see cref="ExecuteQueryAsync(CancellationToken)"/> should also override this + /// overload to call it with <see cref="CancellationToken.None"/>, so that callers of + /// either overload get the driver's asynchronous implementation.</para> + /// </remarks> + public virtual ValueTask<QueryResult> ExecuteQueryAsync() + { + return new ValueTask<QueryResult>(AdbcAsync.Run(ExecuteQuery, Cancel, CancellationToken.None)); + } + + /// <summary> + /// Executes the statement and returns a tuple containing the number + /// of records and the <see cref="IArrowArrayStream"/>. + /// </summary> + /// <param name="cancellationToken"> + /// A token used to cancel the operation. Cancellation calls <see cref="Cancel"/>. + /// </param> + /// <returns> + /// A <see cref="ValueTuple"/> where the first item is the number + /// of records and the second is the <see cref="IArrowArrayStream"/>. + /// </returns> + /// <remarks> + /// The default implementation calls <see cref="ExecuteQueryAsync()"/>, so that drivers + /// written before this overload existed keep their asynchronous behavior, and calls + /// <see cref="Cancel"/> if the token is cancelled. + /// </remarks> + public virtual ValueTask<QueryResult> ExecuteQueryAsync(CancellationToken cancellationToken) { - return await Task.Run(() => ExecuteQuery()); + if (cancellationToken.IsCancellationRequested) + { + return new ValueTask<QueryResult>(Task.FromCanceled<QueryResult>(cancellationToken)); + } + + ValueTask<QueryResult> result = ExecuteQueryAsync(); Review Comment: Agreed. Fixed in ed6512562. `WithCancellation` now takes a function that starts the task, and the token is registered before the function is called, so `Cancel` can reach an override that blocks before returning its task. This applies to both the query and the update overloads. The regression tests are `SynchronouslyBlockingLegacyQueryOverrideCanBeCancelled` and `SynchronouslyBlockingLegacyUpdateOverrideCanBeCancelled`. ########## csharp/src/Apache.Arrow.Adbc/AdbcStatement.cs: ########## @@ -109,10 +166,45 @@ public virtual Schema ExecuteSchema() /// Executes an update command and returns the number of /// records effected. /// </summary> - /// <exception cref="NotImplementedException"></exception> - public virtual async Task<UpdateResult> ExecuteUpdateAsync() + /// <remarks> + /// <para>New code should call <see cref="ExecuteUpdateAsync(CancellationToken)"/>.</para> + /// <para>The default implementation runs <see cref="ExecuteUpdate"/> on the thread pool. + /// It doesn't call <see cref="ExecuteUpdateAsync(CancellationToken)"/>, because that + /// overload's default implementation calls this one. A driver that overrides + /// <see cref="ExecuteUpdateAsync(CancellationToken)"/> should also override this + /// overload to call it with <see cref="CancellationToken.None"/>, so that callers of + /// either overload get the driver's asynchronous implementation.</para> + /// </remarks> + public virtual Task<UpdateResult> ExecuteUpdateAsync() { - return await Task.Run(() => ExecuteUpdate()); + return AdbcAsync.Run(ExecuteUpdate, Cancel, CancellationToken.None); + } + + /// <summary> + /// Executes an update command and returns the number of + /// records effected. + /// </summary> + /// <param name="cancellationToken"> + /// A token used to cancel the operation. Cancellation calls <see cref="Cancel"/>. + /// </param> + /// <remarks> + /// The default implementation calls <see cref="ExecuteUpdateAsync()"/>, so that drivers + /// written before this overload existed keep their asynchronous behavior, and calls + /// <see cref="Cancel"/> if the token is cancelled. + /// </remarks> + public virtual Task<UpdateResult> ExecuteUpdateAsync(CancellationToken cancellationToken) + { + if (cancellationToken.IsCancellationRequested) + { + return Task.FromCanceled<UpdateResult>(cancellationToken); + } + + Task<UpdateResult> result = ExecuteUpdateAsync(); Review Comment: Agreed. Fixed in ed6512562. The parameterless overload's signature can't change, so while the token overload is forwarding, it records the token in a thread-static field together with the statement. The default parameterless implementation picks the token up and passes it to `Task.Run`, so work that is cancelled while still queued never starts. Legacy overrides are still called as before. If an override calls `base` later, from another thread, or calls another statement's base implementation, it gets `CancellationToken.None`, which is today's behavior. The regression test is `QueryCancelledBeforeTheDefaultWorkStartsDoesNotRun`, for both query and update. It uses an override that cancels the token and then calls `base`, which makes the "cancelled while queued" case deterministic. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
