Skip to content

Decouple read buffer from connector #5462

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

Closed
NinoFloris wants to merge 3 commits into from
Closed

Decouple read buffer from connector #5462

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
083db0a
Prepare for decoupling
NinoFloris Nov 27, 2023
694a28a
Decouple ReadBuffer from Connector and remove IO code duplication
NinoFloris Nov 26, 2023
8a482a1
Remove dependency on socket
NinoFloris Nov 27, 2023
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
136 changes: 104 additions & 32 deletions src/Npgsql/Internal/NpgsqlConnector.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,45 @@

namespace Npgsql.Internal;

readonly struct CancellationScope : IDisposable
{
readonly NpgsqlConnector _connector;
readonly CancellationTokenRegistration _registration;
readonly CancellationToken _previousCancellationToken;
readonly bool _previousAttemptPostgresCancellation;

public CancellationScope(NpgsqlConnector connector, CancellationTokenRegistration registration, CancellationToken previousCancellationToken, bool previousAttemptPostgresCancellation)
{
_connector = connector;
_registration = registration;
_previousCancellationToken = previousCancellationToken;
_previousAttemptPostgresCancellation = previousAttemptPostgresCancellation;
}

public void Dispose()
{
if (_connector is null)
return;

_connector.UserCancellationToken = _previousCancellationToken;
_connector.AttemptPostgresCancellation = _previousAttemptPostgresCancellation;
_registration.Dispose();
}
}

interface IConnectionOperationControl
{
bool Aborted { get; }
void Abort(Exception abortReason);
TimeSpan Cancel(Exception? cancellationReason, CancellationToken? cancellationToken = null);
CancellationToken CurrentCancellationToken { get;}
}

