Skip to content

Commit 8b186c6

Browse files
rojiroji
committed
Implement DisposeAsync
The main scenario to be handled asynchronously, is when a connection (or transaction) is disposed by an open reader is still active, with its resultset needing to be consumed. Fixes npgsql#2597
1 parent e59da52 commit 8b186c6

5 files changed

Lines changed: 138 additions & 59 deletions

File tree

‎src/Npgsql/ConnectorPool.cs‎

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -452,16 +452,12 @@ void CloseConnector(NpgsqlConnector connector, bool wasIdle)
452452
openCount = prevState.Open - 1;
453453
}
454454
else
455-
{
456455
openCount = Interlocked.Decrement(ref State.Open);
457-
}
458456

459457
// Unblock a single waiter, if any, to get the slot that just opened up.
460458
while (_waiting.TryDequeue(out var waiter))
461-
{
462459
if (waiter.TaskCompletionSource.TrySetResult(null))
463460
break;
464-
}
465461

466462
// Only turn off the timer one time, when it was this Close that brought Open back to _min.
467463
if (openCount == _min)

‎src/Npgsql/NpgsqlConnection.cs‎

Lines changed: 75 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -593,60 +593,86 @@ public override void EnlistTransaction(Transaction transaction)
593593
#region Close
594594

595595
/// <summary>
596-
/// releases the connection to the database. If the connection is pooled, it will be
597-
/// made available for re-use. If it is non-pooled, the actual connection will be shutdown.
596+
/// Releases the connection. If the connection is pooled, it will be returned to the pull and made available for re-use.
597+
/// If it is non-pooled, the physical connection will be closed.
598598
/// </summary>
599-
public override void Close() => Close(false);
599+
public override void Close() => Close(wasBroken: false, async: false);
600600

601-
internal void Close(bool wasBroken)
601+
/// <summary>
602+
/// Releases the connection. If the connection is pooled, it will be returned to the pull and made available for re-use.
603+
/// If it is non-pooled, the physical connection will be closed.
604+
/// </summary>
605+
#if !NET461 && !NETSTANDARD2_0
606+
public override Task CloseAsync()
607+
#else
608+
public Task CloseAsync()
609+
#endif
610+
{
611+
using (NoSynchronizationContextScope.Enter())
612+
return Close(wasBroken: false, async: true);
613+
}
614+
615+
internal Task Close(bool wasBroken, bool async)
602616
{
603617
if (Connector == null)
604-
return;
618+
return Task.CompletedTask;
605619
var connectorId = Connector.Id;
606620
Log.Trace("Closing connection...", connectorId);
607621
_wasBroken = wasBroken;
608622

609-
Connector.CloseOngoingOperations();
623+
if (Connector.HasOngoingOperation)
624+
return CloseOngoingOperationAndFinish();
610625

611-
// The connector has closed us during CloseOngoingOperations due to an underlying failure.
612-
if (Connector == null)
613-
return;
626+
FinishClose();
627+
return Task.CompletedTask;
614628

615-
if (Settings.Pooling)
629+
async Task CloseOngoingOperationAndFinish()
616630
{
617-
if (EnlistedTransaction == null)
618-
_pool!.Release(Connector);
619-
else
631+
await Connector!.CloseOngoingOperations(async);
632+
633+
// The connector has closed us during CloseOngoingOperations due to an underlying failure.
634+
if (Connector == null)
635+
return;
636+
637+
FinishClose();
638+
}
639+
640+
void FinishClose()
641+
{
642+
var connector = Connector!;
643+
if (Settings.Pooling)
620644
{
621-
// A System.Transactions transaction is still in progress, we need to wait for it to complete.
622-
// Close the connection and disconnect it from the resource manager but leave the connector
623-
// in a enlisted pending list in the pool.
624-
_pool!.AddPendingEnlistedConnector(Connector, EnlistedTransaction);
625-
Connector.Connection = null;
645+
if (EnlistedTransaction == null)
646+
_pool!.Release(connector);
647+
else
648+
{
649+
// A System.Transactions transaction is still in progress, we need to wait for it to complete.
650+
// Close the connection and disconnect it from the resource manager but leave the connector
651+
// in a enlisted pending list in the pool.
652+
_pool!.AddPendingEnlistedConnector(connector, EnlistedTransaction);
653+
connector.Connection = null;
654+
EnlistedTransaction = null;
655+
}
656+
}
657+
else // Non-pooled connection
658+
{
659+
if (EnlistedTransaction == null)
660+
connector.Close();
661+
// If a non-pooled connection is being closed but is enlisted in an ongoing
662+
// TransactionScope, simply detach the connector from the connection and leave
663+
// it open. It will be closed when the TransactionScope is disposed.
664+
connector.Connection = null;
626665
EnlistedTransaction = null;
627666
}
628-
}
629-
else // Non-pooled connection
630-
{
631-
if (EnlistedTransaction == null)
632-
Connector.Close();
633-
// If a non-pooled connection is being closed but is enlisted in an ongoing
634-
// TransactionScope, simply detach the connector from the connection and leave
635-
// it open. It will be closed when the TransactionScope is disposed.
636-
Connector.Connection = null;
637-
EnlistedTransaction = null;
638-
}
639-
640-
Log.Debug("Connection closed", connectorId);
641667

642-
Connector = null;
643-
644-
OnStateChange(OpenToClosedEventArgs);
668+
Log.Debug("Connection closed", connectorId);
669+
Connector = null;
670+
OnStateChange(OpenToClosedEventArgs);
671+
}
645672
}
646673

