Skip to content

Commit 5001097

Browse files
lcheunglcilcheunglci
authored
Fix | Adding disposable stack temp ref struct and use (#1980)
1 parent 9bd90e6 commit 5001097

File tree

5 files changed

+174
-120
lines changed

5 files changed

+174
-120
lines changed

‎src/Microsoft.Data.SqlClient/netcore/src/Microsoft.Data.SqlClient.csproj

+3
Original file line numberDiff line numberDiff line change
@@ -115,6 +115,9 @@
115115
<Compile Include="..\..\src\Microsoft\Data\SqlClient\DataClassification\SensitivityClassification.cs">
116116
<Link>Microsoft\Data\SqlClient\DataClassification\SensitivityClassification.cs</Link>
117117
</Compile>
118+
<Compile Include="..\..\src\Microsoft\Data\SqlClient\DisposableTemporaryOnStack.cs">
119+
<Link>Microsoft\Data\SqlClient\DisposableTemporaryOnStack.cs</Link>
120+
</Compile>
118121
<Compile Include="..\..\src\Microsoft\Data\SqlClient\EnclaveDelegate.cs">
119122
<Link>Microsoft\Data\SqlClient\EnclaveDelegate.cs</Link>
120123
</Compile>

‎src/Microsoft.Data.SqlClient/netcore/src/Microsoft/Data/SqlClient/SqlDataReader.cs

+70-66
Original file line numberDiff line numberDiff line change
@@ -4397,6 +4397,7 @@ private void AssertReaderState(bool requireData, bool permitAsync, int? columnIn
43974397
public override Task<bool> NextResultAsync(CancellationToken cancellationToken)
43984398
{
43994399
using (TryEventScope.Create("SqlDataReader.NextResultAsync | API | Object Id {0}", ObjectID))
4400+
using (var registrationHolder = new DisposableTemporaryOnStack<CancellationTokenRegistration>())
44004401
{
44014402
TaskCompletionSource<bool> source = new TaskCompletionSource<bool>();
44024403

@@ -4406,15 +4407,14 @@ public override Task<bool> NextResultAsync(CancellationToken cancellationToken)
44064407
return source.Task;
44074408
}
44084409

4409-
IDisposable registration = null;
44104410
if (cancellationToken.CanBeCanceled)
44114411
{
44124412
if (cancellationToken.IsCancellationRequested)
44134413
{
44144414
source.SetCanceled();
44154415
return source.Task;
44164416
}
4417-
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
4417+
registrationHolder.Set(cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command));
44184418
}
44194419

44204420
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
@@ -4432,7 +4432,7 @@ public override Task<bool> NextResultAsync(CancellationToken cancellationToken)
44324432
return source.Task;
44334433
}
44344434

4435-
return InvokeAsyncCall(new HasNextResultAsyncCallContext(this, source, registration));
4435+
return InvokeAsyncCall(new HasNextResultAsyncCallContext(this, source, registrationHolder.Take()));
44364436
}
44374437
}
44384438