/// <summary>
/// Represents a connection to a PostgreSQL backend. Unlike NpgsqlConnection objects, which are
/// exposed to users, connectors are internal to Npgsql and are recycled by the connection pool.
/// </summary>
public sealed partial class NpgsqlConnector
public sealed partial class NpgsqlConnector : IConnectionOperationControl
{
#region Fields and Properties

Expand Down Expand Up @@ -273,7 +307,7 @@ internal bool PostgresCancellationPerformed
CancellationTokenRegistration _cancellationTokenRegistration;
internal bool UserCancellationRequested => _userCancellationRequested;
internal CancellationToken UserCancellationToken { get; set; }
internal bool AttemptPostgresCancellation { get; private set; }
internal bool AttemptPostgresCancellation { get; set; }
static readonly TimeSpan _cancelImmediatelyTimeout = TimeSpan.FromMilliseconds(-1);

IDisposable? _certificate;
Expand Down Expand Up @@ -743,7 +777,8 @@ async Task RawOpen(SslMode sslMode, NpgsqlTimeout timeout, bool async, Cancellat
RelaxedTextEncoding = Encoding.GetEncoding(Settings.Encoding, EncoderFallback.ReplacementFallback, DecoderFallback.ReplacementFallback);
}

ReadBuffer = new NpgsqlReadBuffer(this, _stream, _socket, Settings.ReadBufferSize, TextEncoding, RelaxedTextEncoding);
ReadBuffer = new NpgsqlReadBuffer(this, DataSource.MetricsReporter,
_stream, Settings.ReadBufferSize, TextEncoding, RelaxedTextEncoding);
WriteBuffer = new NpgsqlWriteBuffer(this, _stream, _socket, Settings.WriteBufferSize, TextEncoding);

timeout.CheckAndApply(this);
Expand Down Expand Up @@ -1239,6 +1274,8 @@ internal ValueTask<IBackendMessage> ReadMessage(
return new ValueTask<IBackendMessage?>(ParseServerMessage(ReadBuffer, messageCode, len, false))!;
}

bool ReadingNotifications { get; set; }

[AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))]
async ValueTask<IBackendMessage?> ReadMessageLong(
bool async,
Expand Down Expand Up @@ -1269,11 +1306,14 @@ internal ValueTask<IBackendMessage> ReadMessage(

PostgresException? error = null;

if (readingNotifications)
ReadingNotifications = true;
try
{
while (true)
{
await ReadBuffer.Ensure(5, async, readingNotifications).ConfigureAwait(false);

await ReadBuffer.Ensure(5, async).ConfigureAwait(false);
var messageCode = (BackendMessageCode)ReadBuffer.ReadByte();
ValidateBackendMessageCode(messageCode);
var len = ReadBuffer.ReadInt32() - 4; // Transmitted length includes itself
Expand Down Expand Up @@ -1359,7 +1399,8 @@ internal ValueTask<IBackendMessage> ReadMessage(

// Reset flushed bytes after any RFQ or in between potentially long running operations.
// Just in case we'll hit that 15 exbibyte limit of a signed long...
if (messageCode is BackendMessageCode.ReadyForQuery or BackendMessageCode.CopyData or BackendMessageCode.NotificationResponse)
if (messageCode is BackendMessageCode.ReadyForQuery or BackendMessageCode.CopyData
or BackendMessageCode.NotificationResponse)
ReadBuffer.ResetFlushedBytes();

return msg;
Expand Down Expand Up @@ -1387,6 +1428,11 @@ internal ValueTask<IBackendMessage> ReadMessage(
ExceptionDispatchInfo.Capture(error).Throw();
throw;
}
finally
{
if (readingNotifications)
ReadingNotifications = false;
}
}

internal IBackendMessage? ParseResultSetMessage(NpgsqlReadBuffer buf, BackendMessageCode code, int len, bool handleCallbacks = false)
Expand Down Expand Up @@ -1741,6 +1787,9 @@ internal bool PerformPostgresCancellation()
{
Debug.Assert(BackendProcessId != 0, "PostgreSQL cancellation requested by the backend doesn't support it");

if (PostgresCancellationPerformed)
return true;

lock (CancelLock)
{
if (PostgresCancellationPerformed)
Expand Down Expand Up @@ -1798,7 +1847,7 @@ internal CancellationTokenRegistration StartCancellableOperation(
CancellationToken cancellationToken = default,
bool attemptPgCancellation = true)
{
_userCancellationRequested = PostgresCancellationPerformed = false;
_readPgCancellationEffect = _userCancellationRequested = PostgresCancellationPerformed = false;
UserCancellationToken = cancellationToken;
ReadBuffer.Cts.ResetCts();

Expand Down Expand Up @@ -1827,7 +1876,7 @@ internal CancellationTokenRegistration StartCancellableOperation(
/// PostgreSQL cancellation will be skipped and client-socket cancellation will occur immediately.
/// </param>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
internal NestedCancellableScope StartNestedCancellableOperation(
internal CancellationScope StartNestedCancellableOperation(
CancellationToken cancellationToken = default,
bool attemptPgCancellation = true)
{
Expand All @@ -1841,31 +1890,6 @@ internal NestedCancellableScope StartNestedCancellableOperation(
return new(this, registration, currentUserCancellationToken, currentAttemptPostgresCancellation);
}

internal readonly struct NestedCancellableScope : IDisposable
{
readonly NpgsqlConnector _connector;
readonly CancellationTokenRegistration _registration;
readonly CancellationToken _previousCancellationToken;
readonly bool _previousAttemptPostgresCancellation;

public NestedCancellableScope(NpgsqlConnector connector, CancellationTokenRegistration registration, CancellationToken previousCancellationToken, bool previousAttemptPostgresCancellation)
{
_connector = connector;
_registration = registration;
_previousCancellationToken = previousCancellationToken;
_previousAttemptPostgresCancellation = previousAttemptPostgresCancellation;
}

public void Dispose()
{
if (_connector is null)
return;

_connector.UserCancellationToken = _previousCancellationToken;
_connector.AttemptPostgresCancellation = _previousAttemptPostgresCancellation;
_registration.Dispose();
}
}

#endregion Cancel

Expand Down Expand Up @@ -2734,6 +2758,54 @@ void ReadParameterStatus(ReadOnlySpan<byte> incomingName, ReadOnlySpan<byte> inc
}

#endregion Misc

bool IConnectionOperationControl.Aborted => IsBroken;
void IConnectionOperationControl.Abort(Exception abortReason) => Break(abortReason);

bool _readPgCancellationEffect;

bool OperationSupportsPostgresCancellation
=> SupportsPostgresCancellation && State is ConnectorState.Executing or ConnectorState.Fetching or ConnectorState.Replication;

TimeSpan IConnectionOperationControl.Cancel(Exception? cancellationReason, CancellationToken? cancellationToken)
{
// If we should attempt PostgreSQL cancellation do it the first time around.
if (!_readPgCancellationEffect && AttemptPostgresCancellation && OperationSupportsPostgresCancellation && PerformPostgresCancellation())
{
// When the timeout is negative the pg cancellation effect read will be skipped and we'll directly abort in the code below.
if (Settings.CancellationTimeout >= 0)
{
_readPgCancellationEffect = true;
return TimeSpan.FromMilliseconds(Settings.CancellationTimeout);
}
}

// When we have read the pg cancellation effect the subsequent Cancel call with a UserCancellationToken
// is due to a read timing out on the previously returned CancellationTimeout.
// If we have done the cancellation effect read and we're called by a future read UserCancellationToken would not be present.
if (cancellationToken is not { IsCancellationRequested: true } || _readPgCancellationEffect && UserCancellationToken == cancellationToken)
cancellationReason = new TimeoutException("Timeout during reading attempt", cancellationReason);

cancellationReason = new NpgsqlException("Exception while reading from stream", cancellationReason);

// Return an OCE if there is a cancellationToken (async caller) and it's either cancelled or a user performed explicit cancellation.
// Synchronous reads end up with a conventional NpgsqlException(TimeoutException) as we cannot cancel an in-progress sync read.
if (cancellationToken is { } token && (token.IsCancellationRequested || _userCancellationRequested))
{
cancellationReason = new OperationCanceledException(
token.IsCancellationRequested ? null : "Read was cancelled out-of-band by user code", cancellationReason, token);
}

// When reading notifications (Wait) there's nothing to cancel, and no breaking of the connection.
if (!ReadingNotifications)
((IConnectionOperationControl)this).Abort(cancellationReason);
throw cancellationReason;
}

// After a cancellation read was attempted the connection was either broken or we're reading more (due to successful cancellation)
// In the latter case CurrentCancellationToken is usually called again, however we shouldn't return the previously cancelled token.
CancellationToken IConnectionOperationControl.CurrentCancellationToken
=> _readPgCancellationEffect ? CancellationToken.None : UserCancellationToken;
}

#region Enums
Expand Down
41 changes: 11 additions & 30 deletions src/Npgsql/Internal/NpgsqlReadBuffer.Stream.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,25 +10,22 @@ sealed partial class NpgsqlReadBuffer
{
internal sealed class ColumnStream : Stream
{
readonly NpgsqlConnector _connector;
readonly NpgsqlReadBuffer _buf;
long _startPos;
int _start;
int _read;
bool _canSeek;
bool _commandScoped;
/// Does not throw ODE.
internal int CurrentLength { get; private set; }
internal bool IsDisposed { get; private set; }

internal ColumnStream(NpgsqlConnector connector)
internal ColumnStream(NpgsqlReadBuffer buffer)
{
_connector = connector;
_buf = connector.ReadBuffer;
_buf = buffer;
IsDisposed = true;
}

internal void Init(int len, bool canSeek, bool commandScoped)
internal void Init(int len, bool canSeek)
{
Debug.Assert(!canSeek || _buf.ReadBytesLeft >= len,
"Seekable stream constructed but not all data is in buffer (sequential)");
Expand All @@ -40,10 +37,11 @@ internal void Init(int len, bool canSeek, bool commandScoped)
CurrentLength = len;
_read = 0;

_commandScoped = commandScoped;
IsDisposed = false;
}

internal void Advance(int count) => _read += count;

public override bool CanRead => true;

public override bool CanWrite => false;
Expand Down Expand Up @@ -160,29 +158,15 @@ public override int Read(Span<byte> span)
if (count == 0)
return 0;

var read = _buf.Read(_commandScoped, span.Slice(0, count));
_read += read;

return read;
return _buf.StreamRead(this, span.Slice(0, count));
}

public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken = default)
{
CheckDisposed();

var count = Math.Min(buffer.Length, CurrentLength - _read);
return count == 0 ? new ValueTask<int>(0) : ReadLong(this, buffer.Slice(0, count), cancellationToken);

static async ValueTask<int> ReadLong(ColumnStream stream, Memory<byte> buffer, CancellationToken cancellationToken = default)
{
using var registration = cancellationToken.CanBeCanceled
? stream._connector.StartNestedCancellableOperation(cancellationToken, attemptPgCancellation: false)
: default;

var read = await stream._buf.ReadAsync(stream._commandScoped, buffer, cancellationToken).ConfigureAwait(false);
stream._read += read;
return read;
}
return count == 0 ? new ValueTask<int>(0) : _buf.StreamReadAsync(this, buffer.Slice(0, count), cancellationToken);
}

public override void Write(byte[] buffer, int offset, int count)
Expand All @@ -205,13 +189,10 @@ async ValueTask DisposeAsync(bool disposing, bool async)
if (IsDisposed || !disposing)
return;

if (!_connector.IsBroken)
{
var pos = _buf.CumulativeReadPosition - _startPos;
var remaining = checked((int)(CurrentLength - pos));
if (remaining > 0)
await _buf.Skip(remaining, async).ConfigureAwait(false);
}
var pos = _buf.CumulativeReadPosition - _startPos;
var remaining = checked((int)(CurrentLength - pos));
if (remaining > 0)
await _buf.Skip(remaining, async).ConfigureAwait(false);

IsDisposed = true;
}
Expand Down
Loading