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