Skip to content

Commit

Permalink
Fixes ReadAsync() behavior to register Cancellation token action befo…
Browse files Browse the repository at this point in the history
…re streaming results (#1781)
  • Loading branch information
cheenamalhotra authored Oct 4, 2022
1 parent 81055cf commit 871c0d2
Show file tree
Hide file tree
Showing 4 changed files with 98 additions and 12 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -4733,6 +4733,13 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
return Task.FromException<bool>(ADP.ExceptionWithStackTrace(ADP.DataReaderClosed()));
}

// Register first to catch any already expired tokens to be able to trigger cancellation event.
IDisposable registration = null;
if (cancellationToken.CanBeCanceled)
{
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
}

// If user's token is canceled, return a canceled task
if (cancellationToken.IsCancellationRequested)
{
Expand Down Expand Up @@ -4831,12 +4838,6 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
return source.Task;
}

IDisposable registration = null;
if (cancellationToken.CanBeCanceled)
{
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
}

ReadAsyncCallContext context = null;
if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection)
{
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5326,6 +5326,13 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
return ADP.CreatedTaskWithException<bool>(ADP.ExceptionWithStackTrace(ADP.DataReaderClosed("ReadAsync")));
}

// Register first to catch any already expired tokens to be able to trigger cancellation event.
IDisposable registration = null;
if (cancellationToken.CanBeCanceled)
{
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
}

// If user's token is canceled, return a canceled task
if (cancellationToken.IsCancellationRequested)
{
Expand Down Expand Up @@ -5425,12 +5432,6 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
return source.Task;
}

IDisposable registration = null;
if (cancellationToken.CanBeCanceled)
{
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
}

var context = Interlocked.Exchange(ref _cachedReadAsyncContext, null) ?? new ReadAsyncCallContext();

Debug.Assert(context._reader == null && context._source == null && context._disposable == null, "cached ReadAsyncCallContext was not properly disposed");
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -86,6 +86,7 @@
<Compile Include="SQL\ConnectivityTests\TcpDefaultForAzureTest.cs" />
<Compile Include="SQL\DataBaseSchemaTest\ConnectionSchemaTest.cs" />
<Compile Include="SQL\DataClassificationTest\DataClassificationTest.cs" />
<Compile Include="SQL\DataReaderTest\DataReaderCancellationTest.cs" />
<Compile Include="SQL\DataReaderTest\DataReaderStreamsTest.cs" />
<Compile Include="SQL\DataReaderTest\DataReaderTest.cs" />
<Compile Include="SQL\DataStreamTest\DataStreamTest.cs" />
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,83 @@
// Licensed to the .NET Foundation under one or more agreements.
// The .NET Foundation licenses this file to you under the MIT license.
// See the LICENSE file in the project root for more information.

using System.Diagnostics;
using System.Threading;
using System.Threading.Tasks;
using Xunit;

namespace Microsoft.Data.SqlClient.ManualTesting.Tests
{
public class DataReaderCancellationTest
{
/// <summary>
/// Test ensures cancellation token is registered before ReadAsync starts processing results from TDS Stream,
/// such that when Cancel is triggered, the token is capable of canceling reading further results.
/// </summary>
/// <returns>Async Task</returns>
[ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
public static async Task CancellationTokenIsRespected_ReadAsync()
{
const string longRunningQuery = @"
with TenRows as (select Value from (values (1), (2), (3), (4), (5), (6), (7), (8), (9), (10)) as TenRows (Value)),
ThousandRows as (select A.Value as A, B.Value as B, C.Value as C from TenRows as A, TenRows as B, TenRows as C)
select *
from ThousandRows as A, ThousandRows as B, ThousandRows as C;";

using (var source = new CancellationTokenSource())
using (var connection = new SqlConnection(DataTestUtility.TCPConnectionString))
{
await connection.OpenAsync(source.Token);

Stopwatch stopwatch = Stopwatch.StartNew();
await Assert.ThrowsAsync<TaskCanceledException>(async () =>
{
using (var command = new SqlCommand(longRunningQuery, connection))
using (var reader = await command.ExecuteReaderAsync(source.Token))
{
while (await reader.ReadAsync(source.Token))
{
source.Cancel();
}
}
});
Assert.True(stopwatch.ElapsedMilliseconds < 10000, "Cancellation did not trigger on time.");
}
}

/// <summary>
/// Test ensures cancellation token is registered before ReadAsync starts processing results from TDS Stream,
/// such that when Cancel is triggered, the token is capable of canceling reading further results.
/// </summary>
/// <returns>Async Task</returns>
[ConditionalFact(typeof(DataTestUtility), nameof(DataTestUtility.AreConnStringsSetup))]
public static async Task CancelledCancellationTokenIsRespected_ReadAsync()
{
const string longRunningQuery = @"
with TenRows as (select Value from (values (1), (2), (3), (4), (5), (6), (7), (8), (9), (10)) as TenRows (Value)),
ThousandRows as (select A.Value as A, B.Value as B, C.Value as C from TenRows as A, TenRows as B, TenRows as C)
select *
from ThousandRows as A, ThousandRows as B, ThousandRows as C;";

using (var source = new CancellationTokenSource())
using (var connection = new SqlConnection(DataTestUtility.TCPConnectionString))
{
await connection.OpenAsync(source.Token);

Stopwatch stopwatch = Stopwatch.StartNew();
await Assert.ThrowsAsync<TaskCanceledException>(async () =>
{
using (var command = new SqlCommand(longRunningQuery, connection))
using (var reader = await command.ExecuteReaderAsync(source.Token))
{
source.Cancel();
while (await reader.ReadAsync(source.Token))
{ }
}
});
Assert.True(stopwatch.ElapsedMilliseconds < 10000, "Cancellation did not trigger on time.");
}
}
}
}

0 comments on commit 871c0d2

Please sign in to comment.