647674
/// <summary>
648-
/// Releases all resources used by the
649-
/// <see cref="NpgsqlConnection">NpgsqlConnection</see>.
675+
/// Releases all resources used by the <see cref="NpgsqlConnection">NpgsqlConnection</see>.
650676
/// </summary>
651677
/// <param name="disposing"><b>true</b> when called from Dispose();
652678
/// <b>false</b> when being called from the finalizer.</param>
@@ -656,10 +682,22 @@ protected override void Dispose(bool disposing)
656682
return;
657683
if (disposing)
658684
Close();
659-
base.Dispose(disposing);
660685
_disposed = true;
661686
}
662687

688+
#if !NET461 && !NETSTANDARD2_0
689+
/// <summary>
690+
/// Releases all resources used by the <see cref="NpgsqlConnection">NpgsqlConnection</see>.
691+
/// </summary>
692+
public override async ValueTask DisposeAsync()
693+
{
694+
if (_disposed)
695+
return;
696+
await CloseAsync();
697+
_disposed = true;
698+
}
699+
#endif
700+
663701
#endregion
664702

665703
#region Notifications and Notices

‎src/Npgsql/NpgsqlConnector.cs‎

Lines changed: 23 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1208,30 +1208,36 @@ void DoCancelRequest(int backendProcessId, int backendSecretKey)
12081208

12091209
#region Close / Reset
12101210

1211+
internal bool HasOngoingOperation => CurrentReader != null || CurrentCopyOperation != null;
1212+
12111213
/// <summary>
12121214
/// Closes ongoing operations, i.e. an open reader exists or a COPY operation still in progress, as
12131215
/// part of a connection close.
12141216
/// Does nothing if the thread has been aborted - the connector will be closed immediately.
12151217
/// </summary>
1216-
internal void CloseOngoingOperations()
1218+
internal async Task CloseOngoingOperations(bool async)
12171219
{
1218-
CurrentReader?.Close(true, false);
1219-
var currentCopyOperation = CurrentCopyOperation;
1220-
if (currentCopyOperation == null)
1220+
var reader = CurrentReader;
1221+
var copyOperation = CurrentCopyOperation;
1222+
1223+
if (reader != null)
1224+
await reader.Close(connectionClosing: true, async);
1225+
1226+
if (copyOperation == null)
12211227
return;
12221228

12231229
// TODO: There's probably a race condition as the COPY operation may finish on its own during the next few lines
12241230

12251231
// Note: we only want to cancel import operations, since in these cases cancel is safe.
12261232
// Export cancellations go through the PostgreSQL "asynchronous" cancel mechanism and are
12271233
// therefore vulnerable to the race condition in #615.
1228-
if (currentCopyOperation is NpgsqlBinaryImporter ||
1229-
currentCopyOperation is NpgsqlCopyTextWriter ||
1230-
currentCopyOperation is NpgsqlRawCopyStream rawCopyStream && rawCopyStream.CanWrite)
1234+
if (copyOperation is NpgsqlBinaryImporter ||
1235+
copyOperation is NpgsqlCopyTextWriter ||
1236+
copyOperation is NpgsqlRawCopyStream rawCopyStream && rawCopyStream.CanWrite)
12311237
{
12321238
try
12331239
{
1234-
currentCopyOperation.Cancel();
1240+
copyOperation.Cancel();
12351241
}
12361242
catch (Exception e)
12371243
{
@@ -1241,14 +1247,16 @@ currentCopyOperation is NpgsqlCopyTextWriter ||
12411247

12421248
try
12431249
{
1244-
currentCopyOperation.Dispose();
1250+
copyOperation.Dispose();
12451251
}
12461252
catch (Exception e)
12471253
{
12481254
Log.Warn("Error while disposing cancelled COPY on connector close", e, Id);
12491255
}
12501256
}
12511257

1258+
// TODO in theory this should be async-optional, but the only I/O done here is the Terminate Flush, which is
1259+
// very unlikely to block (plus locking would need to be worked out)
12521260
internal void Close()
12531261
{
12541262
lock (this)
@@ -1323,22 +1331,25 @@ internal void Break()
13231331
// Note that the connection's full state is usually calculated from the connector's, but in
13241332
// states closed/broken the connector is null. We therefore need a way to distinguish between
13251333
// Closed and Broken on the connection.
1326-
conn.Close(true);
1334+
conn.Close(wasBroken: true, async: false);
13271335
}
13281336
}
13291337
}
13301338

