From 848987378963eaa320fc1d1f317dc8365abb6d60 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?S=C3=A9rgio=20Silveira?= Date: Tue, 25 Aug 2026 17:48:51 +0200 Subject: [PATCH] fix(csharp/src/Client): stop blocking on async calls in the ADO.NET wrapper AdbcCommand overrode only ExecuteDbDataReader and AdbcDataReader only Read(), so both BCL async entry points fell back to the synchronous bodies and reached .Result. One of those sits inside a private method named ReadNextRecordBatchAsync that returns ValueTask and takes a CancellationToken it never used. On a host with a SynchronizationContext the awaited path deadlocks. ReadNextRecordBatchAsync now awaits the stream. AdbcDataReader overrides ReadAsync, keeping the intra-batch path free of a state machine via a cached task. AdbcCommand overrides ExecuteDbDataReaderAsync and awaits AdbcStatement.ExecuteQueryAsync, sharing behavior validation with ExecuteReader. Where blocking legitimately remains, on the synchronous APIs, it now uses AsTask(). A driver's stream is genuinely asynchronous, so reading .Result on the ValueTask it returns is unsupported. That applied to AdbcDataReader.Read and to the schema-loading loop in AdbcConnection, which had the same defect independently of this change. All additive. No driver changes and no public contract change. Read(), GetSchema() and ExecuteDbDataReader are unchanged, and a driver that overrides only ExecuteQuery still gets the base Task.Run implementation. AdbcStatement.ExecuteQueryAsync takes no CancellationToken, so the initial query call stays uncancellable as it is today. Per-batch fetches are cancellable. Closes #4715 --- csharp/src/Client/AdbcCommand.cs | 31 +- csharp/src/Client/AdbcConnection.cs | 2 +- csharp/src/Client/AdbcDataReader.cs | 35 ++- .../Client/ClientTests.cs | 291 ++++++++++++++++++ 4 files changed, 351 insertions(+), 8 deletions(-) diff --git a/csharp/src/Client/AdbcCommand.cs b/csharp/src/Client/AdbcCommand.cs index d87f427c76..e4eb7fd29e 100644 --- a/csharp/src/Client/AdbcCommand.cs +++ b/csharp/src/Client/AdbcCommand.cs @@ -23,6 +23,7 @@ using System.Data.SqlTypes; using System.Globalization; using System.Linq; +using System.Threading; using System.Threading.Tasks; using Apache.Arrow.Types; @@ -209,6 +210,18 @@ protected override DbDataReader ExecuteDbDataReader(CommandBehavior behavior) return ExecuteReader(behavior); } + protected override async Task ExecuteDbDataReaderAsync(CommandBehavior behavior, CancellationToken cancellationToken) + { + bool closeConnection = ValidateReaderBehavior(behavior); + + cancellationToken.ThrowIfCancellationRequested(); + + BindParameters(); + QueryResult result = await AdbcStatement.ExecuteQueryAsync().ConfigureAwait(false); + + return new AdbcDataReader(this, result, this.DecimalBehavior, this.StructBehavior, closeConnection); + } + /// /// Executes the reader with the default behavior. /// @@ -226,21 +239,33 @@ protected override DbDataReader ExecuteDbDataReader(CommandBehavior behavior) /// /// public new AdbcDataReader ExecuteReader(CommandBehavior behavior) + { + bool closeConnection = ValidateReaderBehavior(behavior); + QueryResult result = this.ExecuteQuery(); + + return new AdbcDataReader(this, result, this.DecimalBehavior, this.StructBehavior, closeConnection); + } + + /// + /// Validates the behavior and reports whether the connection should be closed + /// when the reader is disposed. + /// + private bool ValidateReaderBehavior(CommandBehavior behavior) { if (_disposed) throw new ObjectDisposedException(nameof(AdbcCommand)); - bool closeConnection = (behavior & CommandBehavior.CloseConnection) != 0; switch (behavior & ~CommandBehavior.CloseConnection) { case CommandBehavior.SchemaOnly: // The schema is not known until a read happens case CommandBehavior.Default: - QueryResult result = this.ExecuteQuery(); - return new AdbcDataReader(this, result, this.DecimalBehavior, this.StructBehavior, closeConnection); + break; default: throw new InvalidOperationException($"{behavior} is not supported with this provider"); } + + return (behavior & CommandBehavior.CloseConnection) != 0; } protected override void Dispose(bool disposing) diff --git a/csharp/src/Client/AdbcConnection.cs b/csharp/src/Client/AdbcConnection.cs index 7e0d11c3f3..1a0087f2be 100644 --- a/csharp/src/Client/AdbcConnection.cs +++ b/csharp/src/Client/AdbcConnection.cs @@ -551,7 +551,7 @@ public override DataTable GetSchema(Adbc.AdbcConnection adbcConnection, string?[ State state = new State(result, indices.ToArray(), loaders.ToArray()); while (true) { - using (RecordBatch? batch = stream.ReadNextRecordBatchAsync().Result) + using (RecordBatch? batch = stream.ReadNextRecordBatchAsync().AsTask().GetAwaiter().GetResult()) { if (batch == null) { return result; } diff --git a/csharp/src/Client/AdbcDataReader.cs b/csharp/src/Client/AdbcDataReader.cs index 1b7ac02f27..23516b1b9d 100644 --- a/csharp/src/Client/AdbcDataReader.cs +++ b/csharp/src/Client/AdbcDataReader.cs @@ -44,6 +44,8 @@ namespace Apache.Arrow.Adbc.Client /// public sealed class AdbcDataReader : DbDataReader, IDbColumnSchemaGenerator { + private static readonly Task s_true = Task.FromResult(true); + private readonly AdbcCommand adbcCommand; private readonly bool closeConnection; private readonly QueryResult adbcQueryResult; @@ -336,7 +338,30 @@ public override bool Read() // old batch — they must see the exception again immediately. this.recordBatch?.Dispose(); this.recordBatch = null; - this.recordBatch = ReadNextRecordBatchAsync().Result; + + this.recordBatch = ReadNextRecordBatchAsync().AsTask().GetAwaiter().GetResult(); + + return this.recordBatch != null; + } + + public override Task ReadAsync(CancellationToken cancellationToken) + { + if (this.recordBatch != null && this.currentRowInRecordBatch < this.recordBatch.Length - 1) + { + this.currentRowInRecordBatch++; + return s_true; + } + + return FetchNextBatchAsync(cancellationToken); + } + + private async Task FetchNextBatchAsync(CancellationToken cancellationToken) + { + // Clear the previous batch first: a caller retrying after a mid-stream error + // must see the exception again, never stale rows from the old batch. + this.recordBatch?.Dispose(); + this.recordBatch = null; + this.recordBatch = await ReadNextRecordBatchAsync(cancellationToken).ConfigureAwait(false); return this.recordBatch != null; } @@ -389,18 +414,20 @@ public ReadOnlyCollection GetAdbcColumnSchema() /// /// An optional cancellation token /// or null - private ValueTask ReadNextRecordBatchAsync(CancellationToken cancellationToken = default) + private async ValueTask ReadNextRecordBatchAsync(CancellationToken cancellationToken = default) { this.currentRowInRecordBatch = 0; - RecordBatch? recordBatch = this.adbcQueryResult.Stream?.ReadNextRecordBatchAsync(cancellationToken).Result; + RecordBatch? recordBatch = this.adbcQueryResult.Stream is not null + ? await this.adbcQueryResult.Stream.ReadNextRecordBatchAsync(cancellationToken).ConfigureAwait(false) + : null; if (recordBatch != null) { this.TotalBatches += 1; } - return new ValueTask(recordBatch); + return recordBatch; } } } diff --git a/csharp/test/Apache.Arrow.Adbc.Tests/Client/ClientTests.cs b/csharp/test/Apache.Arrow.Adbc.Tests/Client/ClientTests.cs index b8afa26260..5a49fa7649 100644 --- a/csharp/test/Apache.Arrow.Adbc.Tests/Client/ClientTests.cs +++ b/csharp/test/Apache.Arrow.Adbc.Tests/Client/ClientTests.cs @@ -16,8 +16,10 @@ */ using System; +using System.Collections.Concurrent; using System.Collections.Generic; using System.ComponentModel; +using System.Data.Common; using System.Data.SqlTypes; using System.Linq; using System.Threading; @@ -229,6 +231,147 @@ internal void TestConnectionStringParsing(ConnectionStringExample connectionStri Assert.Null(cmd.AdbcCommandTimeoutProperty); } } + + [Fact] + public void ReadAsyncDoesNotDeadlockOnASynchronizationContext() + { + const int timeoutMilliseconds = 5_000; + const int expectedRows = 4; + + using ManualResetEventSlim finished = new ManualResetEventSlim(false); + int rows = 0; + Exception? failure = null; + + Thread thread = new Thread(() => + { + PumpingSynchronizationContext context = new PumpingSynchronizationContext(); + SynchronizationContext.SetSynchronizationContext(context); + + context.Post(async _ => + { + try + { + // syncOnlyDriver: a regressed build must reach ReadAsync, not throw earlier. + AsyncReaderFixture fixture = AsyncReaderFixture.Create( + batchCount: 2, + rowsPerBatch: 2, + syncOnlyDriver: true); + + using (DbDataReader reader = await fixture.Command.ExecuteReaderAsync()) + { + while (await reader.ReadAsync()) + { + rows++; + } + } + } + catch (Exception ex) + { + failure = ex; + } + + finished.Set(); + context.Complete(); + }, null); + + context.Pump(); + }); + + thread.IsBackground = true; + thread.Start(); + + Assert.True(finished.Wait(timeoutMilliseconds), "ReadAsync deadlocked on a SynchronizationContext"); + Assert.Null(failure); + Assert.Equal(expectedRows, rows); + } + + [Fact] + public async Task ExecuteReaderAsyncUsesTheAsyncStatementPath() + { + AsyncReaderFixture fixture = AsyncReaderFixture.Create(batchCount: 1, rowsPerBatch: 1); + + using (DbDataReader reader = await fixture.Command.ExecuteReaderAsync()) + { + } + + fixture.Statement.Verify(x => x.ExecuteQueryAsync(), Times.Once); + fixture.Statement.Verify(x => x.ExecuteQuery(), Times.Never); + } + + [Fact] + public async Task ReadAsyncPassesTheCancellationTokenToTheStream() + { + AsyncReaderFixture fixture = AsyncReaderFixture.Create(batchCount: 2, rowsPerBatch: 1); + + using CancellationTokenSource cts = new CancellationTokenSource(); + using DbDataReader reader = await fixture.Command.ExecuteReaderAsync(cts.Token); + + Assert.True(await reader.ReadAsync(cts.Token)); + Assert.Equal(cts.Token, fixture.Stream.LastToken); + + cts.Cancel(); + + await Assert.ThrowsAnyAsync(() => reader.ReadAsync(cts.Token)); + } + + [Fact] + public async Task ReadAsyncRethrowsAfterAMidStreamErrorInsteadOfServingStaleRows() + { + AsyncReaderFixture fixture = AsyncReaderFixture.Create(batchCount: 2, rowsPerBatch: 2, throwAtCall: 1); + + using DbDataReader reader = await fixture.Command.ExecuteReaderAsync(); + + Assert.True(await reader.ReadAsync()); + Assert.True(await reader.ReadAsync()); + + await Assert.ThrowsAsync(() => reader.ReadAsync()); + await Assert.ThrowsAsync(() => reader.ReadAsync()); + } + + [Fact] + public async Task ExecuteReaderAsyncStillWorksWhenTheDriverOnlyOverridesExecuteQuery() + { + const int batchCount = 2; + const int rowsPerBatch = 2; + int expectedRows = batchCount * rowsPerBatch; + + AsyncReaderFixture fixture = AsyncReaderFixture.Create(batchCount, rowsPerBatch, syncOnlyDriver: true); + + int rows = 0; + + using (DbDataReader reader = await fixture.Command.ExecuteReaderAsync()) + { + while (await reader.ReadAsync()) + { + rows++; + } + } + + Assert.Equal(expectedRows, rows); + fixture.Statement.Verify(x => x.ExecuteQuery(), Times.Once); + } + + [Fact] + public void ReadStillDrainsAnAsynchronousStream() + { + const int batchCount = 2; + const int rowsPerBatch = 3; + int expectedRows = batchCount * rowsPerBatch; + + AsyncReaderFixture fixture = AsyncReaderFixture.Create(batchCount, rowsPerBatch, syncOnlyDriver: true); + + int rows = 0; + + using (AdbcDataReader reader = fixture.Command.ExecuteReader()) + { + while (reader.Read()) + { + rows++; + } + } + + Assert.Equal(expectedRows, rows); + } } internal class ConnectionStringExample @@ -324,4 +467,152 @@ public ValueTask ReadNextRecordBatchAsync(CancellationToken cancell return new ValueTask(this.recordBatches[calls]); } } + + /// + /// An that suspends before producing a batch, and can + /// fail on a chosen fetch. + /// + class AsyncArrayStream : IArrowArrayStream + { + private readonly List recordBatches; + private readonly Schema schema; + private readonly int throwAtCall; + + // start at -1 to use the count of calls as the index + private int calls = -1; + + public AsyncArrayStream(Schema schema, List recordBatches, int throwAtCall) + { + this.schema = schema; + this.recordBatches = recordBatches; + this.throwAtCall = throwAtCall; + } + + public Schema Schema => this.schema; + + public CancellationToken LastToken { get; private set; } + + public void Dispose() { } + + public async ValueTask ReadNextRecordBatchAsync(CancellationToken cancellationToken = default) + { + this.LastToken = cancellationToken; + + // Yield captures the ambient SynchronizationContext, matching FlightSqlResult. + await Task.Yield(); + + cancellationToken.ThrowIfCancellationRequested(); + + this.calls++; + + if (this.throwAtCall >= 0 && this.calls >= this.throwAtCall) + throw new InvalidOperationException("stream failed mid-read"); + + return this.calls < this.recordBatches.Count ? this.recordBatches[this.calls] : null!; + } + } + + /// + /// A single-threaded with a message pump, as WPF, + /// WinForms and classic ASP.NET install. Continuations run only when the owning thread + /// returns to the pump. + /// + class PumpingSynchronizationContext : SynchronizationContext + { + private readonly BlockingCollection> queue = + new BlockingCollection>(); + + public override void Post(SendOrPostCallback d, object? state) + { + try + { + this.queue.Add(new KeyValuePair(d, state)); + } + catch (InvalidOperationException) + { + // the pump has already been completed + } + } + + public override void Send(SendOrPostCallback d, object? state) => d(state); + + public void Pump() + { + foreach (KeyValuePair item in this.queue.GetConsumingEnumerable()) + { + item.Key(item.Value); + } + } + + public void Complete() => this.queue.CompleteAdding(); + } + + /// + /// Builds an over a mocked statement and an + /// . + /// + class AsyncReaderFixture + { + private AsyncReaderFixture(AdbcCommand command, Mock statement, AsyncArrayStream stream) + { + Command = command; + Statement = statement; + Stream = stream; + } + + public AdbcCommand Command { get; } + + public Mock Statement { get; } + + public AsyncArrayStream Stream { get; } + + /// + /// Stub only, leaving the base + /// to supply the result. + /// + public static AsyncReaderFixture Create( + int batchCount, + int rowsPerBatch, + int throwAtCall = -1, + bool syncOnlyDriver = false) + { + List fields = new List() { new Field("n", Int32Type.Default, true) }; + Schema schema = new Schema(fields, new List>()); + + List batches = new List(); + + for (int batch = 0; batch < batchCount; batch++) + { + Int32Array.Builder builder = new Int32Array.Builder(); + + for (int row = 0; row < rowsPerBatch; row++) + { + builder.Append((batch * rowsPerBatch) + row); + } + + Int32Array array = builder.Build(); + batches.Add(new RecordBatch(schema, new List() { array }, array.Length)); + } + + AsyncArrayStream stream = new AsyncArrayStream(schema, batches, throwAtCall); + QueryResult queryResult = new QueryResult(batchCount * rowsPerBatch, stream); + + Mock mockStatement = new Mock(); + + if (syncOnlyDriver) + { + mockStatement.CallBase = true; + mockStatement.Setup(x => x.ExecuteQuery()).Returns(queryResult); + } + else + { + mockStatement.Setup(x => x.ExecuteQueryAsync()).Returns(new ValueTask(queryResult)); + } + + AdbcClient.AdbcConnection connection = new AdbcClient.AdbcConnection(); + AdbcCommand command = new AdbcCommand(mockStatement.Object, connection); + + return new AsyncReaderFixture(command, mockStatement, stream); + } + } }