diff --git a/src/Npgsql/NpgsqlBinaryExporter.cs b/src/Npgsql/NpgsqlBinaryExporter.cs index f221056119..9473c95959 100644 --- a/src/Npgsql/NpgsqlBinaryExporter.cs +++ b/src/Npgsql/NpgsqlBinaryExporter.cs @@ -24,7 +24,7 @@ public sealed class NpgsqlBinaryExporter : ICancelable NpgsqlConnector _connector; NpgsqlReadBuffer _buf; - bool _isConsumed, _isDisposed; + ExporterState _state = ExporterState.Uninitialized; long _endOfMessagePos; short _column; @@ -91,6 +91,7 @@ internal async Task Init(string copyToCommand, bool async, CancellationToken can throw _connector.UnexpectedMessageReceived(msg.Code); } + _state = ExporterState.Ready; NumColumns = copyOutResponse.NumColumns; _columnInfoCache = new PgConverterInfo[NumColumns]; _rowsExported = 0; @@ -141,7 +142,7 @@ async Task ReadHeader(bool async) async ValueTask StartRow(bool async, CancellationToken cancellationToken = default) { ThrowIfDisposed(); - if (_isConsumed) + if (_state == ExporterState.Consumed) return -1; using var registration = _connector.StartNestedCancellableOperation(cancellationToken); @@ -176,7 +177,7 @@ async ValueTask StartRow(bool async, CancellationToken cancellationToken = Expect(await _connector.ReadMessage(async).ConfigureAwait(false), _connector); Expect(await _connector.ReadMessage(async).ConfigureAwait(false), _connector); _column = BeforeRow; - _isConsumed = true; + _state = ExporterState.Consumed; return -1; } @@ -437,7 +438,7 @@ void ThrowIfNotOnRow() void ThrowIfDisposed() { - if (_isDisposed) + if (_state == ExporterState.Disposed) ThrowHelper.ThrowObjectDisposedException(nameof(NpgsqlBinaryExporter), "The COPY operation has already ended."); } @@ -472,10 +473,10 @@ public Task CancelAsync() async ValueTask DisposeAsync(bool async) { - if (_isDisposed) + if (_state == ExporterState.Disposed) return; - if (_isConsumed) + if (_state is ExporterState.Consumed or ExporterState.Uninitialized) { LogMessages.BinaryCopyOperationCompleted(_copyLogger, _rowsExported, _connector.Id); } @@ -512,7 +513,7 @@ async ValueTask DisposeAsync(bool async) void Cleanup() { - Debug.Assert(!_isDisposed); + Debug.Assert(_state != ExporterState.Disposed); var connector = _connector; if (!ReferenceEquals(connector, null)) @@ -523,9 +524,21 @@ void Cleanup() } _buf = null!; - _isDisposed = true; + _state = ExporterState.Disposed; } } #endregion + + #region Enums + + enum ExporterState + { + Uninitialized, + Ready, + Consumed, + Disposed + } + + #endregion Enums } diff --git a/src/Npgsql/NpgsqlBinaryImporter.cs b/src/Npgsql/NpgsqlBinaryImporter.cs index 633d1bac15..52c9438fde 100644 --- a/src/Npgsql/NpgsqlBinaryImporter.cs +++ b/src/Npgsql/NpgsqlBinaryImporter.cs @@ -25,7 +25,7 @@ public sealed class NpgsqlBinaryImporter : ICancelable NpgsqlConnector _connector; NpgsqlWriteBuffer _buf; - ImporterState _state; + ImporterState _state = ImporterState.Uninitialized; /// /// The number of columns in the current (not-yet-written) row. @@ -99,6 +99,7 @@ internal async Task Init(string copyFromCommand, bool async, CancellationToken c throw _connector.UnexpectedMessageReceived(msg.Code); } + _state = ImporterState.Ready; _params = new NpgsqlParameter[copyInResponse.NumColumns]; _rowsImported = 0; _buf.StartCopyMode(); @@ -512,6 +513,7 @@ async ValueTask CloseAsync(bool async, CancellationToken cancellationToken = def case ImporterState.Ready: await Cancel(async, cancellationToken).ConfigureAwait(false); break; + case ImporterState.Uninitialized: case ImporterState.Cancelled: case ImporterState.Committed: break; @@ -553,6 +555,7 @@ void CheckReady() static void Throw(ImporterState state) => throw (state switch { + ImporterState.Uninitialized => throw new InvalidOperationException("The COPY operation has not been initialized."), ImporterState.Disposed => new ObjectDisposedException(typeof(NpgsqlBinaryImporter).FullName, "The COPY operation has already ended."), ImporterState.Cancelled => new InvalidOperationException("The COPY operation has already been cancelled."), @@ -567,6 +570,7 @@ static void Throw(ImporterState state) enum ImporterState { + Uninitialized, Ready, Committed, Cancelled, diff --git a/src/Npgsql/NpgsqlConnection.cs b/src/Npgsql/NpgsqlConnection.cs index c63f5bc4b6..6bda4f71d5 100644 --- a/src/Npgsql/NpgsqlConnection.cs +++ b/src/Npgsql/NpgsqlConnection.cs @@ -1160,17 +1160,26 @@ async Task BeginBinaryImport(bool async, string copyFromCo LogMessages.StartingBinaryImport(connector.LoggingConfiguration.CopyLogger, connector.Id); // no point in passing a cancellationToken here, as we register the cancellation in the Init method connector.StartUserAction(ConnectorState.Copy, attemptPgCancellation: false); + var importer = new NpgsqlBinaryImporter(connector); try { - var importer = new NpgsqlBinaryImporter(connector); await importer.Init(copyFromCommand, async, cancellationToken).ConfigureAwait(false); connector.CurrentCopyOperation = importer; return importer; } catch { - connector.EndUserAction(); - EndBindingScope(ConnectorBindingScope.Copy); + try + { + if (async) + await importer.DisposeAsync().ConfigureAwait(false); + else + importer.Dispose(); + } + catch + { + // ignored + } throw; } } @@ -1210,17 +1219,26 @@ async Task BeginBinaryExport(bool async, string copyToComm LogMessages.StartingBinaryExport(connector.LoggingConfiguration.CopyLogger, connector.Id); // no point in passing a cancellationToken here, as we register the cancellation in the Init method connector.StartUserAction(ConnectorState.Copy, attemptPgCancellation: false); + var exporter = new NpgsqlBinaryExporter(connector); try { - var exporter = new NpgsqlBinaryExporter(connector); await exporter.Init(copyToCommand, async, cancellationToken).ConfigureAwait(false); connector.CurrentCopyOperation = exporter; return exporter; } catch { - connector.EndUserAction(); - EndBindingScope(ConnectorBindingScope.Copy); + try + { + if (async) + await exporter.DisposeAsync().ConfigureAwait(false); + else + exporter.Dispose(); + } + catch + { + // ignored + } throw; } } @@ -1266,9 +1284,9 @@ async Task BeginTextImport(bool async, string copyFromCommand, Cance LogMessages.StartingTextImport(connector.LoggingConfiguration.CopyLogger, connector.Id); // no point in passing a cancellationToken here, as we register the cancellation in the Init method connector.StartUserAction(ConnectorState.Copy, attemptPgCancellation: false); + var copyStream = new NpgsqlRawCopyStream(connector); try { - var copyStream = new NpgsqlRawCopyStream(connector); await copyStream.Init(copyFromCommand, async, cancellationToken).ConfigureAwait(false); var writer = new NpgsqlCopyTextWriter(connector, copyStream); connector.CurrentCopyOperation = writer; @@ -1276,8 +1294,17 @@ async Task BeginTextImport(bool async, string copyFromCommand, Cance } catch { - connector.EndUserAction(); - EndBindingScope(ConnectorBindingScope.Copy); + try + { + if (async) + await copyStream.DisposeAsync().ConfigureAwait(false); + else + copyStream.Dispose(); + } + catch + { + // ignored + } throw; } } @@ -1323,9 +1350,9 @@ async Task BeginTextExport(bool async, string copyToCommand, Cancell LogMessages.StartingTextExport(connector.LoggingConfiguration.CopyLogger, connector.Id); // no point in passing a cancellationToken here, as we register the cancellation in the Init method connector.StartUserAction(ConnectorState.Copy, attemptPgCancellation: false); + var copyStream = new NpgsqlRawCopyStream(connector); try { - var copyStream = new NpgsqlRawCopyStream(connector); await copyStream.Init(copyToCommand, async, cancellationToken).ConfigureAwait(false); var reader = new NpgsqlCopyTextReader(connector, copyStream); connector.CurrentCopyOperation = reader; @@ -1333,8 +1360,17 @@ async Task BeginTextExport(bool async, string copyToCommand, Cancell } catch { - connector.EndUserAction(); - EndBindingScope(ConnectorBindingScope.Copy); + try + { + if (async) + await copyStream.DisposeAsync().ConfigureAwait(false); + else + copyStream.Dispose(); + } + catch + { + // ignored + } throw; } } @@ -1380,9 +1416,9 @@ async Task BeginRawBinaryCopy(bool async, string copyComman LogMessages.StartingRawCopy(connector.LoggingConfiguration.CopyLogger, connector.Id); // no point in passing a cancellationToken here, as we register the cancellation in the Init method connector.StartUserAction(ConnectorState.Copy, attemptPgCancellation: false); + var stream = new NpgsqlRawCopyStream(connector); try { - var stream = new NpgsqlRawCopyStream(connector); await stream.Init(copyCommand, async, cancellationToken).ConfigureAwait(false); if (!stream.IsBinary) { @@ -1395,8 +1431,17 @@ async Task BeginRawBinaryCopy(bool async, string copyComman } catch { - connector.EndUserAction(); - EndBindingScope(ConnectorBindingScope.Copy); + try + { + if (async) + await stream.DisposeAsync().ConfigureAwait(false); + else + stream.Dispose(); + } + catch + { + // ignored + } throw; } } diff --git a/src/Npgsql/NpgsqlRawCopyStream.cs b/src/Npgsql/NpgsqlRawCopyStream.cs index 3648b24075..664c39a1b8 100644 --- a/src/Npgsql/NpgsqlRawCopyStream.cs +++ b/src/Npgsql/NpgsqlRawCopyStream.cs @@ -28,7 +28,7 @@ public sealed class NpgsqlRawCopyStream : Stream, ICancelable NpgsqlWriteBuffer _writeBuf; int _leftToReadInDataMsg; - bool _isDisposed, _isConsumed; + CopyStreamState _state = CopyStreamState.Uninitialized; bool _canRead; bool _canWrite; @@ -84,12 +84,14 @@ internal async Task Init(string copyCommand, bool async, CancellationToken cance switch (msg.Code) { case BackendMessageCode.CopyInResponse: + _state = CopyStreamState.Ready; var copyInResponse = (CopyInResponseMessage) msg; IsBinary = copyInResponse.IsBinary; _canWrite = true; _writeBuf.StartCopyMode(); break; case BackendMessageCode.CopyOutResponse: + _state = CopyStreamState.Ready; var copyOutResponse = (CopyOutResponseMessage) msg; IsBinary = copyOutResponse.IsBinary; _canRead = true; @@ -245,7 +247,7 @@ async ValueTask ReadAsyncInternal() async ValueTask ReadCore(int count, bool async, CancellationToken cancellationToken = default) { - if (_isConsumed) + if (_state == CopyStreamState.Consumed) return 0; using var registration = _connector.StartNestedCancellableOperation(cancellationToken, attemptPgCancellation: false); @@ -261,7 +263,7 @@ async ValueTask ReadCore(int count, bool async, CancellationToken cancellat } catch { - if (!_isDisposed) + if (_state != CopyStreamState.Disposed) Cleanup(); throw; } @@ -274,7 +276,7 @@ async ValueTask ReadCore(int count, bool async, CancellationToken cancellat case BackendMessageCode.CopyDone: Expect(await _connector.ReadMessage(async).ConfigureAwait(false), _connector); Expect(await _connector.ReadMessage(async).ConfigureAwait(false), _connector); - _isConsumed = true; + _state = CopyStreamState.Consumed; return 0; default: throw _connector.UnexpectedMessageReceived(msg.Code); @@ -331,6 +333,9 @@ async Task Cancel(bool async) } catch (PostgresException e) { + // TODO: NpgsqlBinaryImporter doesn't cleanup on cancellation + // And instead relies on users disposing the object + // We probably should do the same here Cleanup(); if (e.SqlState != PostgresErrorCodes.QueryCanceled) @@ -355,7 +360,7 @@ public override ValueTask DisposeAsync() async ValueTask DisposeAsync(bool disposing, bool async) { - if (_isDisposed || !disposing) + if (_state == CopyStreamState.Disposed || !disposing) return; try @@ -373,7 +378,7 @@ async ValueTask DisposeAsync(bool disposing, bool async) } else { - if (!_isConsumed) + if (_state != CopyStreamState.Consumed && _state != CopyStreamState.Uninitialized) { try { @@ -403,7 +408,7 @@ async ValueTask DisposeAsync(bool disposing, bool async) #pragma warning disable CS8625 void Cleanup() { - Debug.Assert(!_isDisposed); + Debug.Assert(_state != CopyStreamState.Disposed); LogMessages.CopyOperationCompleted(_copyLogger, _connector.Id); _connector.EndUserAction(); _connector.CurrentCopyOperation = null; @@ -411,13 +416,13 @@ void Cleanup() _connector = null; _readBuf = null; _writeBuf = null; - _isDisposed = true; + _state = CopyStreamState.Disposed; } #pragma warning restore CS8625 void CheckDisposed() { - if (_isDisposed) { + if (_state == CopyStreamState.Disposed) { throw new ObjectDisposedException(nameof(NpgsqlRawCopyStream), "The COPY operation has already ended."); } } @@ -452,6 +457,18 @@ static void ValidateArguments(byte[] buffer, int offset, int count) ThrowHelper.ThrowArgumentException("Offset and length were out of bounds for the array or count is greater than the number of elements from index to the end of the source collection."); } #endregion + + #region Enums + + enum CopyStreamState + { + Uninitialized, + Ready, + Consumed, + Disposed + } + + #endregion Enums } ///