Skip to content

Async Generic Helpers #4334

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
benrr101 wants to merge 10 commits into
base: main
Choose a base branch
from
Open

Async Generic Helpers #4334

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
cc8da2f
Extract AsyncHelper from SqlUtil.cs into the utilities namespace
benrr101 Oct 10, 2025
e20fc9b
Implement AsyncHelpers with generic state, add unit tests
benrr101 Jun 2, 2026
ab424fa
Change existing async helper calls in SqlBulkCopy.cs
benrr101 Jun 2, 2026
1608dfd
Change existing async helper calls in TdsParser.cs
benrr101 Jun 2, 2026
80beeff
Change existing async helper calls in TdsParser.cs
benrr101 Jun 2, 2026
734fa77
Change existing async helper calls in SqlCommand
benrr101 Jun 2, 2026
c0127ec
Remove old tests
benrr101 Jun 3, 2026
e516e3b
Comments from copilot round 1
benrr101 Jun 4, 2026
8a3ff94
Address feedback from wraith - including constraining state types to …
benrr101 Jun 5, 2026
e85c55c
Fix tests for last change
benrr101 Jun 5, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 7 additions & 5 deletions ...Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/Connection/SqlConnectionInternal.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -6,20 +6,22 @@
using System.Collections.Generic;
using System.Data.Common;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
using System.Globalization;
using System.Security;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using System.Transactions;
using Microsoft.Data.Common;
using Microsoft.Data.Common.ConnectionString;
using Microsoft.Data.ProviderBase;
using Microsoft.Data.SqlClient.Connection;
using Microsoft.Data.SqlClient.ConnectionPool;
using IsolationLevel = System.Data.IsolationLevel;
using Microsoft.Data.SqlClient.Internal;
using Microsoft.Data.SqlClient.Utilities;
using IsolationLevel = System.Data.IsolationLevel;

#if NETFRAMEWORK
using Microsoft.Data.Common.ConnectionString;
#endif