@@ -4727,17 +4727,17 @@ out bytesRead
47274727
public override Task<bool> ReadAsync(CancellationToken cancellationToken)
47284728
{
47294729
using (TryEventScope.Create("SqlDataReader.ReadAsync | API | Object Id {0}", ObjectID))
4730+
using (var registrationHolder = new DisposableTemporaryOnStack<CancellationTokenRegistration>())
47304731
{
47314732
if (IsClosed)
47324733
{
47334734
return Task.FromException<bool>(ADP.ExceptionWithStackTrace(ADP.DataReaderClosed()));
47344735
}
47354736

47364737
// Register first to catch any already expired tokens to be able to trigger cancellation event.
4737-
IDisposable registration = null;
47384738
if (cancellationToken.CanBeCanceled)
47394739
{
4740-
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
4740+
registrationHolder.Set(cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command));
47414741
}
47424742

47434743
// If user's token is canceled, return a canceled task
@@ -4850,7 +4850,7 @@ public override Task<bool> ReadAsync(CancellationToken cancellationToken)
48504850

48514851
Debug.Assert(context.Reader == null && context.Source == null && context.Disposable == null, "cached ReadAsyncCallContext was not properly disposed");
48524852

4853-
context.Set(this, source, registration);
4853+
context.Set(this, source, registrationHolder.Take());
48544854
context._hasMoreData = more;
48554855
context._hasReadRowToken = rowTokenRead;
48564856

@@ -4988,49 +4988,51 @@ override public Task<bool> IsDBNullAsync(int i, CancellationToken cancellationTo
49884988
return Task.FromException<bool>(ex);
49894989
}
49904990

4991-
// Setup and check for pending task
4992-
TaskCompletionSource<bool> source = new TaskCompletionSource<bool>();
4993-
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
4994-
if (original != null)
4991+
using (var registrationHolder = new DisposableTemporaryOnStack<CancellationTokenRegistration>())
49954992
{
4996-
source.SetException(ADP.ExceptionWithStackTrace(ADP.AsyncOperationPending()));
4997-
return source.Task;
4998-
}
4993+
// Setup and check for pending task
4994+
TaskCompletionSource<bool> source = new TaskCompletionSource<bool>();
4995+
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
4996+
if (original != null)
4997+
{
4998+
source.SetException(ADP.ExceptionWithStackTrace(ADP.AsyncOperationPending()));
4999+
return source.Task;
5000+
}
49995001

5000-
// Check if cancellation due to close is requested (this needs to be done after setting _currentTask)
5001-
if (_cancelAsyncOnCloseToken.IsCancellationRequested)
5002-
{
5003-
source.SetCanceled();
5004-
_currentTask = null;
5005-
return source.Task;
5006-
}
5002+
// Check if cancellation due to close is requested (this needs to be done after setting _currentTask)
5003+
if (_cancelAsyncOnCloseToken.IsCancellationRequested)
5004+
{
5005+
source.SetCanceled();
5006+
_currentTask = null;
5007+
return source.Task;
5008+
}
50075009

5008-
// Setup cancellations
5009-
IDisposable registration = null;
5010-
if (cancellationToken.CanBeCanceled)
5011-
{
5012-
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
5013-
}
5010+
// Setup cancellations
5011+
if (cancellationToken.CanBeCanceled)
5012+
{
5013+
registrationHolder.Set(cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command));
5014+
}
50145015

5015-
IsDBNullAsyncCallContext context = null;
5016-
if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection)
5017-
{
5018-
context = Interlocked.Exchange(ref sqlInternalConnection.CachedDataReaderIsDBNullContext, null);
5019-
}
5020-
if (context is null)
5021-
{
5022-
context = new IsDBNullAsyncCallContext();
5023-
}
5016+
IsDBNullAsyncCallContext context = null;
5017+
if (_connection?.InnerConnection is SqlInternalConnection sqlInternalConnection)
5018+
{
5019+
context = Interlocked.Exchange(ref sqlInternalConnection.CachedDataReaderIsDBNullContext, null);
5020+
}
5021+
if (context is null)
5022+
{
5023+
context = new IsDBNullAsyncCallContext();
5024+
}
50245025

5025-
Debug.Assert(context.Reader == null && context.Source == null && context.Disposable == null, "cached ISDBNullAsync context not properly disposed");
5026+
Debug.Assert(context.Reader == null && context.Source == null && context.Disposable == default, "cached ISDBNullAsync context not properly disposed");
50265027

5027-
context.Set(this, source, registration);
5028-
context._columnIndex = i;
5028+
context.Set(this, source, registrationHolder.Take());
5029+
context._columnIndex = i;
50295030

5030-
// Setup async
5031-
PrepareAsyncInvocation(useSnapshot: true);
5031+
// Setup async
5032+
PrepareAsyncInvocation(useSnapshot: true);
50325033