13311339
/// <summary>
13321340
/// Closes the socket and cleans up client-side resources associated with this connector.
13331341
/// </summary>
1342+
/// <remarks>
1343+
/// This method doesn't actually perform any meaningful I/O, and therefore is sync-only.
1344+
/// </remarks>
13341345
void Cleanup()
13351346
{
13361347
Debug.Assert(Monitor.IsEntered(this));
13371348

13381349
Log.Trace("Cleaning up connector", Id);
13391350
try
13401351
{
1341-
_stream?.Dispose();
1352+
_stream.Dispose();
13421353
}
13431354
catch
13441355
{
@@ -1350,6 +1361,7 @@ void Cleanup()
13501361
CurrentReader.Command.State = CommandState.Idle;
13511362
try
13521363
{
1364+
// Will never complete asynchronously (stream is already closed)
13531365
CurrentReader.Close();
13541366
}
13551367
catch

‎src/Npgsql/NpgsqlDataReader.cs‎

Lines changed: 13 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -760,10 +760,21 @@ async Task Consume(bool async)
760760
/// </summary>
761761
protected override void Dispose(bool disposing) => Close();
762762

763+
#if !NET461 && !NETSTANDARD2_0
764+
/// <summary>
765+
/// Releases the resources used by the <see cref="NpgsqlDataReader">NpgsqlDataReader</see>.
766+
/// </summary>
767+
public override ValueTask DisposeAsync()
768+
{
769+
using (NoSynchronizationContextScope.Enter())
770+
return new ValueTask(Close(connectionClosing: false, async: true));
771+
}
772+
#endif
773+
763774
/// <summary>
764775
/// Closes the <see cref="NpgsqlDataReader"/> reader, allowing a new command to be executed.
765776
/// </summary>
766-
public override void Close() => Close(false, false).GetAwaiter().GetResult();
777+
public override void Close() => Close(connectionClosing: false, async: false).GetAwaiter().GetResult();
767778

768779
/// <summary>
769780
/// Closes the <see cref="NpgsqlDataReader"/> reader, allowing a new command to be executed.
@@ -773,7 +784,7 @@ public override Task CloseAsync()
773784
#else
774785
public Task CloseAsync()
775786
#endif
776-
=> Close(false, true);
787+
=> Close(connectionClosing: false, async: true);
777788

778789
internal async Task Close(bool connectionClosing, bool async)
779790
{

‎src/Npgsql/NpgsqlTransaction.cs‎

Lines changed: 27 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -156,12 +156,12 @@ public Task CommitAsync(CancellationToken cancellationToken = default)
156156
/// </summary>
157157
public override void Rollback() => Rollback(false).GetAwaiter().GetResult();
158158

159-
async Task Rollback(bool async)
159+
Task Rollback(bool async)
160160
{
161161
CheckReady();
162-
if (!_connector.DatabaseInfo.SupportsTransactions)
163-
return;
164-
await _connector.Rollback(async);
162+
return _connector.DatabaseInfo.SupportsTransactions
163+
? _connector.Rollback(async)
164+
: Task.CompletedTask;
165165
}
166166

167167
/// <summary>
@@ -312,13 +312,35 @@ protected override void Dispose(bool disposing)
312312

313313
if (disposing && !IsCompleted)
314314
{
315-
_connector.CloseOngoingOperations();
315+
_connector.CloseOngoingOperations(async: false).GetAwaiter().GetResult();
316316
Rollback();
317317
}
318318

319319
IsDisposed = true;
320320
}
321321

322+
#if !NET461 && !NETSTANDARD2_0
323+
/// <summary>
324+
/// Disposes the transaction, rolling it back if it is still pending.
325+
/// </summary>
326+
public override async ValueTask DisposeAsync()
327+
{
328+
if (IsDisposed)
329+
return;
330+
331+
if (!IsCompleted)
332+
{
333+
using (NoSynchronizationContextScope.Enter())
334+
{
335+
await _connector.CloseOngoingOperations(async: true);
336+
await Rollback(async: true);
337+
}
338+
}
339+
340+
IsDisposed = true;
341+
}
342+
#endif
343+
322344
/// <summary>
323345
/// Disposes the transaction, without rolling back. Used only in special circumstances, e.g. when
324346
/// the connection is broken.

0 commit comments

Comments
 (0)