namespace Microsoft.Data.SqlClient.Connection
{
Expand Down Expand Up @@ -400,7 +402,7 @@ internal SqlConnectionInternal(

try
{
// If we want to consider pool operations against the overall connect timeout,
// If we want to consider pool operations against the overall connect timeout,
// use the provided timeout. Otherwise, start a fresh timeout to receive the full
// connect timeout.
_timeout = ResolveLoginTimeout(timeout, connectionOptions.ConnectTimeout);
Expand Down
62 changes: 36 additions & 26 deletions src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlBulkCopy.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@
using Microsoft.Data.Common;
using Microsoft.Data.SqlClient.Connection;
using Microsoft.Data.SqlClient.Internal;
using Microsoft.Data.SqlClient.Utilities;

namespace Microsoft.Data.SqlClient
{
Expand Down Expand Up @@ -776,7 +777,7 @@ private string AnalyzeTargetAndCreateUpdateBulkCommand(BulkCopySimpleResultSet i

// Remove it from our unmatched set.
unmatchedColumns.Remove(localColumn.MappedDestinationColumn);

// Check for column types that we refuse to bulk load, even
// though we found a match.
//
Expand Down Expand Up @@ -2855,18 +2856,17 @@ private Task CopyBatchesAsyncContinued(BulkCopySimpleResultSet internalResults,
task,
source,
state: this,
onSuccess: (object state) =>
onSuccess: state =>
{
SqlBulkCopy sqlBulkCopy = (SqlBulkCopy)state;
Task continuedTask = sqlBulkCopy.CopyBatchesAsyncContinuedOnSuccess(internalResults, updateBulkCommandText, cts, source);
Task continuedTask = state.CopyBatchesAsyncContinuedOnSuccess(internalResults, updateBulkCommandText, cts, source);
if (continuedTask == null)
{
// Continuation finished sync, recall into CopyBatchesAsync to continue
sqlBulkCopy.CopyBatchesAsync(internalResults, updateBulkCommandText, cts, source);
state.CopyBatchesAsync(internalResults, updateBulkCommandText, cts, source);
}
},
onFailure: static (Exception _, object state) => ((SqlBulkCopy)state).CopyBatchesAsyncContinuedOnError(cleanupParser: false),
onCancellation: static (object state) => ((SqlBulkCopy)state).CopyBatchesAsyncContinuedOnError(cleanupParser: true));
onFailure: static (state, _) => state.CopyBatchesAsyncContinuedOnError(cleanupParser: false),
onCancellation: static state => state.CopyBatchesAsyncContinuedOnError(cleanupParser: true));

return source.Task;
}
Expand Down Expand Up @@ -2917,24 +2917,23 @@ private Task CopyBatchesAsyncContinuedOnSuccess(BulkCopySimpleResultSet internal
writeTask,
source,
state: this,
onSuccess: (object state) =>
onSuccess: state =>
{
SqlBulkCopy sqlBulkCopy = (SqlBulkCopy)state;
try
{
sqlBulkCopy.RunParser();
sqlBulkCopy.CommitTransaction();
state.RunParser();
state.CommitTransaction();
}
catch (Exception)
{
sqlBulkCopy.CopyBatchesAsyncContinuedOnError(cleanupParser: false);
state.CopyBatchesAsyncContinuedOnError(cleanupParser: false);
throw;
}

// Always call back into CopyBatchesAsync
sqlBulkCopy.CopyBatchesAsync(internalResults, updateBulkCommandText, cts, source);
state.CopyBatchesAsync(internalResults, updateBulkCommandText, cts, source);
},
onFailure: static (Exception _, object state) => ((SqlBulkCopy)state).CopyBatchesAsyncContinuedOnError(cleanupParser: false));
onFailure: static (state, _ ) => state.CopyBatchesAsyncContinuedOnError(cleanupParser: false));
return source.Task;
}
}
Expand Down Expand Up @@ -3189,21 +3188,20 @@ private void WriteToServerInternalRestAsync(CancellationToken cts, TaskCompletio

// No need to cancel timer since SqlBulkCopy creates specific task source for reconnection.
AsyncHelper.SetTimeoutExceptionWithState(
completion: cancellableReconnectTS,
timeout: BulkCopyTimeout,
taskCompletionSource: cancellableReconnectTS,
timeoutInSeconds: BulkCopyTimeout,
state: _destinationTableName,
onFailure: static state =>
SQL.BulkLoadInvalidDestinationTable((string)state, SQL.CR_ReconnectTimeout()),
onTimeout: static state => SQL.BulkLoadInvalidDestinationTable(state, SQL.CR_ReconnectTimeout()),
cancellationToken: CancellationToken.None
);

AsyncHelper.ContinueTaskWithState(
task: cancellableReconnectTS.Task,
completion: source,
taskToContinue:cancellableReconnectTS.Task,
taskCompletionSource: source,
state: regReconnectCancel,
onSuccess: (object state) =>
onSuccess: state =>
{
((StrongBox<CancellationTokenRegistration>)state).Value.Dispose();
state.Value.Dispose();
if (_parserLock != null)
{
_parserLock.Release();
Expand All @@ -3213,10 +3211,22 @@ private void WriteToServerInternalRestAsync(CancellationToken cts, TaskCompletio
_parserLock.Wait(canReleaseFromAnyThread: true);
WriteToServerInternalRestAsync(cts, source);
},
onFailure: static (_, state) => ((StrongBox<CancellationTokenRegistration>)state).Value.Dispose(),
onCancellation: static state => ((StrongBox<CancellationTokenRegistration>)state).Value.Dispose(),
exceptionConverter: ex => SQL.BulkLoadInvalidDestinationTable(_destinationTableName, ex)
);
onFailure: (regReconnectCancel2, exception) =>
Copy link
Copy Markdown
Contributor

@Wraith2 Wraith2 Jun 4, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can this be static or is it now capturing instance state? if it's capturing that's a slight memory performance regression.

Copy link
Copy Markdown
Contributor Author

@benrr101 benrr101 Jun 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

No this one is not static. There are still improvements that can be made to make more callbacks static. At this point, having a separate overload for a one-off situation where the exception type needs to be changed isn't worth it. It's also worth noting that the success callback is not static, nor is the exception converter itself. As such, I'm not sure it will make much of a difference.

Copy link
Copy Markdown
Contributor

@Wraith2 Wraith2 Jun 5, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

It depends how often it's called. If it's once per field it could be significant. It's unlikely to cause time problems but it could add to memory thrashing. I only called this one out because it's changed from static to non-static, existing non-statics were already being pinned by some sort of capture.

I did a pass through the whole codebase years ago and made everything that was capable static, that's why there's weird use of Tuple and some TaskCompletionSource state smuggling

{
regReconnectCancel2.Value.Dispose();

// Convert exception and set it on the source
// Note: This is safe because the helper will only try to set the
// exception and b/c it is already set will pass without setting
// to the original exception.
Exception convertedException = SQL.BulkLoadInvalidDestinationTable(
_destinationTableName,
exception);
source.TrySetException(convertedException);
},
onCancellation: static regReconnectCancel2 =>
regReconnectCancel2.Value.Dispose());

return;
}
else
Expand Down
24 changes: 11 additions & 13 deletions src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/SqlCommand.Encryption.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,6 +14,7 @@
using System.Threading.Tasks;
using Microsoft.Data.Common;
using Microsoft.Data.SqlClient.Connection;
using Microsoft.Data.SqlClient.Utilities;

namespace Microsoft.Data.SqlClient
{
Expand Down Expand Up @@ -252,23 +253,22 @@ private SqlDataReader GetParameterEncryptionDataReader(
bool isRetry)
{
returnTask = AsyncHelper.CreateContinuationTaskWithState(
task: fetchInputParameterEncryptionInfoTask,
taskToContinue: fetchInputParameterEncryptionInfoTask,
state: this,
onSuccess: state =>
onSuccess: sqlCommand =>
{
SqlCommand command = (SqlCommand)state;
bool processFinallyBlockAsync = true;
bool decrementAsyncCountInFinallyBlockAsync = true;

try
{
// Check for any exceptions on network write, before reading.
command.CheckThrowSNIException();
sqlCommand.CheckThrowSNIException();

// If it is async, then TryFetchInputParameterEncryptionInfo ->
// RunExecuteReaderTds would have incremented the async count. Decrement it
// when we are about to complete async execute reader.
SqlConnectionInternal internalConnectionTds = command._activeConnection.GetOpenTdsConnection();
SqlConnectionInternal internalConnectionTds = sqlCommand._activeConnection.GetOpenTdsConnection();
if (internalConnectionTds is not null)
{
internalConnectionTds.DecrementAsyncCount();
Expand All @@ -277,13 +277,13 @@ private SqlDataReader GetParameterEncryptionDataReader(

// Complete executereader.
// @TODO: If we can remove this reference, this could be a static lambda
describeParameterEncryptionDataReader = command.CompleteAsyncExecuteReader(
describeParameterEncryptionDataReader = sqlCommand.CompleteAsyncExecuteReader(
isInternal: false,
forDescribeParameterEncryption: true);
Debug.Assert(command._stateObj is null, "non-null state object in PrepareForTransparentEncryption.");
Debug.Assert(sqlCommand._stateObj is null, "non-null state object in PrepareForTransparentEncryption.");

// Read the results of describe parameter encryption.
command.ReadDescribeEncryptionParameterResults(
sqlCommand.ReadDescribeEncryptionParameterResults(
describeParameterEncryptionDataReader,
describeParameterEncryptionRpcOriginalRpcMap,
isRetry);
Expand All @@ -303,7 +303,7 @@ private SqlDataReader GetParameterEncryptionDataReader(
}
finally
{
command.PrepareTransparentEncryptionFinallyBlock(
sqlCommand.PrepareTransparentEncryptionFinallyBlock(
closeDataReader: processFinallyBlockAsync,
decrementAsyncCount: decrementAsyncCountInFinallyBlockAsync,
clearDataStructures: processFinallyBlockAsync,
Expand All @@ -312,11 +312,9 @@ private SqlDataReader GetParameterEncryptionDataReader(
describeParameterEncryptionDataReader: describeParameterEncryptionDataReader);
}
},
onFailure: static (exception, state) =>
onFailure: static (sqlCommand, exception) =>
{
SqlCommand command = (SqlCommand)state;
command.CachedAsyncState?.ResetAsyncState();

sqlCommand.CachedAsyncState?.ResetAsyncState();
if (exception is not null)
{
throw exception;
Expand Down
Loading
Loading