5033-
return InvokeAsyncCall(context);
5034+
return InvokeAsyncCall(context);
5035+
}
50345036
}
50355037
}
50365038

@@ -5135,37 +5137,39 @@ override public Task<T> GetFieldValueAsync<T>(int i, CancellationToken cancellat
51355137
return Task.FromException<T>(ex);
51365138
}
51375139

5138-
// Setup and check for pending task
5139-
TaskCompletionSource<T> source = new TaskCompletionSource<T>();
5140-
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
5141-
if (original != null)
5140+
using (var registrationHolder = new DisposableTemporaryOnStack<CancellationTokenRegistration>())
51425141
{
5143-
source.SetException(ADP.ExceptionWithStackTrace(ADP.AsyncOperationPending()));
5144-
return source.Task;
5145-
}
5142+
// Setup and check for pending task
5143+
TaskCompletionSource<T> source = new TaskCompletionSource<T>();
5144+
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
5145+
if (original != null)
5146+
{
5147+
source.SetException(ADP.ExceptionWithStackTrace(ADP.AsyncOperationPending()));
5148+
return source.Task;
5149+
}
51465150

5147-
// Check if cancellation due to close is requested (this needs to be done after setting _currentTask)
5148-
if (_cancelAsyncOnCloseToken.IsCancellationRequested)
5149-
{
5150-
source.SetCanceled();
5151-
_currentTask = null;
5152-
return source.Task;
5153-
}
5151+
// Check if cancellation due to close is requested (this needs to be done after setting _currentTask)
5152+
if (_cancelAsyncOnCloseToken.IsCancellationRequested)
5153+
{
5154+
source.SetCanceled();
5155+
_currentTask = null;
5156+
return source.Task;
5157+
}
51545158

5155-
// Setup cancellations
5156-
IDisposable registration = null;
5157-
if (cancellationToken.CanBeCanceled)
5158-
{
5159-
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
5160-
}
5159+
// Setup cancellations
5160+
if (cancellationToken.CanBeCanceled)
5161+
{
5162+
registrationHolder.Set(cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command));
5163+
}
51615164

5162-
// Setup async
5163-
PrepareAsyncInvocation(useSnapshot: true);
5165+
// Setup async
5166+
PrepareAsyncInvocation(useSnapshot: true);
51645167

5165-
GetFieldValueAsyncCallContext<T> context = new GetFieldValueAsyncCallContext<T>(this, source, registration);
5166-
context._columnIndex = i;
5168+
GetFieldValueAsyncCallContext<T> context = new GetFieldValueAsyncCallContext<T>(this, source, registrationHolder.Take());
5169+
context._columnIndex = i;
51675170

5168-
return InvokeAsyncCall(context);
5171+
return InvokeAsyncCall(context);
5172+
}
51695173
}
51705174

51715175
private static Task<T> GetFieldValueAsyncExecute<T>(Task task, object state)

‎src/Microsoft.Data.SqlClient/netfx/src/Microsoft.Data.SqlClient.csproj

+3
Original file line numberDiff line numberDiff line change
@@ -185,6 +185,9 @@
185185
<Compile Include="..\..\src\Microsoft\Data\SqlClient\DataClassification\SensitivityClassification.cs">
186186
<Link>Microsoft\Data\SqlClient\DataClassification\SensitivityClassification.cs</Link>
187187
</Compile>
188+
<Compile Include="..\..\src\Microsoft\Data\SqlClient\DisposableTemporaryOnStack.cs">
189+
<Link>Microsoft\Data\SqlClient\DisposableTemporaryOnStack.cs</Link>
190+
</Compile>
188191
<Compile Include="..\..\src\Microsoft\Data\SqlClient\EnclaveDelegate.cs">
189192
<Link>Microsoft\Data\SqlClient\EnclaveDelegate.cs</Link>
190193
</Compile>

0 commit comments

Comments
 (0)