Skip to content
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.

Commit 781325c

Browse files
Wraith2lcheunglci
Wraith2
authored and
lcheunglci
committedMar 31, 2023
Add | Adding disposable stack temp ref struct and use (dotnet#1818)
1 parent 9bd90e6 commit 781325c

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>

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

+58-54
Original file line numberDiff line numberDiff line change
@@ -4983,6 +4983,7 @@ private void AssertReaderState(bool requireData, bool permitAsync, int? columnIn
49834983
public override Task<bool> NextResultAsync(CancellationToken cancellationToken)
49844984
{
49854985
using (TryEventScope.Create("<sc.SqlDataReader.NextResultAsync|API> {0}", ObjectID))
4986+
using (var registrationHolder = new DisposableTemporaryOnStack<CancellationTokenRegistration>())
49864987
{
49874988

49884989
TaskCompletionSource<bool> source = new TaskCompletionSource<bool>();
@@ -4993,15 +4994,14 @@ public override Task<bool> NextResultAsync(CancellationToken cancellationToken)
49934994
return source.Task;
49944995
}
49954996

4996-
IDisposable registration = null;
49974997
if (cancellationToken.CanBeCanceled)
49984998
{
49994999
if (cancellationToken.IsCancellationRequested)
50005000
{
50015001
source.SetCanceled();
50025002
return source.Task;
50035003
}
5004-
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
5004+
registrationHolder.Set(cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command));
50055005
}
50065006

50075007
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
@@ -5019,7 +5019,7 @@ public override Task<bool> NextResultAsync(CancellationToken cancellationToken)
50195019
return source.Task;
50205020
}
50215021

5022-
return InvokeAsyncCall(new HasNextResultAsyncCallContext(this, source, registration));
5022+
return InvokeAsyncCall(new HasNextResultAsyncCallContext(this, source, registrationHolder.Take()));
50235023
}
50245024
}
50255025

@@ -5320,17 +5320,17 @@ out bytesRead
53205320
public override Task<bool> ReadAsync(CancellationToken cancellationToken)
53215321
{
53225322
using (TryEventScope.Create("<sc.SqlDataReader.ReadAsync|API> {0}", ObjectID))
5323+
using (var registrationHolder = new DisposableTemporaryOnStack<CancellationTokenRegistration>())
53235324
{
53245325
if (IsClosed)
53255326
{
53265327
return ADP.CreatedTaskWithException<bool>(ADP.ExceptionWithStackTrace(ADP.DataReaderClosed("ReadAsync")));
53275328
}
53285329

53295330
// Register first to catch any already expired tokens to be able to trigger cancellation event.
5330-
IDisposable registration = null;
53315331
if (cancellationToken.CanBeCanceled)
53325332
{
5333-
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
5333+
registrationHolder.Set(cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command));
53345334
}
53355335

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

54375437
Debug.Assert(context._reader == null && context._source == null && context._disposable == null, "cached ReadAsyncCallContext was not properly disposed");
54385438

5439-
context.Set(this, source, registration);
5439+
context.Set(this, source, registrationHolder.Take());
54405440
context._hasMoreData = more;
54415441
context._hasReadRowToken = rowTokenRead;
54425442

@@ -5568,41 +5568,43 @@ override public Task<bool> IsDBNullAsync(int i, CancellationToken cancellationTo
55685568
return ADP.CreatedTaskWithException<bool>(ex);
55695569
}
55705570

5571-
// Setup and check for pending task
5572-
TaskCompletionSource<bool> source = new TaskCompletionSource<bool>();
5573-
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
5574-
if (original != null)
5571+
using (var registrationHolder = new DisposableTemporaryOnStack<CancellationTokenRegistration>())
55755572
{
5576-
source.SetException(ADP.ExceptionWithStackTrace(ADP.AsyncOperationPending()));
5577-
return source.Task;
5578-
}
5573+
// Setup and check for pending task
5574+
TaskCompletionSource<bool> source = new TaskCompletionSource<bool>();
5575+
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
5576+
if (original != null)
5577+
{
5578+
source.SetException(ADP.ExceptionWithStackTrace(ADP.AsyncOperationPending()));
5579+
return source.Task;
5580+
}
55795581

5580-
// Check if cancellation due to close is requested (this needs to be done after setting _currentTask)
5581-
if (_cancelAsyncOnCloseToken.IsCancellationRequested)
5582-
{
5583-
source.SetCanceled();
5584-
_currentTask = null;
5585-
return source.Task;
5586-
}
5582+
// Check if cancellation due to close is requested (this needs to be done after setting _currentTask)
5583+
if (_cancelAsyncOnCloseToken.IsCancellationRequested)
5584+
{
5585+
source.SetCanceled();
5586+
_currentTask = null;
5587+
return source.Task;
5588+
}
55875589

5588-
// Setup cancellations
5589-
IDisposable registration = null;
5590-
if (cancellationToken.CanBeCanceled)
5591-
{
5592-
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
5593-
}
5590+
// Setup cancellations
5591+
if (cancellationToken.CanBeCanceled)
5592+
{
5593+
registrationHolder.Set(cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command));
5594+
}
55945595

5595-
IsDBNullAsyncCallContext context = Interlocked.Exchange(ref _cachedIsDBNullContext, null) ?? new IsDBNullAsyncCallContext();
5596+
IsDBNullAsyncCallContext context = Interlocked.Exchange(ref _cachedIsDBNullContext, null) ?? new IsDBNullAsyncCallContext();
55965597

5597-
Debug.Assert(context._reader == null && context._source == null && context._disposable == null, "cached ISDBNullAsync context not properly disposed");
5598+
Debug.Assert(context._reader == null && context._source == null && context._disposable == null, "cached ISDBNullAsync context not properly disposed");
55985599

5599-
context.Set(this, source, registration);
5600-
context._columnIndex = i;
5600+
context.Set(this, source, registrationHolder.Take());
5601+
context._columnIndex = i;
56015602

5602-
// Setup async
5603-
PrepareAsyncInvocation(useSnapshot: true);
5603+
// Setup async
5604+
PrepareAsyncInvocation(useSnapshot: true);
56045605

5605-
return InvokeAsyncCall(context);
5606+
return InvokeAsyncCall(context);
5607+
}
56065608
}
56075609
}
56085610

@@ -5704,31 +5706,33 @@ override public Task<T> GetFieldValueAsync<T>(int i, CancellationToken cancellat
57045706
return ADP.CreatedTaskWithException<T>(ex);
57055707
}
57065708

5707-
// Setup and check for pending task
5708-
TaskCompletionSource<T> source = new TaskCompletionSource<T>();
5709-
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
5710-
if (original != null)
5709+
using (var registrationHolder = new DisposableTemporaryOnStack<CancellationTokenRegistration>())
57115710
{
5712-
source.SetException(ADP.ExceptionWithStackTrace(ADP.AsyncOperationPending()));
5713-
return source.Task;
5714-
}
5711+
// Setup and check for pending task
5712+
TaskCompletionSource<T> source = new TaskCompletionSource<T>();
5713+
Task original = Interlocked.CompareExchange(ref _currentTask, source.Task, null);
5714+
if (original != null)
5715+
{
5716+
source.SetException(ADP.ExceptionWithStackTrace(ADP.AsyncOperationPending()));
5717+
return source.Task;
5718+
}
57155719

5716-
// Check if cancellation due to close is requested (this needs to be done after setting _currentTask)
5717-
if (_cancelAsyncOnCloseToken.IsCancellationRequested)
5718-
{
5719-
source.SetCanceled();
5720-
_currentTask = null;
5721-
return source.Task;
5722-
}
5720+
// Check if cancellation due to close is requested (this needs to be done after setting _currentTask)
5721+
if (_cancelAsyncOnCloseToken.IsCancellationRequested)
5722+
{
5723+
source.SetCanceled();
5724+
_currentTask = null;
5725+
return source.Task;
5726+
}
57235727

5724-
// Setup cancellations
5725-
IDisposable registration = null;
5726-
if (cancellationToken.CanBeCanceled)
5727-
{
5728-
registration = cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command);
5729-
}
5728+
// Setup cancellations
5729+
if (cancellationToken.CanBeCanceled)
5730+
{
5731+
registrationHolder.Set(cancellationToken.Register(SqlCommand.s_cancelIgnoreFailure, _command));
5732+
}
57305733

5731-
return InvokeAsyncCall(new GetFieldValueAsyncCallContext<T>(this, source, registration, i));
5734+
return InvokeAsyncCall(new GetFieldValueAsyncCallContext<T>(this, source, registrationHolder.Take(), i));
5735+
}
57325736
}
57335737

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

‎src/Microsoft.Data.SqlClient/src/Microsoft/Data/SqlClient/DisposableTemporaryOnStack.cs

+40
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
// Licensed to the .NET Foundation under one or more agreements.
2+
// The .NET Foundation licenses this file to you under the MIT license.
3+
// See the LICENSE file in the project root for more information.
4+
5+
6+
using System;
7+
8+
namespace Microsoft.Data.SqlClient
9+
{
10+
internal ref struct DisposableTemporaryOnStack<T>
11+
where T : IDisposable
12+
{
13+
private T _value;
14+
private bool _hasValue;
15+
16+
public void Set(T value)
17+
{
18+
_value = value;
19+
_hasValue = true;
20+
}
21+
22+
public T Take()
23+
{
24+
T value = _value;
25+
_value = default;
26+
_hasValue = false;
27+
return value;
28+
}
29+
30+
public void Dispose()
31+
{
32+
if (_hasValue)
33+
{
34+
_value.Dispose();
35+
_value = default;
36+
_hasValue = false;
37+
}
38+
}
39+
}
40+
}

0 commit comments

Comments
 (0)
Please sign in to comment.