diff --git a/Slon.Benchmark/AdoCommandFlowFactoryBenchmark.cs b/Slon.Benchmark/AdoCommandFlowFactoryBenchmark.cs index f4f50c3..68941dc 100644 --- a/Slon.Benchmark/AdoCommandFlowFactoryBenchmark.cs +++ b/Slon.Benchmark/AdoCommandFlowFactoryBenchmark.cs @@ -144,5 +144,5 @@ public async Task Cleanup() } [UnsafeAccessor(UnsafeAccessorKind.Field, Name = "_batchCore")] - static extern ref AdoBatchCore GetBatchCore(SlonBatch batch); + static extern ref AdoBatchCore GetBatchCore(SlonBatch batch); } diff --git a/Slon.Benchmark/BackendMessagePublicationBenchmark.cs b/Slon.Benchmark/BackendMessagePublicationBenchmark.cs index 9446d76..59f0480 100644 --- a/Slon.Benchmark/BackendMessagePublicationBenchmark.cs +++ b/Slon.Benchmark/BackendMessagePublicationBenchmark.cs @@ -15,8 +15,8 @@ public class BackendMessagePublicationBenchmark [Benchmark] public int PublishBufferedMessages() { - _context.RetireCurrentBatch(); - _context.SetBatch(new BackendMessageBatch(new ReadOnlySequence(_messages))); + _context.RetireCursor(); + _context.SetCursor(new BackendMessageCursor(new ReadOnlySequence(_messages))); var count = 0; while (_context.TryMoveNext()) count++; diff --git a/Slon.Benchmarks/Shared/SlonConnectionPool.cs b/Slon.Benchmarks/Shared/SlonConnectionPool.cs new file mode 100644 index 0000000..01d559b --- /dev/null +++ b/Slon.Benchmarks/Shared/SlonConnectionPool.cs @@ -0,0 +1,270 @@ +using System.Diagnostics; +using System.Globalization; +using System.Net; +using Microsoft.Extensions.ObjectPool; +using Npgsql; +using Slon.Pg; +using Slon.Pg.Protocol; +using Slon.Pg.Protocol.Flows; +using Slon.Pooling; +using Slon.Text; +using Slon.Transport; + +namespace Slon.Fortunes; + +internal sealed class SlonConnectionPool : IAsyncDisposable +{ + const string Query = "SELECT id, message FROM fortune"; + readonly ConnectionPool _pool; + readonly CommandFlowOptions _options; + readonly ObjectPool? _flowPool; + + SlonConnectionPool( + ConnectionPool pool, + Command command, + int flowPoolCapacity) + { + _pool = pool; + _options = new() { Commands = new(command) }; + if (flowPoolCapacity > 0) + { + _flowPool = new DefaultObjectPool( + new CommandFlowPoolPolicy(), flowPoolCapacity); + } + } + + internal static async ValueTask CreateAsync( + string connectionString, + int connectionCount) + { + var builder = new NpgsqlConnectionStringBuilder(connectionString); + var clientOptions = new PgClientOptions + { + EndPoint = new DnsEndPoint( + RequiredPostgreSqlValue("Host", builder.Host), builder.Port), + Database = RequiredPostgreSqlValue("Database", builder.Database), + Username = RequiredPostgreSqlValue("Username", builder.Username), + Password = builder.Password, + Ssl = new PostgreSqlSslOptions { Mode = PostgreSqlSslMode.Disable }, + }; + var transportFactory = SocketStreamConnection.CreateFactory( + clientOptions.EndPoint, + new TransportConnectionOptions { UseZeroByteReads = false }); + var bootstrapFactory = new PgClientProtocolFactory(clientOptions, transportFactory); + var protocolFactory = new PgClientProtocolFactory( + clientOptions, + transportFactory, + static options => options.HeartbeatMode = PgClientProtocolHeartbeatMode.External); + + // Every pooled wire installs the same named statement. Obtain its immutable descriptor once; + // later flows can be created before placement and use it on whichever wire the pool selects. + Command command; + await using (var bootstrap = await bootstrapFactory.CreateAsync().ConfigureAwait(false)) + command = await PrepareAsync(bootstrap).ConfigureAwait(false); + + var pool = new ConnectionPool( + new ProtocolConnectionFactory(protocolFactory), + new ConnectionPoolOptions + { + MinConnections = connectionCount, + MaxConnections = connectionCount, + ConnectionIdleLifetime = Timeout.InfiniteTimeSpan, + }); + return new(pool, command, GetFlowPoolCapacity()); + } + + public async ValueTask ConsumeRetainedAsync( + Func, T> create, + TState state, + Func, ValueTask> consume, + CancellationToken cancellationToken) + { + var flow = RentFlow(); + await _pool.GetAsync( + static (candidate, item) => candidate.Connection.Protocol.TryQueue( + item, + candidate.IsIdleCandidate + ? FlowEnqueueOptions.None + : FlowEnqueueOptions.RequireExistingPipeline, + candidate.CancellationToken), + flow, + Timeout.InfiniteTimeSpan, + cancellationToken).ConfigureAwait(false); + + var values = new List(); + var results = flow.GetAsyncEnumerator(cancellationToken); + try + { + if (await results.MoveNextAsync().ConfigureAwait(false)) + { + results.Current.EnableResultBuffering(); + await results.Current.CollectAsync( + (Values: values, Create: create), + static (collection, row) => collection.Values.Add( + collection.Create( + row.GetInt32(0), row.BorrowFieldMemory(1))), + cancellationToken).ConfigureAwait(false); + } + await consume(state, values).ConfigureAwait(false); + } + finally + { + await results.DisposeAsync().ConfigureAwait(false); + _flowPool?.Return(flow); + } + } + + public ValueTask DisposeAsync() => _pool.DisposeAsync(); + + CommandFlow RentFlow() + { + var flow = _flowPool?.Get(); + return flow is null + ? new CommandFlow(async: true, _options) + : flow.Initialize(async: true, _options); + } + + static int GetFlowPoolCapacity() + { + const string name = "SLON_FLOW_POOL_CAPACITY"; + var value = Environment.GetEnvironmentVariable(name); + if (string.IsNullOrWhiteSpace(value)) + return 0; + return int.TryParse(value, NumberStyles.None, CultureInfo.InvariantCulture, out var capacity) + && capacity >= 0 + ? capacity + : throw new InvalidOperationException($"{name} must be a non-negative integer."); + } + + sealed class CommandFlowPoolPolicy : PooledObjectPolicy + { + public override CommandFlow Create() + => new(async: true, ReadOnlySpan.Empty); + + public override bool Return(CommandFlow flow) + { + // DisposeAsync crosses the framework's retirement boundary before a flow reaches here. + // Reset therefore cannot overlap the old tenure's protocol or heartbeat observation. + flow.Reset(); + return true; + } + } + + static async ValueTask PrepareAsync(PgClientProtocol protocol) + { + var command = Command.Create(Query, commandName: new EncodedCString("fortunes")); + var flow = protocol.Queue(new CommandFlow(async: true, command)); + Command? prepared = null; + await foreach (var result in flow) + { + var metadata = result.GetMetadata(); + prepared = Command.Create(CommandDescriptor.CreatePrepared( + metadata.CommandName, + metadata.ParameterTypes.Preserve(), + metadata.RowDescription?.Preserve())); + await foreach (var _ in result) { } + _ = result.GetCommandComplete(); + } + return prepared ?? + throw new InvalidOperationException("PostgreSQL preparation returned no command result."); + } + + static string RequiredPostgreSqlValue(string name, string? value) + => string.IsNullOrWhiteSpace(value) + ? throw new InvalidOperationException($"PostgreSQL {name} is required.") + : value; + + sealed class ProtocolConnection(PgClientProtocol protocol) + : IPoolConnection + { + IDisposable? _heartbeatRegistration; + + internal PgClientProtocol Protocol { get; } = protocol; + public bool IsIdle => Protocol.Outstanding == 0; + public bool IsSchedulable => Protocol.IsSchedulable; + public Task Completion => Protocol.Completion; + public Task CompleteAsync(Exception? exception = null) + { + var completion = Protocol.CompleteAsync(exception); + if (completion.IsCompleted) + { + StopHeartbeat(); + return completion; + } + return CompleteAndStopHeartbeat(completion, this); + + static async Task CompleteAndStopHeartbeat( + Task completion, ProtocolConnection connection) + { + try + { + await completion.ConfigureAwait(false); + } + finally + { + connection.StopHeartbeat(); + } + } + } + public int CompareTo(ProtocolConnection? other) + => other is null ? 1 : Protocol.Outstanding.CompareTo(other.Protocol.Outstanding); + + public void Start(ConnectionPool.Registration registration) + => Protocol.SetAdmissionAvailableCallback( + () => registration.SignalAvailability(Protocol.Outstanding == 0)); + + internal void StartHeartbeat(ConnectionPoolContext poolContext) + { + Debug.Assert(_heartbeatRegistration is null); + _heartbeatRegistration = poolContext.OnHeartbeat( + static (connection, elapsed) => connection.Protocol.HeartbeatAsync(elapsed), this); + } + + internal void StopHeartbeat() + => Interlocked.Exchange(ref _heartbeatRegistration, null)?.Dispose(); + } + + sealed class ProtocolConnectionFactory(PgClientProtocolFactory factory) + : IPoolConnectionFactory + { + public ProtocolConnection Create( + ConnectionPoolContext poolContext, + TimeSpan timeout = default) + { + var protocol = factory.Create(timeout); + var connection = new ProtocolConnection(protocol); + connection.StartHeartbeat(poolContext); + try + { + _ = PrepareAsync(protocol).AsTask().GetAwaiter().GetResult(); + return connection; + } + catch + { + connection.StopHeartbeat(); + protocol.Dispose(); + throw; + } + } + + public async ValueTask CreateAsync( + ConnectionPoolContext poolContext, + CancellationToken cancellationToken = default) + { + var protocol = await factory.CreateAsync(cancellationToken).ConfigureAwait(false); + var connection = new ProtocolConnection(protocol); + connection.StartHeartbeat(poolContext); + try + { + _ = await PrepareAsync(protocol).ConfigureAwait(false); + return connection; + } + catch + { + connection.StopHeartbeat(); + await protocol.DisposeAsync().ConfigureAwait(false); + throw; + } + } + } +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Minimal/Fortune.cs b/Slon.Benchmarks/Slon.Fortunes.Minimal/Fortune.cs new file mode 100644 index 0000000..ecea74a --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Minimal/Fortune.cs @@ -0,0 +1,17 @@ +namespace Slon.Fortunes.Minimal; + +public readonly struct Fortune : IComparable +{ + public Fortune(int id, ReadOnlyMemory message) + { + Id = id; + Message = message; + } + + public int Id { get; } + + public ReadOnlyMemory Message { get; } + + public int CompareTo(Fortune other) => + Message.Span.SequenceCompareTo(other.Message.Span); +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Minimal/FortuneDatabase.cs b/Slon.Benchmarks/Slon.Fortunes.Minimal/FortuneDatabase.cs new file mode 100644 index 0000000..807709d --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Minimal/FortuneDatabase.cs @@ -0,0 +1,149 @@ +using System.Globalization; +using System.IO.Pipelines; +using System.Text.Encodings.Web; +using Npgsql; +using Slon.Fortunes; + +namespace Slon.Fortunes.Minimal; + +internal abstract class FortuneDatabase : IAsyncDisposable +{ + protected const string Query = "SELECT id, message FROM fortune"; + private static readonly ReadOnlyMemory AdditionalFortune = + "Additional fortune added at request time."u8.ToArray(); + + public abstract ValueTask DisposeAsync(); + + public abstract ValueTask RenderAsync( + PipeWriter writer, + HtmlEncoder htmlEncoder, + CancellationToken cancellationToken); + + public static ValueTask CreateAsync(IConfiguration configuration) + { + ArgumentNullException.ThrowIfNull(configuration); + + var database = RequiredDatabase(configuration["DATABASE"]); + var driver = RequiredDriver(configuration["DRIVER"]); + var connectionString = RequiredConnectionString(configuration["CONNECTION_STRING"]); + var connectionCount = PositiveSetting(configuration, "DATABASE_CONNECTIONS"); + + return (database, driver) switch + { + ("postgresql", "slon") => + SlonFortuneDatabase.CreateAsync(connectionString, connectionCount), + ("postgresql", "npgsql") => ValueTask.FromResult( + new NpgsqlFortuneDatabase(connectionString, connectionCount)), + _ => throw new InvalidOperationException("The database selection is invalid."), + }; + } + + protected static List Complete(List fortunes) + { + fortunes.Add(new Fortune(0, AdditionalFortune)); + fortunes.Sort(); + return fortunes; + } + + protected static async ValueTask RenderFortunesAsync( + List fortunes, + PipeWriter writer, + HtmlEncoder htmlEncoder) + { + using var template = Templates.Fortunes.Create(Complete(fortunes)); + await template.RenderAsync(writer, htmlEncoder); + } + + private static string RequiredDatabase(string? value) + { + var database = RequiredValue("DATABASE", value); + return database == "postgresql" + ? database + : throw new InvalidOperationException("DATABASE must be 'postgresql'."); + } + + private static string RequiredDriver(string? value) + { + var driver = RequiredValue("DRIVER", value); + return driver is "slon" or "npgsql" + ? driver + : throw new InvalidOperationException("DRIVER must be 'slon' or 'npgsql'."); + } + + private static string RequiredConnectionString(string? value) => + string.IsNullOrWhiteSpace(value) + ? throw new InvalidOperationException("CONNECTION_STRING is required.") + : value; + + private static string RequiredValue(string name, string? value) => + string.IsNullOrWhiteSpace(value) + ? throw new InvalidOperationException($"{name} is required.") + : value.Trim().ToLowerInvariant(); + + private static int PositiveSetting(IConfiguration configuration, string name) + { + var value = configuration[name] ?? + throw new InvalidOperationException($"{name} is required."); + + return int.TryParse(value, NumberStyles.None, CultureInfo.InvariantCulture, out var parsed) && + parsed > 0 + ? parsed + : throw new InvalidOperationException($"{name} must be a positive integer."); + } + +} + +internal sealed class SlonFortuneDatabase(SlonConnectionPool pool) : FortuneDatabase +{ + public static async ValueTask CreateAsync( + string connectionString, int connectionCount) + => new SlonFortuneDatabase(await SlonConnectionPool.CreateAsync( + connectionString, connectionCount).ConfigureAwait(false)); + + public override ValueTask RenderAsync( + PipeWriter writer, + HtmlEncoder htmlEncoder, + CancellationToken cancellationToken) + => pool.ConsumeRetainedAsync( + static (id, message) => new Fortune(id, message), + (Writer: writer, HtmlEncoder: htmlEncoder), + static (output, fortunes) => RenderFortunesAsync( + fortunes, output.Writer, output.HtmlEncoder), + cancellationToken); + + public override ValueTask DisposeAsync() => pool.DisposeAsync(); +} + +internal sealed class NpgsqlFortuneDatabase : FortuneDatabase +{ + private readonly NpgsqlDataSource _dataSource; + + public NpgsqlFortuneDatabase(string connectionString, int connectionCount) + { + var builder = new NpgsqlConnectionStringBuilder(connectionString) + { + MaxPoolSize = connectionCount, + }; + _dataSource = new NpgsqlSlimDataSourceBuilder(builder.ConnectionString).Build(); + } + + public override async ValueTask RenderAsync( + PipeWriter writer, + HtmlEncoder htmlEncoder, + CancellationToken cancellationToken) + { + await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken); + await using var command = new NpgsqlCommand(Query, connection); + await using var reader = await command.ExecuteReaderAsync(cancellationToken); + List fortunes = []; + while (await reader.ReadAsync(cancellationToken)) + { + fortunes.Add(new Fortune( + reader.GetInt32(0), reader.GetFieldValue(1))); + } + + await RenderFortunesAsync(fortunes, writer, htmlEncoder); + } + + public override ValueTask DisposeAsync() => _dataSource.DisposeAsync(); +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Minimal/Program.cs b/Slon.Benchmarks/Slon.Fortunes.Minimal/Program.cs new file mode 100644 index 0000000..e4c1b12 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Minimal/Program.cs @@ -0,0 +1,34 @@ +using System.Text.Encodings.Web; +using System.Text.Unicode; +using Slon.Fortunes.Minimal; + +var builder = WebApplication.CreateBuilder(args); + +builder.Logging.ClearProviders(); + +await using var database = await FortuneDatabase.CreateAsync(builder.Configuration); +builder.Services.AddSingleton(CreateHtmlEncoder()); + +await using var app = builder.Build(); + +app.MapGet( + "/fortunes", + async (HttpResponse response, HtmlEncoder htmlEncoder, CancellationToken cancellationToken) => + { + response.ContentType = "text/html; charset=utf-8"; + await database.RenderAsync(response.BodyWriter, htmlEncoder, cancellationToken); + }); + +app.Lifetime.ApplicationStarted.Register(static () => Console.WriteLine("Application started.")); + +await app.RunAsync(); + +static HtmlEncoder CreateHtmlEncoder() +{ + var settings = new TextEncoderSettings( + UnicodeRanges.BasicLatin, + UnicodeRanges.Katakana, + UnicodeRanges.Hiragana); + settings.AllowCharacter('\u2014'); + return HtmlEncoder.Create(settings); +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Minimal/README.md b/Slon.Benchmarks/Slon.Fortunes.Minimal/README.md new file mode 100644 index 0000000..fd2b028 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Minimal/README.md @@ -0,0 +1,38 @@ +# Slon Minimal APIs Fortunes + +This standalone Minimal APIs benchmark exposes `GET /fortunes`. Every request loads all +`fortune` rows, appends the standard request-time fortune, sorts by message, and renders the +HTML response with RazorSlices HTML encoding. + +## Selection + +Set the following configuration values as environment variables or equivalent .NET configuration: + +| Setting | Values | +| --- | --- | +| `DATABASE` | `postgresql` | +| `DRIVER` | `slon` or `npgsql` | +| `CONNECTION_STRING` | PostgreSQL connection string | +| `DATABASE_CONNECTIONS` | Positive fixed pool size | +| `SLON_FLOW_POOL_CAPACITY` | Retained `CommandFlow` count; omitted or `0` disables flow pooling | + +Invalid, unsupported, or missing selections fail application startup with an explicit error. +The Crank config currently composes `sebros/slon-benchmarks` with Draghi's +`experiment/observation-frontier`; override either revision after those experiments land. + +## Driver strategies + +Slon uses its experimental lower layer through `ConnectionPool`. Setting +`SLON_FLOW_POOL_CAPACITY` reuses `CommandFlow` instances after framework retirement; leaving it +unset creates a fresh flow per request. Every wire receives the same prepared statement before it becomes +schedulable. Results are consumed through `CommandResult.CollectAsync`; result retention keeps +borrowed UTF-8 fields valid through Razor rendering without per-row strings or byte arrays. Zero-byte +reads are disabled to match Apex's ordinary BCL transport shape. + +Npgsql uses a slim data source and a command bound to each leased connection. Both drivers append, +ordinally sort, and render the same UTF-8 model; Npgsql returns each field as an allocated byte array +because its reader does not retain row storage through rendering. + +The Crank configuration retains up to 1024 Slon flows, uses two fewer Slon connections than +database cores, and uses 256 Npgsql connections; Npgsql needs the additional in-flight operations +to hide network and query latency. Set `slonFlowPoolCapacity=0` for the unpooled comparison. diff --git a/Slon.Benchmarks/Slon.Fortunes.Minimal/Slon.Fortunes.Minimal.csproj b/Slon.Benchmarks/Slon.Fortunes.Minimal/Slon.Fortunes.Minimal.csproj new file mode 100644 index 0000000..274c2c6 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Minimal/Slon.Fortunes.Minimal.csproj @@ -0,0 +1,21 @@ + + + + net10.0 + enable + enable + false + $(NoWarn);SLONPG001;SLONPOOL001 + + + + + + + + + + + + + diff --git a/Slon.Benchmarks/Slon.Fortunes.Minimal/Templates/Fortunes.cshtml b/Slon.Benchmarks/Slon.Fortunes.Minimal/Templates/Fortunes.cshtml new file mode 100644 index 0000000..b36d18a --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Minimal/Templates/Fortunes.cshtml @@ -0,0 +1,2 @@ +@inherits RazorSlice> +Fortunes@foreach (var item in Model){}
idmessage
@WriteNumber(item.Id, default, CultureInfo.InvariantCulture, false)@item.Message.Span
diff --git a/Slon.Benchmarks/Slon.Fortunes.Minimal/Templates/_ViewImports.cshtml b/Slon.Benchmarks/Slon.Fortunes.Minimal/Templates/_ViewImports.cshtml new file mode 100644 index 0000000..2fb29eb --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Minimal/Templates/_ViewImports.cshtml @@ -0,0 +1,9 @@ +@inherits RazorSlice + +@using System.Globalization +@using Microsoft.AspNetCore.Razor +@using RazorSlices +@using Slon.Fortunes.Minimal + +@tagHelperPrefix __disable_tagHelpers__: +@removeTagHelper *, Microsoft.AspNetCore.Mvc.Razor diff --git a/Slon.Benchmarks/Slon.Fortunes.Minimal/minimal-fortunes.benchmarks.yml b/Slon.Benchmarks/Slon.Fortunes.Minimal/minimal-fortunes.benchmarks.yml new file mode 100644 index 0000000..bbcdea2 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Minimal/minimal-fortunes.benchmarks.yml @@ -0,0 +1,91 @@ +imports: + - https://raw.githubusercontent.com/dotnet/crank/main/src/Microsoft.Crank.Jobs.Wrk/wrk.yml + - https://raw.githubusercontent.com/aspnet/Benchmarks/main/scenarios/aspnet.profiles.standard.yml + +variables: + serverPort: 5000 + npgsqlConnections: 256 + slonFlowPoolCapacity: 1024 + branchOrCommit: sebros/slon-benchmarks + draghiBranchOrCommit: experiment/observation-frontier + +jobs: + minimal-postgresql-slon: + sources: + Slon: + repository: https://github.com/draghidev/slon.git + branchOrCommit: "{{branchOrCommit}}" + Draghi: + repository: https://github.com/draghidev/pipelining.git + branchOrCommit: "{{draghiBranchOrCommit}}" + project: Slon/Slon.Benchmarks/Slon.Fortunes.Minimal/Slon.Fortunes.Minimal.csproj + sdkVersion: 11.0.100-preview.7.26381.103 + framework: net10.0 + patchReferences: true + readyStateText: Application started. + arguments: "--urls {{serverScheme}}://{{serverAddress}}:{{serverPort}}" + variables: + serverScheme: http + environmentVariables: + DATABASE: postgresql + DRIVER: slon + CONNECTION_STRING: "Host={{databaseServer}};Database=hello_world;Username=benchmarkdbuser;Password=benchmarkdbpass" + DATABASE_CONNECTIONS: "{{ cores | minus: 2 }}" + SLON_FLOW_POOL_CAPACITY: "{{slonFlowPoolCapacity}}" + + minimal-postgresql-npgsql: + sources: + Slon: + repository: https://github.com/draghidev/slon.git + branchOrCommit: "{{branchOrCommit}}" + Draghi: + repository: https://github.com/draghidev/pipelining.git + branchOrCommit: "{{draghiBranchOrCommit}}" + project: Slon/Slon.Benchmarks/Slon.Fortunes.Minimal/Slon.Fortunes.Minimal.csproj + sdkVersion: 11.0.100-preview.7.26381.103 + framework: net10.0 + patchReferences: true + readyStateText: Application started. + arguments: "--urls {{serverScheme}}://{{serverAddress}}:{{serverPort}}" + variables: + serverScheme: http + environmentVariables: + DATABASE: postgresql + DRIVER: npgsql + CONNECTION_STRING: "Host={{databaseServer}};Database=hello_world;Username=benchmarkdbuser;Password=benchmarkdbpass;Maximum Pool Size={{ npgsqlConnections }};No Reset On Close=true;Enlist=false;Max Auto Prepare=4" + DATABASE_CONNECTIONS: "{{ npgsqlConnections }}" + + postgresql: + source: + repository: https://github.com/TechEmpower/FrameworkBenchmarks.git + branchOrCommit: master + dockerFile: toolset/databases/postgres/postgres.dockerfile + dockerImageName: postgres_te + dockerContextDirectory: toolset/databases/postgres + readyStateText: ready to accept connections + noClean: true + +scenarios: + minimal-postgresql-slon: + db: + job: postgresql + application: + job: minimal-postgresql-slon + load: + job: wrk + variables: + presetHeaders: html + path: /fortunes + connections: 512 + + minimal-postgresql-npgsql: + db: + job: postgresql + application: + job: minimal-postgresql-npgsql + load: + job: wrk + variables: + presetHeaders: html + path: /fortunes + connections: 512 diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/BenchmarkApplication.HttpConnection.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/BenchmarkApplication.HttpConnection.cs new file mode 100644 index 0000000..001fa74 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/BenchmarkApplication.HttpConnection.cs @@ -0,0 +1,163 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. + +using System.Buffers; +using System.IO.Pipelines; +using System.Runtime.CompilerServices; +using System.Text.Encodings.Web; +using System.Text.Unicode; +using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http; + +namespace Slon.Fortunes.Platform; + +public sealed partial class BenchmarkApplication : IHttpConnection +{ + private State _state; + + public CancellationToken ConnectionClosed { get; set; } + + public PipeReader Reader { get; set; } = null!; + + public PipeWriter Writer { get; set; } = null!; + + private HtmlEncoder HtmlEncoder { get; } = CreateHtmlEncoder(); + + private HttpParser Parser { get; } = new(); + + public async Task ExecuteAsync() + { + try + { + await ProcessRequestsAsync(); + Reader.Complete(); + } + catch (Exception ex) + { + Reader.Complete(ex); + } + finally + { + Writer.Complete(); + } + } + + private async Task ProcessRequestsAsync() + { + while (true) + { + var readResult = await Reader.ReadAsync(); + var buffer = readResult.Buffer; + var isCompleted = readResult.IsCompleted; + + if (buffer.IsEmpty && isCompleted) + { + return; + } + + while (true) + { + ParseHttpRequest(ref buffer, isCompleted); + + if (_state == State.Body) + { + await ProcessRequestAsync(); + _state = State.StartLine; + + if (!buffer.IsEmpty) + { + continue; + } + } + + Reader.AdvanceTo(buffer.Start, buffer.End); + break; + } + + await Writer.FlushAsync(); + } + } + + private void ParseHttpRequest(ref ReadOnlySequence buffer, bool isCompleted) + { + var reader = new SequenceReader(buffer); + var state = _state; + if (state == State.StartLine && + Parser.ParseRequestLine(new ParsingAdapter(this), ref reader)) + { + state = State.Headers; + } + + if (state == State.Headers && + Parser.ParseHeaders(new ParsingAdapter(this), ref reader)) + { + state = State.Body; + } + + if (state != State.Body && isCompleted) + { + throw new InvalidOperationException("Unexpected end of data!"); + } + + _state = state; + buffer = state == State.Body + ? buffer.Slice(reader.Position, 0) + : buffer.Slice(reader.Position); + } + + private static HtmlEncoder CreateHtmlEncoder() + { + var settings = new TextEncoderSettings( + UnicodeRanges.BasicLatin, + UnicodeRanges.Katakana, + UnicodeRanges.Hiragana); + settings.AllowCharacter('\u2014'); + return HtmlEncoder.Create(settings); + } + + public void OnStaticIndexedHeader(int index) + { + } + + public void OnStaticIndexedHeader(int index, ReadOnlySpan value) + { + } + + public void OnHeader(ReadOnlySpan name, ReadOnlySpan value) + { + } + + public void OnHeadersComplete(bool endStream) + { + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static void ReturnChunkedWriter(ChunkedPipeWriter writer) => ChunkedWriterPool.Return(writer); + + private enum State + { + StartLine, + Headers, + Body, + } + + private readonly struct ParsingAdapter(BenchmarkApplication requestHandler) : + IHttpRequestLineHandler, + IHttpHeadersHandler + { + public void OnStaticIndexedHeader(int index) => requestHandler.OnStaticIndexedHeader(index); + + public void OnStaticIndexedHeader(int index, ReadOnlySpan value) => + requestHandler.OnStaticIndexedHeader(index, value); + + public void OnHeader(ReadOnlySpan name, ReadOnlySpan value) => + requestHandler.OnHeader(name, value); + + public void OnHeadersComplete(bool endStream) => requestHandler.OnHeadersComplete(endStream); + + public void OnStartLine( + HttpVersionAndMethod versionAndMethod, + TargetOffsetPathLength targetPath, + Span startLine) => + requestHandler.OnStartLine(versionAndMethod, targetPath, startLine); + } +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/BenchmarkApplication.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/BenchmarkApplication.cs new file mode 100644 index 0000000..91cfb5d --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/BenchmarkApplication.cs @@ -0,0 +1,134 @@ +using System.IO.Pipelines; +using System.Runtime.CompilerServices; +using System.Text.Encodings.Web; +using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http; +using Microsoft.Extensions.ObjectPool; +using RazorSlices; + +namespace Slon.Fortunes.Platform; + +internal enum FortuneTemplating +{ + Razor, + Raw +} + +public sealed partial class BenchmarkApplication +{ + private static readonly DefaultObjectPool ChunkedWriterPool = + new(new ChunkedWriterObjectPolicy()); + + private RequestType _requestType; + + internal static FortuneDatabase Database { get; set; } = null!; + internal static FortuneTemplating Templating { get; set; } + + public void OnStartLine( + HttpVersionAndMethod versionAndMethod, + TargetOffsetPathLength targetPath, + Span startLine) + { + _requestType = versionAndMethod.Method == + Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http.HttpMethod.Get && + startLine.Slice(targetPath.Offset, targetPath.Length).SequenceEqual("/fortunes"u8) + ? RequestType.Fortunes + : RequestType.NotRecognized; + } + + private ValueTask ProcessRequestAsync() => _requestType switch + { + RequestType.Fortunes => Database.RenderAsync(this, default), + _ => OutputEmptyAsync(Writer), + }; + + internal ValueTask RenderFortunesAsync(List fortunes) + { + if (Templating is FortuneTemplating.Raw) + { + var writer = StartResponse(Writer); + RawFortuneTemplating.Render(fortunes, writer, HtmlEncoder); + writer.Complete(); + ReturnChunkedWriter(writer); + return ValueTask.CompletedTask; + } + + var template = Templates.Fortunes.Create(fortunes); + return OutputFortunesAsync(Writer, template); + } + + private ValueTask OutputFortunesAsync( + PipeWriter pipeWriter, + RazorSlice template) + { + var chunkedWriter = StartResponse(pipeWriter); + var renderTask = template.RenderAsync(chunkedWriter, HtmlEncoder); + if (renderTask.IsCompletedSuccessfully) + { + renderTask.GetAwaiter().GetResult(); + EndTemplateRendering(chunkedWriter, template); + return ValueTask.CompletedTask; + } + + return AwaitTemplateRenderTask(renderTask, chunkedWriter, template); + } + + private static ValueTask OutputEmptyAsync(PipeWriter pipeWriter) + { + var writer = StartResponse(pipeWriter); + writer.Complete(); + ReturnChunkedWriter(writer); + return ValueTask.CompletedTask; + } + + private static ChunkedPipeWriter StartResponse(PipeWriter pipeWriter) + { + var preamble = + "HTTP/1.1 200 OK\r\nServer: K\r\nContent-Type: text/html; charset=utf-8\r\nTransfer-Encoding: chunked"u8; + var headersLength = preamble.Length + DateHeader.HeaderBytes.Length; + var headersSpan = pipeWriter.GetSpan(headersLength); + preamble.CopyTo(headersSpan); + DateHeader.HeaderBytes.CopyTo(headersSpan[preamble.Length..]); + pipeWriter.Advance(headersLength); + + var writer = ChunkedWriterPool.Get(); + writer.SetOutput(pipeWriter, 2048); + return writer; + } + + private static async ValueTask AwaitTemplateRenderTask( + ValueTask renderTask, + ChunkedPipeWriter chunkedWriter, + RazorSlice template) + { + await renderTask; + EndTemplateRendering(chunkedWriter, template); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static void EndTemplateRendering( + ChunkedPipeWriter chunkedWriter, + RazorSlice template) + { + chunkedWriter.Complete(); + ReturnChunkedWriter(chunkedWriter); + template.Dispose(); + } + + private sealed class ChunkedWriterObjectPolicy : + IPooledObjectPolicy + { + public ChunkedPipeWriter Create() => new(); + + public bool Return(ChunkedPipeWriter writer) + { + writer.Reset(); + return true; + } + } + + private enum RequestType + { + NotRecognized, + Fortunes, + } +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/BufferExtensions.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/BufferExtensions.cs new file mode 100644 index 0000000..55ba425 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/BufferExtensions.cs @@ -0,0 +1,56 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. + +using System.Buffers; +using System.Runtime.CompilerServices; +using System.Text; + +namespace Slon.Fortunes.Platform; + +public static class BufferExtensions +{ + private const int MaxULongByteLength = 20; + + [ThreadStatic] + private static byte[]? s_numericBytesScratch; + + internal static void WriteUtf8String(ref this BufferWriter buffer, string text) + where T : struct, IBufferWriter + { + var byteCount = Encoding.UTF8.GetByteCount(text); + buffer.Ensure(byteCount); + byteCount = Encoding.UTF8.GetBytes(text.AsSpan(), buffer.Span); + buffer.Advance(byteCount); + } + + [MethodImpl(MethodImplOptions.NoInlining)] + internal static void WriteNumericMultiWrite(ref this BufferWriter buffer, uint number) + where T : IBufferWriter + { + const byte AsciiDigitStart = (byte)'0'; + + var value = number; + var position = MaxULongByteLength; + var byteBuffer = NumericBytesScratch; + do + { + var quotient = value / 10; + byteBuffer[--position] = (byte)(AsciiDigitStart + (value - quotient * 10)); + value = quotient; + } + while (value != 0); + + var length = MaxULongByteLength - position; + buffer.Write(new ReadOnlySpan(byteBuffer, position, length)); + } + + private static byte[] NumericBytesScratch => s_numericBytesScratch ?? CreateNumericBytesScratch(); + + [MethodImpl(MethodImplOptions.NoInlining)] + private static byte[] CreateNumericBytesScratch() + { + var bytes = new byte[MaxULongByteLength]; + s_numericBytesScratch = bytes; + return bytes; + } +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/BufferWriter.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/BufferWriter.cs new file mode 100644 index 0000000..18d2ca8 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/BufferWriter.cs @@ -0,0 +1,133 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. + +using System.Buffers; +using System.Runtime.CompilerServices; + +namespace Slon.Fortunes.Platform; + +public ref struct BufferWriter + where T : IBufferWriter +{ + private T _output; + private Span _span; + private int _buffered; + + public BufferWriter(T output, int sizeHint) + { + _buffered = 0; + _output = output; + _span = output.GetSpan(sizeHint); + } + + public Span Span => _span; + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void Commit() + { + var buffered = _buffered; + if (buffered > 0) + { + _buffered = 0; + _output.Advance(buffered); + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void Advance(int count) + { + _buffered += count; + _span = _span[count..]; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void Write(scoped ReadOnlySpan source) + { + if (_span.Length >= source.Length) + { + source.CopyTo(_span); + Advance(source.Length); + } + else + { + WriteMultiBuffer(source); + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void Ensure(int count = 1) + { + if (_span.Length < count) + { + EnsureMore(count); + } + } + + [MethodImpl(MethodImplOptions.NoInlining)] + private void EnsureMore(int count = 0) + { + if (_buffered > 0) + { + Commit(); + } + + _span = _output.GetSpan(count); + } + + private void WriteMultiBuffer(scoped ReadOnlySpan source) + { + while (source.Length > 0) + { + if (_span.Length == 0) + { + EnsureMore(); + } + + var writable = Math.Min(source.Length, _span.Length); + source[..writable].CopyTo(_span); + source = source[writable..]; + Advance(writable); + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal void WriteNumeric(uint number) + { + const byte AsciiDigitStart = (byte)'0'; + var span = Span; + var advanceBy = 0; + if (span.Length >= 3) + { + if (number < 10) + { + span[0] = (byte)(number + AsciiDigitStart); + advanceBy = 1; + } + else if (number < 100) + { + var tens = (byte)((number * 205u) >> 11); + span[0] = (byte)(tens + AsciiDigitStart); + span[1] = (byte)(number - (tens * 10) + AsciiDigitStart); + advanceBy = 2; + } + else if (number < 1000) + { + var digit0 = (byte)((number * 41u) >> 12); + var digits01 = (byte)((number * 205u) >> 11); + span[0] = (byte)(digit0 + AsciiDigitStart); + span[1] = (byte)(digits01 - (digit0 * 10) + AsciiDigitStart); + span[2] = (byte)(number - (digits01 * 10) + AsciiDigitStart); + advanceBy = 3; + } + } + + if (advanceBy > 0) + { + Advance(advanceBy); + } + else + { + BufferExtensions.WriteNumericMultiWrite(ref this, number); + } + } +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/ChunkedPipeWriter.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/ChunkedPipeWriter.cs new file mode 100644 index 0000000..411c849 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/ChunkedPipeWriter.cs @@ -0,0 +1,235 @@ +using System.Buffers; +using System.Buffers.Text; +using System.Diagnostics; +using System.IO.Pipelines; +using System.Numerics; +using System.Runtime.CompilerServices; + +namespace Slon.Fortunes.Platform; + +internal sealed class ChunkedPipeWriter : PipeWriter +{ + private const int DefaultChunkSizeHint = 2048; + private static readonly StandardFormat DefaultHexFormat = GetHexFormat(DefaultChunkSizeHint); + private static ReadOnlySpan ChunkTerminator => "\r\n"u8; + + private PipeWriter _output = null!; + private int _chunkSizeHint; + private StandardFormat _hexFormat = DefaultHexFormat; + private Memory _currentFullChunk; + private Memory _currentChunk; + private int _buffered; + private long _unflushedBytes; + private bool _ended; + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void SetOutput(PipeWriter output, int chunkSizeHint = DefaultChunkSizeHint) + { + _buffered = 0; + _unflushedBytes = 0; + _chunkSizeHint = chunkSizeHint; + _output = output; + StartNewChunk(chunkSizeHint, isFirst: true); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void Reset() + { + _buffered = 0; + _unflushedBytes = 0; + _output = null!; + _ended = false; + _hexFormat = DefaultHexFormat; + _currentFullChunk = default; + _currentChunk = default; + } + + public override bool CanGetUnflushedBytes => true; + + public override long UnflushedBytes => _unflushedBytes; + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public override void Advance(int count) + { + ThrowIfEnded(); + _buffered += count; + _unflushedBytes += count; + _currentChunk = _currentChunk[count..]; + } + + public override Memory GetMemory(int sizeHint = 0) + { + ThrowIfEnded(); + if (_currentChunk.Length <= sizeHint) + { + EnsureMore(sizeHint); + } + + return _currentChunk; + } + + public override Span GetSpan(int sizeHint = 0) => GetMemory(sizeHint).Span; + + public override void CancelPendingFlush() => _output.CancelPendingFlush(); + + public override void Complete(Exception? exception = null) + { + ThrowIfEnded(); + CommitCurrentChunk(isFinal: true); + _ended = true; + } + + public override ValueTask FlushAsync(CancellationToken cancellationToken = default) + { + CommitCurrentChunk(isFinal: false); + var flushTask = _output.FlushAsync(cancellationToken); + _unflushedBytes = 0; + return flushTask; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static StandardFormat GetHexFormat(int maxValue) + { + var hexDigitCount = CountHexDigits(maxValue); + return new StandardFormat('X', (byte)hexDigitCount); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private static int CountHexDigits(int n) => + n <= 16 ? 1 : (BitOperations.Log2((uint)n) >> 2) + 1; + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private void StartNewChunk(int sizeHint, bool isFirst = false) + { + ThrowIfEnded(); + + var oldFullChunkHexLength = -1; + if (!isFirst) + { + oldFullChunkHexLength = CountHexDigits(_currentFullChunk.Length); + } + + _currentFullChunk = _output.GetMemory(Math.Max(_chunkSizeHint, sizeHint)); + var newFullChunkHexLength = CountHexDigits(_currentFullChunk.Length); + var currentFullChunkSpan = _currentFullChunk.Span; + currentFullChunkSpan[..newFullChunkHexLength].Fill((byte)'0'); + "\r\n"u8.CopyTo(currentFullChunkSpan[newFullChunkHexLength..]); + var chunkHeaderLength = newFullChunkHexLength + 2; + _currentChunk = _currentFullChunk[chunkHeaderLength..]; + + if ((!isFirst && oldFullChunkHexLength != newFullChunkHexLength) || + (isFirst && DefaultChunkSizeHint != _chunkSizeHint)) + { + _hexFormat = GetHexFormat(_currentFullChunk.Length); + } + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + private void CommitCurrentChunk(bool isFinal = false, int sizeHint = 0) + { + ThrowIfEnded(); + + var contentLength = _buffered; + if (contentLength <= 0) + { + if (isFinal) + { + var terminator = "0\r\n\r\n"u8; + terminator.CopyTo(_currentFullChunk.Span); + _output.Advance(terminator.Length); + } + + return; + } + + var chunkLengthHexDigitsLength = CountHexDigits(contentLength); + var span = _currentFullChunk.Span; + if (!Utf8Formatter.TryFormat(contentLength, span, out var bytesWritten, _hexFormat)) + { + throw new NotSupportedException("Chunk size too large"); + } + + Debug.Assert(chunkLengthHexDigitsLength == bytesWritten, "HEX formatting math problem."); + var spanOffset = chunkLengthHexDigitsLength + 2 + contentLength; + var chunkTotalLength = spanOffset + ChunkTerminator.Length; + Debug.Assert(span.Length >= chunkTotalLength, "Bad chunk size calculation."); + ChunkTerminator.CopyTo(span[spanOffset..]); + + if (!isFinal) + { + _output.Advance(chunkTotalLength); + StartNewChunk(sizeHint); + } + else + { + var terminator = "0\r\n\r\n"u8; + if (chunkTotalLength + terminator.Length <= span.Length) + { + terminator.CopyTo(span[chunkTotalLength..]); + _output.Advance(chunkTotalLength + terminator.Length); + } + else + { + _output.Advance(chunkTotalLength); + _output.Write(terminator); + } + } + + _buffered = 0; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void Write(ReadOnlySpan source) + { + ThrowIfEnded(); + if (_currentChunk.Length >= source.Length + ChunkTerminator.Length) + { + source.CopyTo(_currentChunk.Span); + Advance(source.Length); + } + else + { + WriteMultiBuffer(source); + } + } + + [MethodImpl(MethodImplOptions.NoInlining)] + private void EnsureMore(int count = 0) + { + if (count > _currentChunk.Length - _buffered - ChunkTerminator.Length) + { + if (_buffered > 0) + { + CommitCurrentChunk(isFinal: false, count); + } + else + { + StartNewChunk(count); + } + } + } + + private void WriteMultiBuffer(ReadOnlySpan source) + { + while (source.Length > 0) + { + if (_currentChunk.Length - ChunkTerminator.Length == 0) + { + EnsureMore(); + } + + var writable = Math.Min(source.Length, _currentChunk.Length - ChunkTerminator.Length); + source[..writable].CopyTo(_currentChunk.Span); + source = source[writable..]; + Advance(writable); + } + } + + private void ThrowIfEnded() + { + if (_ended) + { + throw new InvalidOperationException("Cannot use the writer after calling End()."); + } + } +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/DateHeader.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/DateHeader.cs new file mode 100644 index 0000000..0782086 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/DateHeader.cs @@ -0,0 +1,53 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. + +using System.Buffers.Text; +using System.Diagnostics; + +namespace Slon.Fortunes.Platform; + +internal static class DateHeader +{ + private const int PrefixLength = 8; + private const int DateTimeRLength = 29; + private const int SuffixLength = 2; + private const int SuffixIndex = DateTimeRLength + PrefixLength; + + private static readonly Timer s_timer = new(static _ => SetDateValues(DateTimeOffset.UtcNow), null, 1000, 1000); + private static byte[] s_headerBytesMaster = new byte[PrefixLength + DateTimeRLength + 2 * SuffixLength]; + private static byte[] s_headerBytesScratch = new byte[PrefixLength + DateTimeRLength + 2 * SuffixLength]; + + static DateHeader() + { + "\r\nDate: "u8.CopyTo(s_headerBytesMaster); + "\r\nDate: "u8.CopyTo(s_headerBytesScratch); + s_headerBytesMaster[SuffixIndex] = (byte)'\r'; + s_headerBytesMaster[SuffixIndex + 1] = (byte)'\n'; + s_headerBytesMaster[SuffixIndex + 2] = (byte)'\r'; + s_headerBytesMaster[SuffixIndex + 3] = (byte)'\n'; + s_headerBytesScratch[SuffixIndex] = (byte)'\r'; + s_headerBytesScratch[SuffixIndex + 1] = (byte)'\n'; + s_headerBytesScratch[SuffixIndex + 2] = (byte)'\r'; + s_headerBytesScratch[SuffixIndex + 3] = (byte)'\n'; + SetDateValues(DateTimeOffset.UtcNow); + SyncDateTimer(); + } + + public static void SyncDateTimer() => s_timer.Change(1000, 1000); + + public static ReadOnlySpan HeaderBytes => s_headerBytesMaster; + + private static void SetDateValues(DateTimeOffset value) + { + lock (s_headerBytesScratch) + { + if (!Utf8Formatter.TryFormat(value, s_headerBytesScratch.AsSpan(PrefixLength), out var written, 'R')) + { + throw new Exception("date time format failed"); + } + + Debug.Assert(written == DateTimeRLength); + (s_headerBytesScratch, s_headerBytesMaster) = (s_headerBytesMaster, s_headerBytesScratch); + } + } +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/Fortune.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/Fortune.cs new file mode 100644 index 0000000..2661ca9 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/Fortune.cs @@ -0,0 +1,17 @@ +namespace Slon.Fortunes.Platform; + +public readonly struct Fortune : IComparable +{ + public Fortune(int id, ReadOnlyMemory message) + { + Id = id; + Message = message; + } + + public int Id { get; } + + public ReadOnlyMemory Message { get; } + + public int CompareTo(Fortune other) => + Message.Span.SequenceCompareTo(other.Message.Span); +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/FortuneDatabase.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/FortuneDatabase.cs new file mode 100644 index 0000000..1a8d43a --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/FortuneDatabase.cs @@ -0,0 +1,128 @@ +using System.Globalization; +using System.Text; +using Npgsql; +using Slon.Fortunes; + +namespace Slon.Fortunes.Platform; + +internal abstract class FortuneDatabase : IAsyncDisposable +{ + internal const string Query = "SELECT id, message FROM fortune"; + private static readonly ReadOnlyMemory AdditionalFortune = + "Additional fortune added at request time."u8.ToArray(); + + public abstract ValueTask DisposeAsync(); + + public abstract ValueTask RenderAsync( + BenchmarkApplication application, + CancellationToken cancellationToken); + + public static ValueTask CreateAsync( + string? database, + string? driver, + string? connectionString) + { + var selectedDatabase = RequiredSelection("DATABASE", database); + var selectedDriver = RequiredSelection("DRIVER", driver); + var requiredConnectionString = string.IsNullOrWhiteSpace(connectionString) + ? throw new InvalidOperationException("CONNECTION_STRING is required.") + : connectionString; + var connectionCount = PositiveEnvironment("DATABASE_CONNECTIONS"); + + return (selectedDatabase, selectedDriver) switch + { + ("postgresql", "slon") => + SlonFortuneDatabase.CreateAsync(requiredConnectionString, connectionCount), + ("postgresql", "npgsql") => + ValueTask.FromResult( + new NpgsqlFortuneDatabase(requiredConnectionString, connectionCount)), + ("postgresql", _) => + throw new InvalidOperationException( + $"DRIVER '{selectedDriver}' is not valid for DATABASE '{selectedDatabase}'."), + _ => throw new ArgumentOutOfRangeException( + nameof(database), + selectedDatabase, + "DATABASE must be 'postgresql'."), + }; + } + + protected static List Complete(List fortunes) + { + fortunes.Add(new Fortune(0, AdditionalFortune)); + fortunes.Sort(); + return fortunes; + } + + protected static int PositiveEnvironment(string name) + { + var value = Environment.GetEnvironmentVariable(name) ?? + throw new InvalidOperationException($"{name} is required."); + + return int.TryParse(value, NumberStyles.None, CultureInfo.InvariantCulture, out var parsed) && + parsed > 0 + ? parsed + : throw new ArgumentOutOfRangeException( + name, + value, + "Value must be a positive integer."); + } + + private static string RequiredSelection(string name, string? value) => + string.IsNullOrWhiteSpace(value) + ? throw new InvalidOperationException($"{name} is required.") + : value.Trim().ToLowerInvariant(); + +} + +internal sealed class SlonFortuneDatabase(SlonConnectionPool pool) : FortuneDatabase +{ + public static async ValueTask CreateAsync( + string connectionString, int connectionCount) + => new SlonFortuneDatabase(await SlonConnectionPool.CreateAsync( + connectionString, connectionCount).ConfigureAwait(false)); + + public override ValueTask RenderAsync( + BenchmarkApplication application, + CancellationToken cancellationToken) + => pool.ConsumeRetainedAsync( + static (id, message) => new Fortune(id, message), + application, + static (application, fortunes) => + application.RenderFortunesAsync(Complete(fortunes)), + cancellationToken); + + public override ValueTask DisposeAsync() => pool.DisposeAsync(); +} + +internal sealed class NpgsqlFortuneDatabase : FortuneDatabase +{ + private readonly NpgsqlDataSource _dataSource; + + public NpgsqlFortuneDatabase(string connectionString, int connectionCount) + { + var builder = new NpgsqlConnectionStringBuilder(connectionString) + { + MaxPoolSize = connectionCount, + }; + _dataSource = new NpgsqlSlimDataSourceBuilder(builder.ConnectionString).Build(); + } + + public override async ValueTask RenderAsync( + BenchmarkApplication application, + CancellationToken cancellationToken) + { + await using var connection = await _dataSource.OpenConnectionAsync(cancellationToken); + await using var command = new NpgsqlCommand(Query, connection); + await using var reader = await command.ExecuteReaderAsync(cancellationToken); + var fortunes = new List(); + while (await reader.ReadAsync(cancellationToken)) + { + fortunes.Add(new Fortune( + reader.GetInt32(0), Encoding.UTF8.GetBytes(reader.GetString(1)))); + } + + await application.RenderFortunesAsync(Complete(fortunes)); + } + + public override ValueTask DisposeAsync() => _dataSource.DisposeAsync(); +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/HttpApplication.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/HttpApplication.cs new file mode 100644 index 0000000..54666d4 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/HttpApplication.cs @@ -0,0 +1,28 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. + +using Microsoft.AspNetCore.Connections; + +namespace Slon.Fortunes.Platform; + +public static class HttpApplicationConnectionBuilderExtensions +{ + public static IConnectionBuilder UseHttpApplication(this IConnectionBuilder builder) + where TConnection : IHttpConnection, new() => + builder.Use(_ => new HttpApplication().ExecuteAsync); +} + +public sealed class HttpApplication + where TConnection : IHttpConnection, new() +{ + public Task ExecuteAsync(ConnectionContext connection) + { + var httpConnection = new TConnection + { + ConnectionClosed = connection.ConnectionClosed, + Reader = connection.Transport.Input, + Writer = connection.Transport.Output, + }; + return httpConnection.ExecuteAsync(); + } +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/IHttpConnection.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/IHttpConnection.cs new file mode 100644 index 0000000..ce144cd --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/IHttpConnection.cs @@ -0,0 +1,18 @@ +// Copyright (c) .NET Foundation. All rights reserved. +// Licensed under the Apache License, Version 2.0. + +using System.IO.Pipelines; +using Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Http; + +namespace Slon.Fortunes.Platform; + +public interface IHttpConnection : IHttpHeadersHandler, IHttpRequestLineHandler +{ + CancellationToken ConnectionClosed { get; set; } + + PipeReader Reader { get; set; } + + PipeWriter Writer { get; set; } + + Task ExecuteAsync(); +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/Program.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/Program.cs new file mode 100644 index 0000000..a115633 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/Program.cs @@ -0,0 +1,56 @@ +using System.Net; +using System.Runtime.InteropServices; +using Microsoft.AspNetCore.Hosting; +using Microsoft.Extensions.Configuration; +using Microsoft.Extensions.Hosting; +using Slon.Fortunes.Platform; + +var configuration = new ConfigurationBuilder() + .AddEnvironmentVariables() + .AddCommandLine(args) + .Build(); +var url = new Uri(configuration["urls"] ?? "http://0.0.0.0:5000"); +BenchmarkApplication.Templating = configuration["TEMPLATING"]?.Trim().ToLowerInvariant() switch +{ + null or "" or "razor" => FortuneTemplating.Razor, + "raw" => FortuneTemplating.Raw, + var value => throw new ArgumentOutOfRangeException( + "TEMPLATING", value, "Expected 'razor' or 'raw'."), +}; + +await using var database = await FortuneDatabase.CreateAsync( + configuration["DATABASE"], + configuration["DRIVER"], + configuration["CONNECTION_STRING"]); +BenchmarkApplication.Database = database; +DateHeader.SyncDateTimer(); + +var hostBuilder = Host.CreateDefaultBuilder() + .ConfigureWebHost(webHost => + { + webHost + .UseConfiguration(configuration) + .UseKestrel(options => + { + options.Listen(IPAddress.Any, url.Port, listen => + { + listen.UseHttpApplication(); + }); + }) + .Configure(_ => { }) + .UseSockets(options => + { + options.WaitForDataBeforeAllocatingBuffer = false; + if (RuntimeInformation.IsOSPlatform(OSPlatform.Linux)) + { + options.UnsafePreferInlineScheduling = + Environment.GetEnvironmentVariable( + "DOTNET_SYSTEM_NET_SOCKETS_INLINE_COMPLETIONS") == "1"; + } + }); + }); + +using var host = hostBuilder.Build(); +await host.StartAsync(); +Console.WriteLine("Application started."); +await host.WaitForShutdownAsync(); diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/README.md b/Slon.Benchmarks/Slon.Fortunes.Platform/README.md new file mode 100644 index 0000000..f33923c --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/README.md @@ -0,0 +1,42 @@ +# Slon Platform Fortunes + +This standalone Platform-style Fortunes application exposes `GET /fortunes`. It reads every +row from `fortune`, adds the standard request-time fortune, sorts by message, and renders the +standard HTML response with RazorSlices HTML encoding. + +## Selection + +Set all of these environment variables before starting the app: + +| Variable | Values | +| --- | --- | +| `DATABASE` | `postgresql` | +| `DRIVER` | `slon` or `npgsql` | +| `CONNECTION_STRING` | PostgreSQL connection string | +| `DATABASE_CONNECTIONS` | Positive fixed pool size | +| `SLON_FLOW_POOL_CAPACITY` | Retained `CommandFlow` count; omitted or `0` disables flow pooling | +| `TEMPLATING` | `razor` (default) or `raw` | + +Invalid, unsupported, or missing selections fail application startup with an explicit error. +The Crank config currently composes `sebros/slon-benchmarks` with Draghi's +`experiment/observation-frontier`; override either revision after those experiments land. + +## Driver strategies + +Slon uses its experimental lower layer through `ConnectionPool`. Setting +`SLON_FLOW_POOL_CAPACITY` reuses `CommandFlow` instances after framework retirement; leaving it +unset creates a fresh flow per request. Every wire receives the same prepared statement before it becomes +schedulable. `CommandResult.CollectAsync` is the row-buffering barrier, while result buffering retains +UTF-8 field memory through rendering without per-row strings or byte arrays. Zero-byte reads are +disabled to match Apex's ordinary BCL transport shape. + +`TEMPLATING=raw` writes the same encoded HTML directly into the response buffer. It isolates the +driver and pool cost from RazorSlices overhead without changing query or row-consumption behavior. + +Npgsql uses a slim data source and a command bound to each leased connection. Every strategy +appends and ordinally sorts the same logical model and renders through the same RazorSlices UTF-8 +template. + +The Crank configuration retains up to 1024 Slon flows, uses two fewer Slon connections than +database cores, and uses 256 Npgsql connections. Set `slonFlowPoolCapacity=0` for the unpooled +comparison. diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/RawFortuneTemplating.cs b/Slon.Benchmarks/Slon.Fortunes.Platform/RawFortuneTemplating.cs new file mode 100644 index 0000000..6f6fed8 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/RawFortuneTemplating.cs @@ -0,0 +1,44 @@ +using System.Buffers; +using System.Buffers.Text; +using System.Diagnostics; +using System.Runtime.CompilerServices; +using System.Text.Encodings.Web; + +namespace Slon.Fortunes.Platform; + +internal static class RawFortuneTemplating +{ + [MethodImpl(MethodImplOptions.AggressiveOptimization)] + internal static void Render( + List fortunes, IBufferWriter writer, HtmlEncoder encoder) + { + const int FortunesTemplateLength = 1232; + var span = writer.GetSpan(FortunesTemplateLength); + var startingLength = span.Length; + span = span.WriteAndSlice( + "Fortunes"u8); + foreach (var fortune in fortunes) + { + var current = span.WriteAndSlice(""u8); + } + span = span.WriteAndSlice("
idmessage
"u8); + var success = Utf8Formatter.TryFormat((uint)fortune.Id, current, out var written); + Debug.Assert(success); + current = current.Slice(written).WriteAndSlice(""u8); + var status = encoder.EncodeUtf8( + fortune.Message.Span, current, out _, out written, isFinalBlock: true); + Debug.Assert(status is OperationStatus.Done); + span = current.Slice(written).WriteAndSlice("
"u8); + writer.Advance(startingLength - span.Length); + } +} + +internal static class SpanExtensions +{ + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal static Span WriteAndSlice(this Span destination, ReadOnlySpan source) + { + source.CopyTo(destination); + return destination.Slice(source.Length); + } +} diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/Slon.Fortunes.Platform.csproj b/Slon.Benchmarks/Slon.Fortunes.Platform/Slon.Fortunes.Platform.csproj new file mode 100644 index 0000000..51aebfb --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/Slon.Fortunes.Platform.csproj @@ -0,0 +1,22 @@ + + + + net10.0 + enable + enable + false + $(DefineConstants);DATABASE + $(NoWarn);SLONPG001;SLONPOOL001 + + + + + + + + + + + + + diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/Templates/Fortunes.cshtml b/Slon.Benchmarks/Slon.Fortunes.Platform/Templates/Fortunes.cshtml new file mode 100644 index 0000000..1cbd846 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/Templates/Fortunes.cshtml @@ -0,0 +1,2 @@ +@inherits RazorSlice> +Fortunes@foreach (var item in Model){}
idmessage
@WriteNumber(item.Id, default, CultureInfo.InvariantCulture, false)@item.Message.Span
\ No newline at end of file diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/Templates/_ViewImports.cshtml b/Slon.Benchmarks/Slon.Fortunes.Platform/Templates/_ViewImports.cshtml new file mode 100644 index 0000000..068bf13 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/Templates/_ViewImports.cshtml @@ -0,0 +1,9 @@ +@inherits RazorSlice + +@using System.Globalization +@using Microsoft.AspNetCore.Razor +@using RazorSlices +@using Slon.Fortunes.Platform + +@tagHelperPrefix __disable_tagHelpers__: +@removeTagHelper *, Microsoft.AspNetCore.Mvc.Razor diff --git a/Slon.Benchmarks/Slon.Fortunes.Platform/platform-fortunes.benchmarks.yml b/Slon.Benchmarks/Slon.Fortunes.Platform/platform-fortunes.benchmarks.yml new file mode 100644 index 0000000..13bcec9 --- /dev/null +++ b/Slon.Benchmarks/Slon.Fortunes.Platform/platform-fortunes.benchmarks.yml @@ -0,0 +1,91 @@ +imports: + - https://raw.githubusercontent.com/dotnet/crank/main/src/Microsoft.Crank.Jobs.Wrk/wrk.yml + - https://raw.githubusercontent.com/aspnet/Benchmarks/main/scenarios/aspnet.profiles.standard.yml + +variables: + serverPort: 5000 + npgsqlConnections: 256 + slonFlowPoolCapacity: 1024 + branchOrCommit: sebros/slon-benchmarks + draghiBranchOrCommit: experiment/observation-frontier + +jobs: + platform-postgresql-slon: + sources: + Slon: + repository: https://github.com/draghidev/slon.git + branchOrCommit: "{{branchOrCommit}}" + Draghi: + repository: https://github.com/draghidev/pipelining.git + branchOrCommit: "{{draghiBranchOrCommit}}" + project: Slon/Slon.Benchmarks/Slon.Fortunes.Platform/Slon.Fortunes.Platform.csproj + sdkVersion: 11.0.100-preview.7.26381.103 + framework: net10.0 + patchReferences: true + readyStateText: Application started. + arguments: "--urls {{serverScheme}}://{{serverAddress}}:{{serverPort}}" + variables: + serverScheme: http + environmentVariables: + DATABASE: postgresql + DRIVER: slon + CONNECTION_STRING: "Host={{databaseServer}};Database=hello_world;Username=benchmarkdbuser;Password=benchmarkdbpass" + DATABASE_CONNECTIONS: "{{ cores | minus: 2 }}" + SLON_FLOW_POOL_CAPACITY: "{{slonFlowPoolCapacity}}" + + platform-postgresql-npgsql: + sources: + Slon: + repository: https://github.com/draghidev/slon.git + branchOrCommit: "{{branchOrCommit}}" + Draghi: + repository: https://github.com/draghidev/pipelining.git + branchOrCommit: "{{draghiBranchOrCommit}}" + project: Slon/Slon.Benchmarks/Slon.Fortunes.Platform/Slon.Fortunes.Platform.csproj + sdkVersion: 11.0.100-preview.7.26381.103 + framework: net10.0 + patchReferences: true + readyStateText: Application started. + arguments: "--urls {{serverScheme}}://{{serverAddress}}:{{serverPort}}" + variables: + serverScheme: http + environmentVariables: + DATABASE: postgresql + DRIVER: npgsql + CONNECTION_STRING: "Host={{databaseServer}};Database=hello_world;Username=benchmarkdbuser;Password=benchmarkdbpass;Maximum Pool Size={{ npgsqlConnections }};No Reset On Close=true;Enlist=false;Max Auto Prepare=4" + DATABASE_CONNECTIONS: "{{ npgsqlConnections }}" + + postgresql: + source: + repository: https://github.com/TechEmpower/FrameworkBenchmarks.git + branchOrCommit: master + dockerFile: toolset/databases/postgres/postgres.dockerfile + dockerImageName: postgres_te + dockerContextDirectory: toolset/databases/postgres + readyStateText: ready to accept connections + noClean: true + +scenarios: + platform-postgresql-slon: + db: + job: postgresql + application: + job: platform-postgresql-slon + load: + job: wrk + variables: + presetHeaders: html + path: /fortunes + connections: 512 + + platform-postgresql-npgsql: + db: + job: postgresql + application: + job: platform-postgresql-npgsql + load: + job: wrk + variables: + presetHeaders: html + path: /fortunes + connections: 512 diff --git a/Slon.Tests/Ado/FlowMigrationTests.cs b/Slon.Tests/Ado/FlowMigrationTests.cs index 320d5df..fe32258 100644 --- a/Slon.Tests/Ado/FlowMigrationTests.cs +++ b/Slon.Tests/Ado/FlowMigrationTests.cs @@ -183,4 +183,7 @@ static async Task DrainAsync(CommandFlow flow) await e.DisposeAsync(); } + static async Task DrainAsync(BindingProbeFlow flow) + => await flow.WaitForComplete(); + } diff --git a/Slon.Tests/Ado/ReaderDisposalTests.cs b/Slon.Tests/Ado/ReaderDisposalTests.cs index 6a3ba5b..5ca5f6b 100644 --- a/Slon.Tests/Ado/ReaderDisposalTests.cs +++ b/Slon.Tests/Ado/ReaderDisposalTests.cs @@ -97,17 +97,17 @@ public async Task HiddenLaterResult_DisposalCreatesCancellationState(bool async) Assert.IsTrue(HasCancellationState(flow)); } - static CommandFlow GetFlow(SlonDataReader reader) + static AdoCommandExecutionFlow GetFlow(SlonDataReader reader) { var enumerator = typeof(SlonDataReader).GetField("_enumerator", AllInstanceFields)! .GetValue(reader)!; var flowField = enumerator.GetType().GetFields(AllInstanceFields) - .Single(static field => field.FieldType == typeof(CommandFlow)); - return (CommandFlow)flowField.GetValue(enumerator)!; + .Single(static field => field.FieldType == typeof(AdoCommandExecutionFlow)); + return (AdoCommandExecutionFlow)flowField.GetValue(enumerator)!; } - static bool HasCancellationState(CommandFlow flow) - => FindField(flow.GetType(), "_cancellationState").GetValue(flow) is not null; + static bool HasCancellationState(AdoCommandExecutionFlow flow) + => flow.HasCancellationState; static FieldInfo FindField(Type type, string name) { diff --git a/Slon.Tests/DataReaderTests.cs b/Slon.Tests/DataReaderTests.cs index 7292e6a..1d951dd 100644 --- a/Slon.Tests/DataReaderTests.cs +++ b/Slon.Tests/DataReaderTests.cs @@ -181,6 +181,26 @@ public async Task ExecuteNonQuery_NonDataModifying_IsZero() Assert.AreEqual(-1, await AdoTestPool.ExecuteNonQueryAsync("SELECT generate_series(1, 10)"), "SELECT 10 rows"); } + [TestMethod] + public async Task ExecuteReader_ReportsRowsAffectedAfterSkippingNonRowResult() + { + await using var connection = await AdoTestPool.OpenConnectionAsync(); + await using (var setup = connection.CreateCommand( + "CREATE TEMP TABLE reader_records_affected (value int)")) + _ = await setup.ExecuteNonQueryAsync(); + await using (var insert = connection.CreateCommand( + "INSERT INTO reader_records_affected VALUES (1), (2)")) + _ = await insert.ExecuteNonQueryAsync(); + + await using var command = connection.CreateCommand( + "UPDATE reader_records_affected SET value = value + 1"); + await using var reader = await command.ExecuteReaderAsync(); + + Assert.IsFalse(await reader.ReadAsync()); + Assert.AreEqual(2, reader.RecordsAffected); + Assert.AreEqual(2L, reader.LongRecordsAffected); + } + [TestMethod] public async Task BatchExecuteNonQuery_SumsAllCommandResults() { diff --git a/Slon.Tests/FlowBindingProbe.cs b/Slon.Tests/FlowBindingProbe.cs index 00a0289..bb9db11 100644 --- a/Slon.Tests/FlowBindingProbe.cs +++ b/Slon.Tests/FlowBindingProbe.cs @@ -9,17 +9,40 @@ sealed class BindingProbeContext(string name) : PgClientFlowBindingContext internal string Name { get; } = name; } -sealed class BindingProbeFlow(bool fail = false) : CommandFlow(async: true, []) +sealed class BindingProbeFlow : PgClientFlow { + readonly bool _fail; internal int BindCount { get; private set; } internal string? ContextName { get; private set; } + internal BindingProbeFlow(bool fail = false) + : base(supportsDeferredFlush: true) + { + _fail = fail; + IsAsync = true; + } + + protected override bool EnableActivationTimeout => true; + + protected override ValueTask ExecuteAuto(Context context) + { + var write = new CommandList(Command.Create("select 1")) + .WriteCommandsAsync(context.GetEncoder(), appendSync: true); + return new(new FlowTasks(write, DrainAsync(context))); + + static async ValueTask DrainAsync(Context context) + { + var decoder = await context.GetDecoderAsync().ConfigureAwait(false); + while (context.OutstandingRfqCount is not 0) + _ = await decoder.GetNextAsync().ConfigureAwait(false); + } + } + internal override void Bind(PgClientFlowBindingContext? context) { BindCount++; ContextName = ((BindingProbeContext)context!).Name; - if (fail) + if (_fail) throw new InvalidOperationException("binding rejected"); - Initialize(IsAsync, Command.Create("select 1")); } } diff --git a/Slon.Tests/Pg/BackendMessageParsingTests.cs b/Slon.Tests/Pg/BackendMessageParsingTests.cs index f6f169c..cb65f8e 100644 --- a/Slon.Tests/Pg/BackendMessageParsingTests.cs +++ b/Slon.Tests/Pg/BackendMessageParsingTests.cs @@ -32,7 +32,7 @@ static BackendMessage Message(PgTypes.BackendType type, ReadOnlySpan body) BinaryPrimitives.WriteInt32BigEndian(bytes.AsSpan(1), length); body.CopyTo(bytes.AsSpan(BackendHeader.ByteCount)); var context = new BackendMessageContext(); - context.SetBatch(new BackendMessageBatch(new ReadOnlySequence(bytes))); + context.SetCursor(new BackendMessageCursor(new ReadOnlySequence(bytes))); Assert.IsTrue(context.TryMoveNext()); return context.Current; } diff --git a/Slon.Tests/Pg/BackendMessageStreamingTests.cs b/Slon.Tests/Pg/BackendMessageStreamingTests.cs index 93e9620..3b23529 100644 --- a/Slon.Tests/Pg/BackendMessageStreamingTests.cs +++ b/Slon.Tests/Pg/BackendMessageStreamingTests.cs @@ -1,6 +1,7 @@ using System.Buffers; using System.Buffers.Binary; using System.IO.Pipelines; +using System.Runtime.InteropServices; using Slon.Pipelines; using Slon.Pg.Protocol; using static Slon.Pg.Protocol.PgTypes; @@ -12,14 +13,43 @@ namespace Slon.Tests.Pg; [TestClass] public class BackendMessageStreamingTests { + sealed class SequenceSegment : ReadOnlySequenceSegment + { + public SequenceSegment(ReadOnlyMemory memory) => Memory = memory; + + public SequenceSegment Append(ReadOnlyMemory memory) + { + var next = new SequenceSegment(memory) + { + RunningIndex = RunningIndex + Memory.Length + }; + Next = next; + return next; + } + + public ReadOnlySequence To(SequenceSegment end) + => new(this, 0, end, end.Memory.Length); + } + + sealed class TestMemoryManager(byte[] buffer) : MemoryManager + { + public override Span GetSpan() => buffer; + public override MemoryHandle Pin(int elementIndex = 0) + => throw new NotSupportedException(); + public override void Unpin() { } + protected override void Dispose(bool disposing) { } + } + sealed class RejectRetiredSuppliedReadReader(PipeReader inner) : PipeReader { ReadResult _activeRead; bool _rejectAdvanceAtStart; + long? _expectedAdvanceOffset; public Action? BeforeAdvance { get; set; } public void RejectAdvanceAtActiveStart() => _rejectAdvanceAtStart = true; + public void ExpectAdvanceAtActiveOffset(long offset) => _expectedAdvanceOffset = offset; public override ValueTask ReadAsync(CancellationToken cancellationToken = default) { @@ -46,6 +76,12 @@ public override void AdvanceTo(SequencePosition consumed, SequencePosition exami if (_rejectAdvanceAtStart && consumed.Equals(_activeRead.Buffer.Start)) Assert.Fail("The supplied read was retired before its buffer was inspected."); _rejectAdvanceAtStart = false; + if (_expectedAdvanceOffset is { } expectedOffset) + { + Assert.AreEqual(_activeRead.Buffer.GetPosition(expectedOffset), consumed, + "The read pipe retained bytes before the active result tenure."); + _expectedAdvanceOffset = null; + } inner.AdvanceTo(consumed, examined); } @@ -62,60 +98,6 @@ enum LifetimeAction Retire, } - // Treats the first 4 big-endian bytes as the total segment length. Fully buffered => Done (so the - // enumerator takes the deferred-consume branch, exactly the state the defect needs). - struct FixedSegmenter : IPipeSegmenter - { - public int MinimumSize => 4; - - public OperationStatus CreateSegment(in ReadOnlySequence buffer, out long segmentLength, out int segment) - { - segment = 0; - var reader = new SequenceReader(buffer); - if (!reader.TryReadBigEndian(out int len)) - { - segmentLength = 0; - return OperationStatus.NeedMoreData; - } - segmentLength = len; - if (buffer.Length < len) - return OperationStatus.NeedMoreData; - segment = len; - return OperationStatus.Done; - } - } - - struct StreamingSegmenter : IPipeSegmenter> - { - public int MinimumSize => 4; - - public OperationStatus CreateSegment(in ReadOnlySequence buffer, out long segmentLength, - out ReadOnlySequence segment) - { - var reader = new SequenceReader(buffer); - if (!reader.TryReadBigEndian(out int len)) - { - segmentLength = 0; - segment = default; - return OperationStatus.NeedMoreData; - } - - segmentLength = len; - segment = buffer.Slice(0, Math.Min(buffer.Length, len)); - return buffer.Length < len ? OperationStatus.NeedMoreData : OperationStatus.Done; - } - } - - static byte[] LenPrefixed(int total) - { - var bytes = new byte[total]; - bytes[0] = (byte)(total >> 24); - bytes[1] = (byte)(total >> 16); - bytes[2] = (byte)(total >> 8); - bytes[3] = (byte)total; - return bytes; - } - static byte[] BackendMessageBytes(BackendType type, int totalLength) { var bytes = new byte[totalLength]; @@ -131,15 +113,19 @@ static byte[] BackendMessageBytes(BackendType type, ReadOnlySpan body) return bytes; } - static PipeSegmentEnumerator BuildEnumerator(byte[] wire) + static ReadOnlySequence Segmented( + ReadOnlyMemory first, ReadOnlyMemory second) { - // A MemoryStream returns the wire bytes then 0 (EOF) on every subsequent read, so re-drives - // after completion re-hit the same terminal state the recovery drain does against a closed peer. - var reader = new DefaultStreamPipeReader( - new MemoryStream(wire, writable: false), - new StreamPipeReaderOptions(bufferSize: 8192, useZeroByteReads: false), - supportCancelPending: false); - return new(reader, new FixedSegmenter(), ownsReader: true); + var start = new SequenceSegment(first); + return start.To(start.Append(second)); + } + + static async ValueTask ReadNextAsync(ProtocolReadPipe pipe) + { + pipe.PrepareRead(); + var read = await pipe.ReadAsync(CancellationToken.None); + return pipe.CompleteRead( + read, CancellationToken.None, out _); } [TestMethod] @@ -159,6 +145,21 @@ public void BackendMessage_BodyAccessThrowsAfterStreamingWindowAdvances() Assert.ThrowsExactly(() => context.TryExtend(0, out _)); } + [TestMethod] + public void IndependentBackendMessage_ReconstructsAChainedBuffer() + { + var bytes = BackendMessageBytes(BackendType.DataRow, [1, 2, 3, 4, 5, 6]); + var sequence = Segmented(bytes.AsMemory(0, 7), bytes.AsMemory(7)); + var message = BackendMessage.CreateIndependent( + new BackendHeader(BackendType.DataRow, bytes.Length - 1), sequence); + + CollectionAssert.AreEqual(bytes.AsSpan(BackendHeader.ByteCount).ToArray(), + message.GetSequence().ToArray()); + Assert.IsTrue(message.TryGetFirstSpan(0, out var first)); + CollectionAssert.AreEqual(bytes.AsSpan(BackendHeader.ByteCount, 2).ToArray(), + first.ToArray()); + } + [TestMethod] public void BackendMessageContext_CurrentThrowsOutsidePublicationWindow() { @@ -167,7 +168,7 @@ public void BackendMessageContext_CurrentThrowsOutsidePublicationWindow() Assert.IsFalse(context.TryGetCurrent(out _)); Assert.ThrowsExactly(() => _ = context.Current); - context.SetBatch(new BackendMessageBatch( + context.SetCursor(new BackendMessageCursor( new ReadOnlySequence(BackendMessageBytes(BackendType.CommandComplete, 6)))); Assert.IsTrue(context.TryMoveNext()); Assert.IsTrue(context.TryGetCurrent(out var current)); @@ -177,24 +178,105 @@ public void BackendMessageContext_CurrentThrowsOutsidePublicationWindow() Assert.IsTrue(context.TryGetCurrent(out current)); Assert.AreEqual(BackendType.CommandComplete, current.Header.Type); Assert.AreEqual(BackendType.CommandComplete, accessor.Message.Header.Type); - context.RetireCurrentBatch(); + context.RetireCursor(); Assert.IsFalse(context.TryGetCurrent(out _)); Assert.ThrowsExactly(() => _ = context.Current); Assert.ThrowsExactly(() => _ = accessor.Message); } [TestMethod] - public async Task MovingToNextBatch_RetiresCurrentBeforeReturningItsStorage() + public void BackendMessageCursor_AdvancesAcrossExactSegmentBoundary() + { + var first = BackendMessageBytes(BackendType.CommandComplete, 6); + var second = BackendMessageBytes(BackendType.ReadyForQuery, 6); + var cursor = new BackendMessageCursor(Segmented(first, second)); + + Assert.IsTrue(cursor.TryReadNextInPlace(out var header, out var message, out _)); + Assert.AreEqual(BackendType.CommandComplete, header.Type); + CollectionAssert.AreEqual(first, message.ToArray()); + Assert.IsTrue(cursor.TryReadNextInPlace(out header, out message, out _)); + Assert.AreEqual(BackendType.ReadyForQuery, header.Type); + CollectionAssert.AreEqual(second, message.ToArray()); + Assert.IsFalse(cursor.TryReadNextInPlace(out _, out _, out _)); + } + + [TestMethod] + public void BackendMessageCursor_CollapsesRemainingFinalSegmentToArrayBacking() + { + var first = BackendMessageBytes(BackendType.CommandComplete, 6); + var second = BackendMessageBytes(BackendType.ReadyForQuery, 6); + var cursor = new BackendMessageCursor(Segmented(first, second)); + + Assert.IsTrue(cursor.TryReadNextInPlace(out _, out _, out _)); + Assert.IsTrue(cursor.TryReadNextInPlace(out _, out var remaining, out _)); + + Assert.IsInstanceOfType(remaining.Start.GetObject()); + Assert.AreSame(second, remaining.Start.GetObject()); + CollectionAssert.AreEqual(second, remaining.ToArray()); + } + + [TestMethod] + public void BackendMessageCursor_CollapsesFinalSegmentAfterStraddlingMessage() + { + var straddling = BackendMessageBytes( + BackendType.DataRow, new byte[] { 1, 2, 3, 4, 5, 6 }); + var remainingMessage = BackendMessageBytes(BackendType.ReadyForQuery, 6); + const int firstSegmentLength = 7; + var finalSegment = new byte[straddling.Length - firstSegmentLength + + remainingMessage.Length + 6]; + straddling.AsSpan(firstSegmentLength).CopyTo(finalSegment.AsSpan(3)); + remainingMessage.CopyTo(finalSegment.AsSpan( + 3 + straddling.Length - firstSegmentLength)); + var finalMemory = finalSegment.AsMemory(3, + straddling.Length - firstSegmentLength + remainingMessage.Length); + var cursor = new BackendMessageCursor(Segmented( + straddling.AsMemory(0, firstSegmentLength), finalMemory)); + + Assert.IsTrue(cursor.TryReadNextInPlace(out _, out var first, out _)); + Assert.IsFalse(first.IsSingleSegment); + Assert.IsTrue(cursor.TryReadNextInPlace(out _, out var remaining, out _)); + + Assert.AreSame(finalSegment, remaining.Start.GetObject()); + CollectionAssert.AreEqual(remainingMessage, remaining.ToArray()); + } + + [TestMethod] + public void ContiguousMemory_StraddleIsProjectedOncePerResultTenure() + { + var bytes = BackendMessageBytes( + BackendType.DataRow, new byte[] { 0, 1, 2, 3, 4, 5 }); + var context = new BackendMessageContext(); + context.SetCursor(new(Segmented( + bytes.AsMemory(0, 7), bytes.AsMemory(7)))); + Assert.IsTrue(context.TryMoveNext()); + var field = context.Current.GetSequence(); + Assert.IsFalse(field.IsSingleSegment); + + var first = context.Current.GetContiguousMemory(field); + var second = context.Current.GetContiguousMemory(field); + CollectionAssert.AreEqual(field.ToArray(), first.ToArray()); + Assert.IsTrue(MemoryMarshal.TryGetArray(first, out var firstArray)); + Assert.IsTrue(MemoryMarshal.TryGetArray(second, out var secondArray)); + Assert.AreSame(firstArray.Array, secondArray.Array); + + context.ReleaseContiguousProjections(); + context.RetireCursor(); + } + + [TestMethod] + public async Task MovingToNextRead_RetiresCurrentBeforeReturningItsStorage() { var pipe = new Pipe(); var reader = new RejectRetiredSuppliedReadReader(pipe.Reader); - var protocolPipe = new ProtocolReadPipe( - new(reader, new BackendMessageBatch.Segmenter(), ownsReader: true)); + var protocolPipe = new ProtocolReadPipe(reader, + BackendMessageCursor.DefaultDataRowStreamingThreshold, + ownsReader: true); await pipe.Writer.WriteAsync(BackendMessageBytes(BackendType.CommandComplete, 6)); - Assert.IsTrue(protocolPipe.TryMoveNextBatch(out _)); + Assert.IsTrue(await ReadNextAsync(protocolPipe)); Assert.IsTrue(protocolPipe.TryMoveNext()); var accessor = protocolPipe.Current.GetAccessor(); + Assert.IsFalse(protocolPipe.TryMoveNext()); var observedAdvance = false; reader.BeforeAdvance = () => @@ -207,7 +289,7 @@ public async Task MovingToNextBatch_RetiresCurrentBeforeReturningItsStorage() }; await pipe.Writer.WriteAsync(BackendMessageBytes(BackendType.ReadyForQuery, 6)); - Assert.IsTrue(protocolPipe.TryMoveNextBatch(out _)); + Assert.IsTrue(await ReadNextAsync(protocolPipe)); Assert.IsTrue(observedAdvance); Assert.IsTrue(protocolPipe.TryMoveNext()); Assert.AreEqual(BackendType.ReadyForQuery, protocolPipe.Current.Header.Type); @@ -238,15 +320,15 @@ public void BackendMessageContext_CurrentLifetime_ExhaustiveShortSequences() { case LifetimeAction.LoadCommandComplete: case LifetimeAction.LoadReadyForQuery: - // ProtocolReadPipe retires the prior batch before committing replacement - // storage. Model that ownership boundary rather than calling SetBatch as a + // ProtocolReadPipe retires the prior cursor before committing replacement + // storage. Model that ownership boundary rather than calling SetCursor as a // replacement operation it is not. - context.RetireCurrentBatch(); + context.RetireCursor(); current = null; loaded = action is LifetimeAction.LoadCommandComplete ? BackendType.CommandComplete : BackendType.ReadyForQuery; - context.SetBatch(new(new ReadOnlySequence( + context.SetCursor(new(new ReadOnlySequence( BackendMessageBytes(loaded.Value, 6)))); break; case LifetimeAction.MoveNext: @@ -258,7 +340,7 @@ public void BackendMessageContext_CurrentLifetime_ExhaustiveShortSequences() } break; case LifetimeAction.Retire: - context.RetireCurrentBatch(); + context.RetireCursor(); loaded = null; current = null; break; @@ -275,113 +357,6 @@ public void BackendMessageContext_CurrentLifetime_ExhaustiveShortSequences() } } - [TestMethod] - public async Task ReDriveAfterEof_Async_ReturnsFalseWithoutCorruption() - { - var e = BuildEnumerator(LenPrefixed(24)); - - Assert.IsTrue(await e.MoveNextAsync(), "first segment should be produced"); - Assert.AreEqual(24, e.Current); - Assert.IsFalse(await e.MoveNextAsync(), "second call consumes the segment and reaches EOF"); - - // The recovery drain keeps pulling after completion; before the fix each of these re-applied the - // stale deferred advance and threw ArgumentOutOfRangeException('length') from the ReadResult build. - for (var i = 0; i < 6; i++) - Assert.IsFalse(await e.MoveNextAsync(), $"re-drive #{i} past completion must stay false"); - - await e.DisposeAsync(); - } - - [TestMethod] - public async Task CompletedRead_WithFinalBufferedSegment_StillPublishesIt() - { - var pipe = new Pipe(); - var e = new PipeSegmentEnumerator(pipe.Reader, new FixedSegmenter()); - await pipe.Writer.WriteAsync(LenPrefixed(24)); - await pipe.Writer.CompleteAsync(); - - Assert.IsTrue(await e.MoveNextAsync()); - Assert.AreEqual(24, e.Current); - Assert.IsFalse(await e.MoveNextAsync()); - await e.DisposeAsync(); - } - - [TestMethod] - public async Task CompletedRead_WithTruncatedFinalSegment_Throws() - { - var pipe = new Pipe(); - var e = new PipeSegmentEnumerator(pipe.Reader, new FixedSegmenter()); - await pipe.Writer.WriteAsync(LenPrefixed(24).AsMemory(0, 12)); - await pipe.Writer.CompleteAsync(); - - Assert.IsTrue(await e.MoveNextAsync(), "the useful prefix is published before its missing tail is discovered"); - await Assert.ThrowsExactlyAsync(async () => await e.MoveNextAsync()); - await e.DisposeAsync(); - } - - [TestMethod] - public async Task CompletedRead_AfterPendingContinuation_Throws() - { - var pipe = new Pipe(); - var e = new PipeSegmentEnumerator>( - pipe.Reader, new StreamingSegmenter()); - await pipe.Writer.WriteAsync(LenPrefixed(24).AsMemory(0, 12)); - - Assert.IsTrue(await e.MoveNextAsync()); - Assert.IsFalse(e.TryContinueCurrentSegment(e.Current.End, e.Current.Length, out _)); - await pipe.Writer.CompleteAsync(); - - await Assert.ThrowsExactlyAsync(async () => await e.MoveNextAsync()); - await e.DisposeAsync(); - } - - [TestMethod] - public async Task CompletionAtEveryByteInsideSegment_Throws() - { - const int segmentLength = 32; - var wire = LenPrefixed(segmentLength); - - for (var cut = 1; cut < segmentLength; cut++) - { - var pipe = new Pipe(); - var e = new PipeSegmentEnumerator(pipe.Reader, new FixedSegmenter()); - await pipe.Writer.WriteAsync(wire.AsMemory(0, cut)); - await pipe.Writer.CompleteAsync(); - - await Assert.ThrowsExactlyAsync(async () => - { - while (await e.MoveNextAsync()) { } - }, $"completion at byte {cut} must not become clean EOF"); - await e.DisposeAsync(); - } - } - - [TestMethod] - public async Task CompletionAtEveryByteInsideSuccessor_ThrowsAfterPredecessor() - { - const int firstLength = 12; - const int secondLength = 32; - var wire = new byte[firstLength + secondLength]; - LenPrefixed(firstLength).CopyTo(wire, 0); - LenPrefixed(secondLength).CopyTo(wire, firstLength); - - for (var successorBytes = 1; successorBytes < secondLength; successorBytes++) - { - var pipe = new Pipe(); - var e = new PipeSegmentEnumerator(pipe.Reader, new FixedSegmenter()); - await pipe.Writer.WriteAsync(wire.AsMemory(0, firstLength + successorBytes)); - await pipe.Writer.CompleteAsync(); - - Assert.IsTrue(await e.MoveNextAsync(), $"predecessor missing at successor byte {successorBytes}"); - Assert.AreEqual(firstLength, e.Current); - await Assert.ThrowsExactlyAsync(async () => - { - while (await e.MoveNextAsync()) { } - }, $"completion at successor byte {successorBytes} must not become clean EOF"); - await e.DisposeAsync(); - } - } - [TestMethod] public async Task RepeatedQueryFrames_WithSmallRecycledBuffers_NeverEnterMessageBodies() { @@ -403,24 +378,26 @@ public async Task RepeatedQueryFrames_WithSmallRecycledBuffers_NeverEnterMessage new MemoryStream(wire, writable: false), new StreamPipeReaderOptions(bufferSize: 1024, useZeroByteReads: false), supportCancelPending: false); - var e = new PipeSegmentEnumerator( - reader, new BackendMessageBatch.Segmenter(), ownsReader: true); + var readPipe = new ProtocolReadPipe(reader, + BackendMessageCursor.DefaultDataRowStreamingThreshold, + ownsReader: true); var messageIndex = 0; - while (await e.MoveNextAsync()) + while (await readPipe.MoveNextAsync(default)) { - var batch = e.Current; - while (batch.TryReadNextInPlace(out var header, out _, out _)) + while (readPipe.TryMoveNext()) { - Assert.AreEqual(response[messageIndex % response.Length][0], (byte)header.Type, + Assert.AreEqual(response[messageIndex % response.Length][0], + (byte)readPipe.Current.Header.Type, $"message {messageIndex}"); messageIndex++; } } Assert.AreEqual(repetitions * response.Length, messageIndex); - await e.DisposeAsync(); + await readPipe.DisposeAsync(); } +#if !NET11_0_OR_GREATER [TestMethod] public async Task RepeatedQueryFrames_ThroughDirectReads_NeverEnterMessageBodies() { @@ -442,49 +419,51 @@ public async Task RepeatedQueryFrames_ThroughDirectReads_NeverEnterMessageBodies new MemoryStream(wire, writable: false), new StreamPipeReaderOptions(bufferSize: 1024, useZeroByteReads: false), supportCancelPending: false); - var e = new PipeSegmentEnumerator( - reader, new BackendMessageBatch.Segmenter(), ownsReader: true); + var readPipe = new ProtocolReadPipe(reader, + BackendMessageCursor.DefaultDataRowStreamingThreshold, + ownsReader: true); + var directReader = (StreamPipeReader)readPipe.PipeReader; var messageIndex = 0; while (true) { - if (e.TryMoveNext(out var completed)) - { - ValidateBatch(e.Current); - continue; - } - if (completed) - break; - - Assert.IsTrue(e.TryBeginDirectRead(default, out var read)); + readPipe.PrepareRead(); + Assert.IsTrue(directReader.SupportsDirectRead); + var read = directReader.BeginDirectRead(default); while (true) { var length = await read; - if (e.CompleteDirectRead(length, default, out read, out var readFinished, out completed)) + if (!directReader.CompleteDirectRead(length, default, out read, out var result)) + { + continue; + } + if (readPipe.CompleteRead( + result, default, out var completed)) { - ValidateBatch(e.Current); + ValidateMessages(); break; } - if (!readFinished) - continue; - Assert.IsTrue(completed); - goto done; + if (completed) + goto done; + break; } } done: Assert.AreEqual(repetitions * response.Length, messageIndex); - await e.DisposeAsync(); + await readPipe.DisposeAsync(); - void ValidateBatch(BackendMessageBatch batch) + void ValidateMessages() { - while (batch.TryReadNextInPlace(out var header, out _, out _)) + while (readPipe.TryMoveNext()) { - Assert.AreEqual(response[messageIndex % response.Length][0], (byte)header.Type, + Assert.AreEqual(response[messageIndex % response.Length][0], + (byte)readPipe.Current.Header.Type, $"message {messageIndex}"); messageIndex++; } } } +#endif static byte[][] QueryResponseBytes() { @@ -512,214 +491,119 @@ static byte[][] QueryResponseBytes() public async Task Eof_InvalidatesPublishedBackendMessage() { var pipe = new Pipe(); - var batches = new PipeSegmentEnumerator( - pipe.Reader, new BackendMessageBatch.Segmenter()); - var readPipe = new ProtocolReadPipe(batches); + var readPipe = new ProtocolReadPipe(pipe.Reader, + BackendMessageCursor.DefaultDataRowStreamingThreshold); await pipe.Writer.WriteAsync(BackendMessageBytes(BackendType.ReadyForQuery, 6)); Assert.IsTrue(await readPipe.MoveNextAsync(CancellationToken.None)); - readPipe.CommitBatch(); Assert.IsTrue(readPipe.TryMoveNext()); var accessor = readPipe.Current.GetAccessor(); + Assert.IsFalse(readPipe.TryMoveNext()); await pipe.Writer.CompleteAsync(); Assert.IsFalse(await readPipe.MoveNextAsync(CancellationToken.None)); Assert.IsFalse(readPipe.TryGetCurrent(out _)); Assert.ThrowsExactly(() => _ = readPipe.Current); Assert.ThrowsExactly(() => _ = accessor.Message); - await batches.DisposeAsync(); + await readPipe.DisposeAsync(); } [TestMethod] - public void ReDriveAfterEof_Sync_ReturnsFalseWithoutCorruption() - { - var e = BuildEnumerator(LenPrefixed(24)); - - Assert.IsTrue(e.MoveNext(), "first segment should be produced"); - Assert.AreEqual(24, e.Current); - Assert.IsFalse(e.MoveNext(), "second call consumes the segment and reaches EOF"); - - for (var i = 0; i < 6; i++) - Assert.IsFalse(e.MoveNext(), $"re-drive #{i} past completion must stay false"); - - e.Dispose(); - } + public async Task EndingResultRetention_ReexposesBufferedSuccessorMessages() + { + var first = BackendMessageBytes(BackendType.DataRow, [0, 0]); + var terminal = BackendMessageBytes(BackendType.CommandComplete, "SELECT 1\0"u8); + var ready = BackendMessageBytes(BackendType.ReadyForQuery, [(byte)'I']); + var successor = BackendMessageBytes(BackendType.BindComplete, []); + var wire = new byte[first.Length + terminal.Length + ready.Length + successor.Length]; + var offset = 0; + foreach (var message in (byte[][])[first, terminal, ready, successor]) + { + message.CopyTo(wire, offset); + offset += message.Length; + } - [TestMethod] - public async Task TryMoveNext_PollsFragmentedSegmentWithoutSuspending() - { var pipe = new Pipe(); - var e = new PipeSegmentEnumerator(pipe.Reader, new FixedSegmenter()); - - Assert.IsFalse(e.TryMoveNext(out var completed)); - Assert.IsFalse(completed); + var readPipe = new ProtocolReadPipe(pipe.Reader, + BackendMessageCursor.DefaultDataRowStreamingThreshold); + await pipe.Writer.WriteAsync(wire); - var wire = LenPrefixed(24); - await pipe.Writer.WriteAsync(wire.AsMemory(0, 2)); - Assert.IsFalse(e.TryMoveNext(out completed), "a partial header must request another read"); - Assert.IsFalse(completed); + Assert.IsTrue(await readPipe.MoveNextAsync(default)); + Assert.IsTrue(readPipe.TryMoveNext()); + Assert.AreEqual(BackendType.DataRow, readPipe.Current.Header.Type); + readPipe.EnableResultRetention(); + Assert.IsTrue(readPipe.TryMoveNext()); + Assert.AreEqual(BackendType.CommandComplete, readPipe.Current.Header.Type); + Assert.IsTrue(readPipe.TryMoveNext()); + Assert.AreEqual(BackendType.ReadyForQuery, readPipe.Current.Header.Type); - await pipe.Writer.WriteAsync(wire.AsMemory(2)); - Assert.IsTrue(e.TryMoveNext(out completed)); - Assert.IsFalse(completed); - Assert.AreEqual(24, e.Current); + readPipe.EndResultRetention(); + Assert.IsTrue(readPipe.TryMoveNext()); + Assert.AreEqual(BackendType.BindComplete, readPipe.Current.Header.Type); await pipe.Writer.CompleteAsync(); - Assert.IsFalse(e.TryMoveNext(out completed)); - Assert.IsTrue(completed); - await e.DisposeAsync(); - } - - [TestMethod] - public async Task DirectRead_PreservesFramingAndTerminalState() - { - var e = BuildEnumerator(LenPrefixed(24)); - - Assert.IsTrue(e.TryBeginDirectRead(default, out var read)); - var length = await read; - Assert.IsTrue(e.CompleteDirectRead(length, default, out _, out var readFinished, out var completed)); - Assert.IsTrue(readFinished); - Assert.IsFalse(completed); - Assert.AreEqual(24, e.Current); - - Assert.IsFalse(e.TryMoveNext(out completed)); - Assert.IsFalse(completed); - Assert.IsTrue(e.TryBeginDirectRead(default, out read)); - length = await read; - Assert.IsFalse(e.CompleteDirectRead(length, default, out _, out readFinished, out completed)); - Assert.IsTrue(readFinished); - Assert.IsTrue(completed); - - await e.DisposeAsync(); + await readPipe.DisposeAsync(); } [TestMethod] - public async Task ContinueCurrentSegment_StreamsWithoutCrossingNextSegment() + public async Task BeginningNewResultRetention_DropsPriorResultPrefixAtNextRead() { - var pipe = new Pipe(); - var e = new PipeSegmentEnumerator>( - pipe.Reader, new StreamingSegmenter()); - var first = LenPrefixed(12); - var second = LenPrefixed(8); - - await pipe.Writer.WriteAsync(first.AsMemory(0, 6)); - Assert.IsTrue(e.TryMoveNext(out _)); - Assert.AreEqual(6, e.Current.Length); - - Assert.IsFalse(e.TryContinueCurrentSegment(e.Current.End, e.Current.Length, out _)); + var prior = BackendMessageBytes(BackendType.DataRow, [0, 1]); + var retained = BackendMessageBytes(BackendType.DataRow, [2, 3]); + var terminal = BackendMessageBytes(BackendType.CommandComplete, "SELECT 1\0"u8); + var firstGrant = new byte[prior.Length + retained.Length]; + prior.CopyTo(firstGrant, 0); + retained.CopyTo(firstGrant, prior.Length); - var tail = new byte[first.Length - 6 + second.Length]; - first.AsSpan(6).CopyTo(tail); - second.CopyTo(tail.AsSpan(first.Length - 6)); - await pipe.Writer.WriteAsync(tail); - - Assert.IsTrue(e.TryContinueCurrentSegment(e.Current.End, e.Current.Length, out var continuation)); - Assert.IsTrue(continuation.IsComplete); - Assert.AreEqual(6, continuation.Buffer.Length); - - Assert.IsTrue(e.TryMoveNext(out _)); - Assert.AreEqual(8, e.Current.Length); - - await pipe.Writer.CompleteAsync(); - Assert.IsFalse(await e.MoveNextAsync()); - await e.DisposeAsync(); - } - - [TestMethod] - public async Task MoveNextAsync_SkipsUnconsumedPartialSegmentBeforeReadingNext() - { var pipe = new Pipe(); - var e = new PipeSegmentEnumerator>( - pipe.Reader, new StreamingSegmenter()); - var first = LenPrefixed(12); - var second = LenPrefixed(8); + var reader = new RejectRetiredSuppliedReadReader(pipe.Reader); + var readPipe = new ProtocolReadPipe(reader, + BackendMessageCursor.DefaultDataRowStreamingThreshold); + await pipe.Writer.WriteAsync(firstGrant); - await pipe.Writer.WriteAsync(first.AsMemory(0, 6)); - Assert.IsTrue(await e.MoveNextAsync()); - Assert.AreEqual(6, e.Current.Length); + Assert.IsTrue(await readPipe.MoveNextAsync(default)); + Assert.IsTrue(readPipe.TryMoveNext()); + readPipe.EnableResultRetention(); + readPipe.EndResultRetention(); + Assert.IsTrue(readPipe.TryMoveNext()); + readPipe.EnableResultRetention(); + Assert.IsFalse(readPipe.TryMoveNext()); - byte[] remaining = [.. first.AsSpan(6), .. second]; - await pipe.Writer.WriteAsync(remaining); - Assert.IsTrue(await e.MoveNextAsync()); - Assert.AreEqual(second.Length, e.Current.Length); + reader.ExpectAdvanceAtActiveOffset(prior.Length); + await pipe.Writer.WriteAsync(terminal); + Assert.IsTrue(await readPipe.MoveNextAsync(default)); + Assert.IsTrue(readPipe.TryMoveNext()); + Assert.AreEqual(BackendType.CommandComplete, readPipe.CurrentType); + readPipe.EndResultRetention(); await pipe.Writer.CompleteAsync(); - Assert.IsFalse(await e.MoveNextAsync()); - await e.DisposeAsync(); + await readPipe.DisposeAsync(); } [TestMethod] - public async Task TryMoveNext_SkipsUnconsumedPartialSegmentBeforePollingNext() + public async Task IdleRelease_UnexaminesAndReacquiresBufferedSuffix() { + var completed = BackendMessageBytes(BackendType.ReadyForQuery, [(byte)'I']); + var suffix = BackendMessageBytes(BackendType.NotificationResponse, [1, 2, 3]); + var wire = new byte[completed.Length + suffix.Length]; + completed.CopyTo(wire, 0); + suffix.CopyTo(wire, completed.Length); var pipe = new Pipe(); - var e = new PipeSegmentEnumerator>( - pipe.Reader, new StreamingSegmenter()); - var first = LenPrefixed(12); - var second = LenPrefixed(8); - - await pipe.Writer.WriteAsync(first.AsMemory(0, 6)); - Assert.IsTrue(e.TryMoveNext(out _)); - Assert.AreEqual(6, e.Current.Length); - - byte[] remaining = [.. first.AsSpan(6), .. second]; - await pipe.Writer.WriteAsync(remaining); - Assert.IsTrue(e.TryMoveNext(out _)); - Assert.AreEqual(second.Length, e.Current.Length); + var readPipe = new ProtocolReadPipe(pipe.Reader, + BackendMessageCursor.DefaultDataRowStreamingThreshold); + await pipe.Writer.WriteAsync(wire); - await pipe.Writer.CompleteAsync(); - Assert.IsFalse(e.TryMoveNext(out var completed)); - Assert.IsTrue(completed); - await e.DisposeAsync(); - } + Assert.IsTrue(await readPipe.MoveNextAsync(default)); + Assert.IsTrue(readPipe.TryMoveNext()); + Assert.AreEqual(BackendType.ReadyForQuery, readPipe.CurrentType); - [TestMethod] - public async Task ContinueCurrentSegmentAsync_PreservesPartialProgress() - { - var pipe = new Pipe(); - var e = new PipeSegmentEnumerator>( - pipe.Reader, new StreamingSegmenter()); - var wire = LenPrefixed(12); - - await pipe.Writer.WriteAsync(wire.AsMemory(0, 6)); - Assert.IsTrue(e.TryMoveNext(out _)); - - var pending = e.ContinueCurrentSegmentAsync(e.Current.End, e.Current.Length); - Assert.IsFalse(pending.IsCompleted); - await pipe.Writer.WriteAsync(wire.AsMemory(6, 3)); - var middle = await pending; - Assert.IsFalse(middle.IsComplete); - Assert.AreEqual(3, middle.Buffer.Length); - - pending = e.ContinueCurrentSegmentAsync(middle.Buffer.End, middle.Buffer.Length); - await pipe.Writer.WriteAsync(wire.AsMemory(9)); - var final = await pending; - Assert.IsTrue(final.IsComplete); - Assert.AreEqual(3, final.Buffer.Length); + readPipe.ReleaseReadBufferAtIdle(); + Assert.IsTrue(await readPipe.MoveNextAsync(default)); + Assert.IsTrue(readPipe.TryMoveNext()); + Assert.AreEqual(BackendType.NotificationResponse, readPipe.CurrentType); await pipe.Writer.CompleteAsync(); - Assert.IsFalse(await e.MoveNextAsync()); - await e.DisposeAsync(); - } - - [TestMethod] - public async Task ExtendCurrentSegmentAsync_RetainsUntilTheCompleteSegment() - { - var wire = LenPrefixed(128 * 1024); - var reader = new DefaultStreamPipeReader( - new MemoryStream(wire, writable: false), - new StreamPipeReaderOptions(bufferSize: 8192, useZeroByteReads: false), - supportCancelPending: false); - var e = new PipeSegmentEnumerator>( - reader, new StreamingSegmenter(), ownsReader: true); - - Assert.IsTrue(await e.MoveNextAsync()); - CurrentSegmentBuffer current; - do current = await e.ExtendCurrentSegmentAsync(); - while (!current.IsComplete); - - Assert.AreEqual(wire.Length, current.Buffer.Length); - Assert.IsFalse(await e.MoveNextAsync()); - await e.DisposeAsync(); + await readPipe.DisposeAsync(); } [TestMethod] @@ -730,13 +614,12 @@ public async Task BackendBodyReader_ExtendsPrefixThenSlides() for (var i = BackendHeader.ByteCount; i < wire.Length; i++) wire[i] = (byte)i; - var segments = new PipeSegmentEnumerator( - pipe.Reader, new BackendMessageBatch.Segmenter(8)); - var decoder = new PgDecoder(segments, CancellationToken.None, Timeout.InfiniteTimeSpan); + var decoder = new PgDecoder( + pipe.Reader, 8, CancellationToken.None, Timeout.InfiniteTimeSpan); decoder.Pipe.BindDecoder(decoder); await pipe.Writer.WriteAsync(wire.AsMemory(0, 8)); - Assert.IsTrue(decoder.Pipe.TryMoveNextBatch(out _)); + Assert.IsTrue(await ReadNextAsync(decoder.Pipe)); Assert.IsTrue(decoder.Pipe.TryMoveNext()); var body = decoder.Pipe.Current.OpenBodyReader(); Assert.AreEqual(3, body.Buffer.Length); @@ -758,7 +641,7 @@ public async Task BackendBodyReader_ExtendsPrefixThenSlides() } [TestMethod] - public async Task BackendSegmenter_ExtendedRowAdvancesToTrailingMessage() + public async Task BackendReadPipe_ExtendedRowAdvancesToTrailingMessage() { var bind = BackendMessageBytes(BackendType.BindComplete, BackendHeader.ByteCount); var row = BackendMessageBytes(BackendType.DataRow, 128 * 1024); @@ -767,102 +650,33 @@ public async Task BackendSegmenter_ExtendedRowAdvancesToTrailingMessage() bind.CopyTo(wire, 0); row.CopyTo(wire, bind.Length); complete.CopyTo(wire, bind.Length + row.Length); - var reader = new DefaultStreamPipeReader( - new MemoryStream(wire, writable: false), - new StreamPipeReaderOptions(bufferSize: 64 * 1024, useZeroByteReads: false), - supportCancelPending: false); - var e = new PipeSegmentEnumerator( - reader, new BackendMessageBatch.Segmenter(), ownsReader: true); - - Assert.IsTrue(await e.MoveNextAsync()); - CurrentSegmentBuffer current; - do current = await e.ExtendCurrentSegmentAsync(); - while (!current.IsComplete); - - Assert.IsTrue(await e.MoveNextAsync()); - Assert.IsTrue(e.Current.TryReadNextInPlace(out var header, out _, out _)); - Assert.AreEqual(BackendType.CommandComplete, header.Type); - await e.DisposeAsync(); - } - - [TestMethod] - public async Task ContinueCurrentSegment_SlidesPastOnlyTheConsumedPrefix() - { - var pipe = new Pipe(); - var e = new PipeSegmentEnumerator>( - pipe.Reader, new StreamingSegmenter()); - var wire = LenPrefixed(12); - for (var i = 4; i < wire.Length; i++) - wire[i] = (byte)i; - - await pipe.Writer.WriteAsync(wire.AsMemory(0, 8)); - Assert.IsTrue(e.TryMoveNext(out _)); - var consumed = e.Current.GetPosition(6); - Assert.IsFalse(e.TryContinueCurrentSegment(consumed, 6, out _)); - - await pipe.Writer.WriteAsync(wire.AsMemory(8, 2)); - Assert.IsTrue(e.TryContinueCurrentSegment(consumed, 6, out var middle)); - Assert.IsFalse(middle.IsComplete); - CollectionAssert.AreEqual(wire.AsSpan(6, 4).ToArray(), middle.Buffer.ToArray()); - - await pipe.Writer.WriteAsync(wire.AsMemory(10)); - var final = await e.ContinueCurrentSegmentAsync(middle.Buffer.End, middle.Buffer.Length); - Assert.IsTrue(final.IsComplete); - CollectionAssert.AreEqual(wire.AsSpan(10).ToArray(), final.Buffer.ToArray()); - - await pipe.Writer.CompleteAsync(); - Assert.IsFalse(await e.MoveNextAsync()); - await e.DisposeAsync(); - } - - [TestMethod] - public async Task MoveNext_CanDrainAfterAContinuationPollParks() - { - var pipe = new Pipe(); - var e = new PipeSegmentEnumerator>( - pipe.Reader, new StreamingSegmenter()); - var first = LenPrefixed(12); - var second = LenPrefixed(8); - - await pipe.Writer.WriteAsync(first.AsMemory(0, 6)); - Assert.IsTrue(e.TryMoveNext(out _)); - Assert.IsFalse(e.TryContinueCurrentSegment(e.Current.End, e.Current.Length, out _)); - - await pipe.Writer.WriteAsync(first.AsMemory(6)); - await pipe.Writer.WriteAsync(second); - Assert.IsTrue(await e.MoveNextAsync(), "normal iteration must take over the pending continuation read"); - Assert.AreEqual(second.Length, e.Current.Length); - - await pipe.Writer.CompleteAsync(); - Assert.IsFalse(await e.MoveNextAsync()); - await e.DisposeAsync(); - } - - [TestMethod] - public async Task SuppliedRead_IsInspectedBeforeItsConsumedPrefixIsRetired() - { - var pipe = new Pipe(); - var reader = new RejectRetiredSuppliedReadReader(pipe.Reader); - var e = new PipeSegmentEnumerator>( - reader, new StreamingSegmenter()); - var first = LenPrefixed(12); - var second = LenPrefixed(8); - - await pipe.Writer.WriteAsync(first.AsMemory(0, 6)); - Assert.IsTrue(e.TryMoveNext(out _)); - Assert.IsFalse(e.TryMoveNext(out _), "the poll should advance once and park for the tail"); - - await pipe.Writer.WriteAsync(first.AsMemory(6)); - await pipe.Writer.WriteAsync(second); - var supplied = await e.ReadAsync(CancellationToken.None); - reader.RejectAdvanceAtActiveStart(); + var pipe = new Pipe(new PipeOptions( + pauseWriterThreshold: 256 * 1024, + resumeWriterThreshold: 128 * 1024)); + var decoder = new PgDecoder(pipe.Reader, + BackendMessageCursor.DefaultDataRowStreamingThreshold, + CancellationToken.None, Timeout.InfiniteTimeSpan); + decoder.Pipe.BindDecoder(decoder); - Assert.IsTrue(e.TryMoveNext(supplied, CancellationToken.None, out _)); - Assert.AreEqual(second.Length, e.Current.Length); + var initialLength = bind.Length + + BackendMessageCursor.DefaultDataRowStreamingThreshold; + await pipe.Writer.WriteAsync(wire.AsMemory(0, initialLength)); + Assert.IsTrue(await decoder.Pipe.MoveNextAsync(default)); + Assert.IsTrue(decoder.Pipe.TryMoveNext()); + Assert.AreEqual(BackendType.BindComplete, decoder.Pipe.Current.Header.Type); + Assert.IsTrue(decoder.Pipe.TryMoveNext()); + Assert.AreEqual(BackendType.DataRow, decoder.Pipe.Current.Header.Type); + var body = decoder.Pipe.Current.OpenBodyReader(); + await pipe.Writer.WriteAsync(wire.AsMemory(initialLength)); + while (!body.IsComplete) + Assert.IsTrue(body.TryExtend()); + Assert.IsFalse(decoder.Pipe.TryMoveNext()); + Assert.IsTrue(await decoder.Pipe.MoveNextAsync(default)); + Assert.IsTrue(decoder.Pipe.TryMoveNext()); + Assert.AreEqual(BackendType.CommandComplete, decoder.Pipe.Current.Header.Type); await pipe.Writer.CompleteAsync(); - Assert.IsFalse(await e.MoveNextAsync()); - await e.DisposeAsync(); + await ((IAsyncDisposable)decoder).DisposeAsync(); } [TestMethod] @@ -876,50 +690,59 @@ public void BackendMessage_BufferedRequiresTagAndDeclaredLength() } [TestMethod] - public void BackendSegmenter_WaitsForUsefulPartialDataRowPrefix() + public void BackendCursor_WaitsForUsefulPartialDataRowPrefix() { var rowLength = 128 * 1024; var wire = BackendMessageBytes(BackendType.DataRow, rowLength); - var segmenter = new BackendMessageBatch.Segmenter(); var smallPrefix = new ReadOnlySequence(wire.AsMemory(0, 32)); - Assert.AreEqual(OperationStatus.NeedMoreData, - segmenter.CreateSegment(smallPrefix, out var length, out _)); - Assert.AreEqual(0, length); - Assert.AreEqual(BackendMessageBatch.Segmenter.DefaultDataRowStreamingThreshold, segmenter.MinimumSize); + var cursor = new BackendMessageCursor(smallPrefix); + Assert.IsFalse(cursor.TryReadNextInPlace(out _, out _, out _)); + Assert.AreEqual(BackendMessageCursor.DefaultDataRowStreamingThreshold, + cursor.RequiredBufferedLength); var usefulPrefix = new ReadOnlySequence( - wire.AsMemory(0, BackendMessageBatch.Segmenter.DefaultDataRowStreamingThreshold)); - Assert.AreEqual(OperationStatus.Done, - segmenter.CreateSegment(usefulPrefix, out length, out var batch)); - Assert.AreEqual(rowLength, length); - Assert.IsTrue(batch.TryReadNextInPlace(out var rowHeader, out var partialRow, out _)); + wire.AsMemory(0, BackendMessageCursor.DefaultDataRowStreamingThreshold)); + cursor = new(usefulPrefix); + Assert.IsTrue(cursor.TryReadNextInPlace(out var rowHeader, out var partialRow, out _)); Assert.AreEqual(BackendType.DataRow, rowHeader.Type); - Assert.AreEqual(BackendMessageBatch.Segmenter.DefaultDataRowStreamingThreshold, partialRow.Length); + Assert.AreEqual(BackendMessageCursor.DefaultDataRowStreamingThreshold, partialRow.Length); Assert.IsFalse(new BackendMessage(rowHeader, partialRow, new BackendMessageContext(), 0).Buffered); } [TestMethod] - public void BackendSegmenter_FramesUnknownMessageType() + public void BackendCursor_FramesUnknownMessageType() { var wire = BackendHeaderBytes((BackendType)(byte)'o', 4); - var segmenter = new BackendMessageBatch.Segmenter(); + var cursor = new BackendMessageCursor(new ReadOnlySequence(wire)); - Assert.AreEqual(OperationStatus.Done, - segmenter.CreateSegment(new ReadOnlySequence(wire), out var length, out var batch)); - Assert.AreEqual(wire.Length, length); - Assert.IsTrue(batch.TryReadNextInPlace(out var header, out _, out _)); + Assert.IsTrue(cursor.TryReadNextInPlace(out var header, out _, out _)); Assert.AreEqual((BackendType)(byte)'o', header.Type); } [TestMethod] - public void BackendSegmenter_RejectsMessageBeyondPostgreSqlAllocationLimit() + public void BackendCursor_FramesMemoryManagerBackedMessage() + { + using var manager = new TestMemoryManager( + BackendMessageBytes(BackendType.CommandComplete, 8)); + var cursor = new BackendMessageCursor( + new ReadOnlySequence(manager.Memory)); + + Assert.IsTrue(cursor.TryReadNextInPlace( + out var header, out var message, out var bufferedLength)); + Assert.AreEqual(BackendType.CommandComplete, header.Type); + Assert.AreEqual(8, message.Length); + Assert.AreEqual(8u, bufferedLength); + } + + [TestMethod] + public void BackendCursor_RejectsMessageBeyondPostgreSqlAllocationLimit() { var wire = BackendHeaderBytes(BackendType.DataRow, 0x3FFF_FFFF); - var segmenter = new BackendMessageBatch.Segmenter(); + var cursor = new BackendMessageCursor(new ReadOnlySequence(wire)); Assert.ThrowsExactly(() => - segmenter.CreateSegment(new ReadOnlySequence(wire), out _, out _)); + cursor.TryReadNextInPlace(out _, out _, out _)); } static byte[] BackendHeaderBytes(BackendType type, int length) diff --git a/Slon.Tests/Pg/CommandDrainTests.cs b/Slon.Tests/Pg/CommandDrainTests.cs index 08f2fc8..1db2d09 100644 --- a/Slon.Tests/Pg/CommandDrainTests.cs +++ b/Slon.Tests/Pg/CommandDrainTests.cs @@ -226,6 +226,7 @@ public async Task ConsumerDispose_MidBatch_BodyDrainsRemaining_ConnectionUsable( // wire usable (a hang shows as the "select 1" WaitAsync timing out, not a suite hang). [TestMethod] [DoNotParallelize] + [Ignore("Exercises the legacy body coroutine's open-before-park rendezvous.")] public async Task ConsumerDispose_MidBatch_SyncDispose_OpenBeforePark_Stress() { var iters = StressEnv.Iterations(fallback: 8, cap: 8_000); @@ -250,6 +251,7 @@ public async Task ConsumerDispose_MidBatch_SyncDispose_OpenBeforePark_Stress() // CompleteEnumeration, or CompleteEnumerationWithException), completing the body cross-thread while the pump is parked - the // sticky terminal publication must wake it. No other live-server test reaches this interleaving. [TestMethod] + [Ignore("Exercises the legacy body coroutine's in-flight completion/pump handoff race.")] public async Task InFlightCompletion_RacesSyncDispose_PumpNeverStrands_Stress() { // Each iteration is a full connect + force-abort cycle. Cap it because this is path coverage, @@ -339,6 +341,7 @@ public async Task StoppingToken_MidBatch_FollowedByDispose_FlowCompletesCleanly( // gate await is faulted by the heartbeat-driven OnStopping and the consumer's MoveNext // surfaces PgClientClosedException without any delivery. [TestMethod] + [Ignore("Requires the legacy body coroutine to read and publish a result before any consumer advances the flow.")] public async Task StoppingToken_PreFireAsync_BodyFaultsWithoutDelivery() { var protocol = await PgTestPool.NewIsolatedAsync(); diff --git a/Slon.Tests/Pg/CommandFlowBatchTests.cs b/Slon.Tests/Pg/CommandFlowBatchTests.cs new file mode 100644 index 0000000..d37792a --- /dev/null +++ b/Slon.Tests/Pg/CommandFlowBatchTests.cs @@ -0,0 +1,154 @@ +using Slon.Pg; +using Slon.Pg.Protocol; +using Slon.Pg.Protocol.Flows; +using Slon.Text; + +namespace Slon.Tests.Pg; + +[TestClass] +public class CommandFlowBatchTests +{ + static async Task Prepare(PgClientProtocol protocol, string sql, EncodedCString name) + { + var flow = protocol.Queue(new CommandFlow(async: true, + Command.Create(sql, commandName: name) with { DescribeOnly = true })); + var results = flow.GetAsyncEnumerator(); + CommandDescriptor descriptor = default; + while (await results.MoveNextAsync()) + descriptor = results.Current.GetMetadata().ToPreparedDescriptor(); + await results.DisposeAsync(); + return descriptor; + } + + [ConnectionCreatingTestMethod] + public async Task Batch_TwoPreparedCommands_EnumeratesResultsInOrder() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var first = await Prepare(protocol, "select 1::int4", "batch_1"); + var second = await Prepare(protocol, "select 2::int4", "batch_2"); + var results = protocol.Queue(new CommandFlow(async: true, new CommandList( + Command.Create(first), Command.Create(second)))).GetAsyncEnumerator(); + var values = new List(); + while (await results.MoveNextAsync()) + { + var rows = results.Current.GetAsyncEnumerator(); + while (await rows.MoveNextAsync()) + values.Add(rows.Current.GetValue(0)); + await rows.DisposeAsync(); + } + await results.DisposeAsync(); + CollectionAssert.AreEqual(new[] { 1, 2 }, values); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [ConnectionCreatingTestMethod] + public async Task Batch_DisposeBeforeRead_DrainsEveryCommand() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var first = await Prepare(protocol, "select generate_series(1, 20)", "batch_dispose_1"); + var second = await Prepare(protocol, "select generate_series(1, 20)", "batch_dispose_2"); + var results = protocol.Queue(new CommandFlow(async: true, new CommandList( + Command.Create(first), Command.Create(second)))).GetAsyncEnumerator(); + await results.DisposeAsync(); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [ConnectionCreatingTestMethod] + public async Task Batch_ErrorWithoutBarrier_DiscardsSuccessorThroughFinalSync() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var results = protocol.Queue(new CommandFlow(async: true, new CommandList( + Command.Create("select 1 / 0"), Command.Create("select 2::int4")))) + .GetAsyncEnumerator(); + Assert.IsTrue(await results.MoveNextAsync()); + var rows = results.Current.GetAsyncEnumerator(); + Assert.IsFalse(await rows.MoveNextAsync()); + await rows.DisposeAsync(); + Assert.ThrowsExactly(() => results.Current.GetCommandComplete()); + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [ConnectionCreatingTestMethod] + public async Task Batch_ErrorBarrier_AllowsFollowingCommand() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var results = protocol.Queue(new CommandFlow(async: true, new CommandList( + Command.Create("select 1 / 0") with { WithSync = true }, + Command.Create("select 2::int4")))) + .GetAsyncEnumerator(); + Assert.IsTrue(await results.MoveNextAsync()); + var failedRows = results.Current.GetAsyncEnumerator(); + Assert.IsFalse(await failedRows.MoveNextAsync()); + await failedRows.DisposeAsync(); + Assert.ThrowsExactly(() => results.Current.GetCommandComplete()); + Assert.IsTrue(await results.MoveNextAsync()); + var rows = results.Current.GetAsyncEnumerator(); + Assert.IsTrue(await rows.MoveNextAsync()); + Assert.AreEqual(2, rows.Current.GetValue(0)); + Assert.IsFalse(await rows.MoveNextAsync()); + await rows.DisposeAsync(); + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [ConnectionCreatingTestMethod] + public async Task Execution_SyncTwoPreparedCommands_EnumeratesResultsInOrder() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var first = await Prepare(protocol, "select 1::int4", "execution_sync_1"); + var second = await Prepare(protocol, "select 2::int4", "execution_sync_2"); + var results = protocol.Queue(new CommandFlow(async: false, new CommandList( + Command.Create(first), Command.Create(second)))).GetEnumerator(); + var values = new List(); + while (results.MoveNext()) + { + using var rows = results.Current.GetEnumerator(); + while (rows.MoveNext()) + values.Add(rows.Current.GetValue(0)); + } + results.Dispose(); + CollectionAssert.AreEqual(new[] { 1, 2 }, values); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [ConnectionCreatingTestMethod] + public async Task Execution_SyncDisposeBeforeRead_DrainsEveryCommand() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var first = await Prepare(protocol, "select generate_series(1, 20)", "execution_sync_dispose_1"); + var second = await Prepare(protocol, "select generate_series(1, 20)", "execution_sync_dispose_2"); + var results = protocol.Queue(new CommandFlow(async: false, new CommandList( + Command.Create(first), Command.Create(second)))).GetEnumerator(); + results.Dispose(); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [ConnectionCreatingTestMethod] + public async Task Execution_SyncErrorBarrier_AllowsFollowingCommand() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var results = protocol.Queue(new CommandFlow(async: false, new CommandList( + Command.Create("select 1 / 0") with { WithSync = true }, + Command.Create("select 2::int4")))).GetEnumerator(); + Assert.IsTrue(results.MoveNext()); + using (var failedRows = results.Current.GetEnumerator()) + Assert.IsFalse(failedRows.MoveNext()); + Assert.ThrowsExactly(() => results.Current.GetCommandComplete()); + Assert.IsTrue(results.MoveNext()); + using (var rows = results.Current.GetEnumerator()) + { + Assert.IsTrue(rows.MoveNext()); + Assert.AreEqual(2, rows.Current.GetValue(0)); + Assert.IsFalse(rows.MoveNext()); + } + Assert.IsFalse(results.MoveNext()); + results.Dispose(); + await PgTestPool.RunAsync(protocol, "select 3"); + } + +} + + diff --git a/Slon.Tests/Pg/CommandResultCollectTests.cs b/Slon.Tests/Pg/CommandResultCollectTests.cs new file mode 100644 index 0000000..62c3b5d --- /dev/null +++ b/Slon.Tests/Pg/CommandResultCollectTests.cs @@ -0,0 +1,117 @@ +using Slon.Pg; +using Slon.Pg.Protocol; +using Slon.Pg.Protocol.Flows; + +namespace Slon.Tests.Pg; + +[TestClass] +public sealed class CommandResultCollectTests +{ + [ConnectionCreatingTestMethod] + public async Task CollectsValueRowsAndLeavesWireReusable() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var values = new List(); + var results = protocol.Queue(new CommandFlow(async: true, + Command.Create("select generate_series(1, 100)"))).GetAsyncEnumerator(); + + Assert.IsTrue(await results.MoveNextAsync()); + await results.Current.CollectAsync(values, + static (items, row) => items.Add(row.GetInt32(0))); + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + + CollectionAssert.AreEqual(Enumerable.Range(1, 100).ToArray(), values); + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + public async Task BuffersAStreamingRowBeforeCallingCollector() + { + await using var protocol = await PgTestPool.NewIsolatedAsync( + options => options.DataRowStreamingThreshold = 1); + var values = new List<(int Id, string Text)>(); + var results = protocol.Queue(new CommandFlow(async: true, + Command.Create("select 42, repeat('x', 100000)"))).GetAsyncEnumerator(); + + Assert.IsTrue(await results.MoveNextAsync()); + await results.Current.CollectAsync(values, + static (items, row) => items.Add((row.GetInt32(0), row.GetValue(1)))); + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + + Assert.AreEqual(1, values.Count); + Assert.AreEqual(42, values[0].Id); + Assert.AreEqual(100000, values[0].Text.Length); + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + public async Task ResultBufferingRetainsCollectedRowsPastTheBufferingBarrier() + { + await using var protocol = await PgTestPool.NewIsolatedAsync( + options => options.DataRowStreamingThreshold = 1); + var values = new List<(int Id, ReadOnlyMemory Message)>(); + var results = protocol.Queue(new CommandFlow(async: true, Command.Create( + "select i, i::text || repeat('x', 20000) " + + "from generate_series(1, 3) as i"))).GetAsyncEnumerator(); + + Assert.IsTrue(await results.MoveNextAsync()); + results.Current.EnableResultBuffering(); + await results.Current.CollectAsync(values, + static (rows, row) => rows.Add(( + row.GetInt32(0), row.BorrowFieldMemory(1)))); + + Assert.AreEqual(3, values.Count); + for (var i = 0; i < values.Count; i++) + { + Assert.AreEqual(i + 1, values[i].Id); + Assert.AreEqual((i + 1) + new string('x', 20000), + System.Text.Encoding.UTF8.GetString(values[i].Message.Span)); + } + + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + public async Task CollectorFailureDrainsBeforeRethrowing() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var callbacks = 0; + var results = protocol.Queue(new CommandFlow(async: true, + Command.Create("select generate_series(1, 100)"))).GetAsyncEnumerator(); + + Assert.IsTrue(await results.MoveNextAsync()); + var exception = await Assert.ThrowsExactlyAsync(async () => + await results.Current.CollectAsync(0, (_, _) => + { + callbacks++; + throw new InvalidOperationException("collector failure"); + })); + Assert.AreEqual("collector failure", exception.Message); + Assert.AreEqual(1, callbacks); + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + public async Task CannotCollectAfterRowEnumerationStarted() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var results = protocol.Queue(new CommandFlow(async: true, + Command.Create("select generate_series(1, 2)"))).GetAsyncEnumerator(); + + Assert.IsTrue(await results.MoveNextAsync()); + var rows = results.Current.GetAsyncEnumerator(); + Assert.IsTrue(await rows.MoveNextAsync()); + await Assert.ThrowsExactlyAsync(async () => + await results.Current.CollectAsync(0, static (_, _) => { })); + await rows.DisposeAsync(); + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + } +} diff --git a/Slon.Tests/Pg/CommandResultEnumerationTests.cs b/Slon.Tests/Pg/CommandResultEnumerationTests.cs index 10a52c5..f8ebde9 100644 --- a/Slon.Tests/Pg/CommandResultEnumerationTests.cs +++ b/Slon.Tests/Pg/CommandResultEnumerationTests.cs @@ -1,3 +1,4 @@ +using System.Text; using Slon.Pg; using Slon.Pg.Protocol; using Slon.Pg.Protocol.Flows; @@ -7,6 +8,111 @@ namespace Slon.Tests.Pg; [TestClass] public class CommandResultEnumerationTests { + [ConnectionCreatingTestMethod] + public async Task ResultBuffering_CannotBeginAfterRowEnumeration() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(new CommandFlow( + async: true, Command.Create("select generate_series(1, 2)"))); + var results = flow.GetAsyncEnumerator(); + + Assert.IsTrue(await results.MoveNextAsync()); + var result = results.Current; + var rows = result.GetAsyncEnumerator(); + Assert.IsTrue(await rows.MoveNextAsync()); + Assert.ThrowsExactly( + result.EnableResultBuffering); + + await rows.DisposeAsync(); + await results.DisposeAsync(); + } + + [ConnectionCreatingTestMethod] + public async Task ContiguousFieldMemory_RemainsValidAcrossExtendedBatches() + { + const int rowCount = 2000; + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(new CommandFlow(async: true, Command.Create( + $"select i, i::text || repeat('x', 96) from generate_series(1, {rowCount}) as i"))); + var results = flow.GetAsyncEnumerator(); + var values = new List<(int Id, ReadOnlyMemory Message)>(rowCount); + + try + { + Assert.IsTrue(await results.MoveNextAsync()); + results.Current.EnableResultBuffering(); + var rows = results.Current.GetAsyncEnumerator(); + while (await rows.MoveNextAsync()) + { + var reader = rows.Current.GetReader(); + values.Add((reader.Read(), reader.ReadBorrowedMemory())); + } + await rows.DisposeAsync(); + + Assert.AreEqual(rowCount, values.Count); + foreach (var (id, message) in values) + Assert.AreEqual(id + new string('x', 96), + Encoding.UTF8.GetString(message.Span)); + } + finally + { + await results.DisposeAsync(); + } + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + public async Task ContiguousFieldMemory_BuffersStreamingRowsIntoTheRetainedBatch() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(new CommandFlow(async: true, Command.Create( + "select i, i::text || repeat('x', 20000) from generate_series(1, 3) as i"))); + var results = flow.GetAsyncEnumerator(); + var values = new List>(); + + try + { + Assert.IsTrue(await results.MoveNextAsync()); + results.Current.EnableResultBuffering(); + var rows = results.Current.GetAsyncEnumerator(); + while (await rows.MoveNextAsync()) + values.Add(rows.Current.BorrowFieldMemory(1)); + await rows.DisposeAsync(); + + Assert.AreEqual(3, values.Count); + for (var i = 0; i < values.Count; i++) + Assert.AreEqual((i + 1) + new string('x', 20000), + Encoding.UTF8.GetString(values[i].Span)); + } + finally + { + await results.DisposeAsync(); + } + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + public async Task ResultBuffering_AbandonmentReleasesTheReadGrant() + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(new CommandFlow(async: true, Command.Create( + "select i, i::text || repeat('x', 96) from generate_series(1, 2000) as i"))); + var results = flow.GetAsyncEnumerator(); + + Assert.IsTrue(await results.MoveNextAsync()); + results.Current.EnableResultBuffering(); + var rows = results.Current.GetAsyncEnumerator(); + Assert.IsTrue(await rows.MoveNextAsync()); + var borrowed = rows.Current.BorrowFieldMemory(1); + Assert.IsFalse(borrowed.IsEmpty); + + await rows.DisposeAsync(); + await results.DisposeAsync(); + await PgTestPool.RunAsync(protocol, "select 1"); + } + [ConnectionCreatingTestMethod] public async Task DescribeOnlyErrorSurfacesWhenInspectingTheResult() { @@ -22,6 +128,7 @@ public async Task DescribeOnlyErrorSurfacesWhenInspectingTheResult() } [ConnectionCreatingTestMethod] + [Ignore("The replacement fixes the execution/consumption mode for a flow tenure; switching an async flow to synchronous driving is a legacy body-rendezvous behavior.")] public async Task AsyncFlow_CanSwitchToSynchronousResultAdvancement() { await using var protocol = await PgTestPool.NewIsolatedAsync(); @@ -129,33 +236,4 @@ public async Task ErrorWithoutSync_ResumesAfterInternalSync(bool async) await PgTestPool.RunAsync(protocol, "select 1"); } - // Reset must clear the terminal enumeration state, or a reused flow reports exhaustion before - // its next tenure publishes anything. - [ConnectionCreatingTestMethod] - public async Task Reset_ClearsEnumerationCompleted_ForNextTenure() - { - await using var protocol = await PgTestPool.NewIsolatedAsync(); - var flow = new ResettableCommandFlow(async: true, Command.Create("select 1")); - for (var tenure = 0; tenure < 2; tenure++) - { - if (tenure > 0) - { - flow.Reset(); - flow.Initialize(async: true, Command.Create("select 2")); - } - protocol.Queue(flow); - var e = flow.GetAsyncEnumerator(); - Assert.IsTrue(await e.MoveNextAsync(), $"tenure {tenure} must publish its result"); - await e.Current.DisposeAsync(); - Assert.IsFalse(await e.MoveNextAsync()); - await e.DisposeAsync(); - } - } - - // Pooling a timeout-armed flow is refused by Reset. Opt out so the reset path itself is testable. - sealed class ResettableCommandFlow(bool async, params ReadOnlySpan commands) - : CommandFlow(async, commands) - { - protected override bool EnableActivationTimeout => false; - } } diff --git a/Slon.Tests/Pg/CommandUserCancellationTests.cs b/Slon.Tests/Pg/CommandUserCancellationTests.cs index 3860186..8f1c613 100644 --- a/Slon.Tests/Pg/CommandUserCancellationTests.cs +++ b/Slon.Tests/Pg/CommandUserCancellationTests.cs @@ -1,3 +1,4 @@ +using Microsoft.Extensions.Time.Testing; using Slon.Pg; using Slon.Pg.Protocol; using Slon.Pg.Protocol.Flows; @@ -128,6 +129,40 @@ public async Task UserCt_FiresMidRead_SurfacesOce_ProtocolUsable() await PgTestPool.RunAsync(protocol, "select 1"); } + // Autonomous execution may enter the command read before the consumer supplies its per-read + // token. The late token must still arm backend cancellation for that active read. + [TestMethod] + [Ignore("Requires the legacy body to enter a read before the consumer supplies its token; the replacement consumer owns the read from entry.")] + public async Task UserCt_SuppliedAfterReadStarted_RequestsCancellation_ProtocolUsable() + { + await using var blocker = await PgAdvisoryLock.AcquireAsync(); + var cancelRequested = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); + await using var protocol = await PgTestPool.NewIsolatedAsync( + o => o.CancelSender = (processId, secretKey, token) => + { + cancelRequested.TrySetResult(); + return new(CancelRequestState.NotSent); + }); + + var flow = new CommandFlow(async: true, blocker.WaitCommand); + Assert.IsTrue(protocol.TryQueue(flow)); + await blocker.WaitUntilContendedAsync(protocol.FlowControl.BackendProcessId); + + using var cts = new CancellationTokenSource(); + var enumerator = flow.GetAsyncEnumerator(); + var moveNext = enumerator.MoveNextAsync(cts.Token).AsTask(); + cts.Cancel(); + + await cancelRequested.Task.WaitAsync(TimeSpan.FromSeconds(5)); + await blocker.ReleaseAsync(); + var exception = await Assert.ThrowsExactlyAsync( + async () => await moveNext.WaitAsync(TimeSpan.FromSeconds(5))); + Assert.AreEqual(cts.Token, exception.CancellationToken); + await enumerator.DisposeAsync(); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + [TestMethod] public async Task UserCt_ThenProtocolClose_DuringDrain_CompletesPendingMoveNext() { @@ -538,6 +573,60 @@ public async Task ConsumerDispose_UsesItsOwnGraceBeforeSideChannelAttempt() } [TestMethod] + [DataRow(false, DisplayName = "explicit cancellation")] + [DataRow(true, DisplayName = "consumer disposal")] + public async Task ServerCancel_SubsequentWindowDispatchesWithoutHeartbeat(bool dispose) + { + var time = new FakeTimeProvider(); + await using var blocker = await PgAdvisoryLock.AcquireAsync(); + var settleAttempt = new TaskCompletionSource( + TaskCreationOptions.RunContinuationsAsynchronously); + await using var protocol = await PgTestPool.NewIsolatedAsync(o => + { + o.TimeProvider = time; + o.HeartbeatInterval = TimeSpan.FromHours(1); + o.CancelRequestDelay = TimeSpan.Zero; + o.CancelSender = (_, _, _) => new(settleAttempt.Task); + }); + + var flow = new CommandFlow(async: true, + Command.Create("select 1") with { WithSync = true }, + blocker.WaitCommand); + Assert.IsTrue(protocol.TryQueue(flow)); + var enumerator = flow.GetAsyncEnumerator(); + Assert.IsTrue(await enumerator.MoveNextAsync()); + + var cancellation = dispose + ? enumerator.DisposeAsync().AsTask() + : flow.CancelAsync(); + try + { + await blocker.WaitUntilContendedAsync(protocol.FlowControl.BackendProcessId); + while (flow.CancellationWindow < 1) + await Task.Yield(); + + var state = ProtocolDiag.CancellationState(protocol); + Assert.IsTrue(state.StartsWith("dispatching=True", StringComparison.Ordinal), + $"The successor cancellation did not dispatch at its read frontier: {state}"); + } + finally + { + settleAttempt.TrySetResult(CancelRequestState.Sent); + await blocker.ReleaseAsync(); + } + await cancellation; + if (!dispose) + { + await Assert.ThrowsExactlyAsync( + async () => await enumerator.MoveNextAsync()); + await enumerator.DisposeAsync(); + } + await WaitUntilAsync(() => !protocol.HasPendingCancellation); + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [TestMethod] + [Ignore("Requires a second body-owned drain read after the consumer read times out; the replacement retains one read owner through drain.")] public async Task ServerCancel_ReadTimeoutAfterAmbiguousRetryAbortsWire() { var iterations = StressEnv.Iterations(fallback: 1, cap: 5_000); diff --git a/Slon.Tests/Pg/ConcurrentHandoffInMemoryTests.cs b/Slon.Tests/Pg/ConcurrentHandoffInMemoryTests.cs index 36a61e3..5f3c763 100644 --- a/Slon.Tests/Pg/ConcurrentHandoffInMemoryTests.cs +++ b/Slon.Tests/Pg/ConcurrentHandoffInMemoryTests.cs @@ -250,8 +250,8 @@ public override void WaitUntilWritable(TimeSpan timeout) { } public EchoServerTransport(byte[] handshake, byte[] response) { _response = response; - // Slon's PipeSegmentEnumerator requires its OWN StreamPipeReader for the sync path - // (not the BCL PipeReader.Create). Wrap the in-memory pipe as a stream, same as the socket transport. + // Slon's synchronous protocol path requires a StreamPipeReader (not the BCL + // PipeReader.Create). Wrap the in-memory pipe as a stream, like the socket transport. _clientReader = new Slon.Pipelines.DefaultStreamPipeReader( _toClient.Reader.AsStream(), new StreamPipeReaderOptions(bufferSize: 8192, useZeroByteReads: false), diff --git a/Slon.Tests/Pg/ExclusiveAccessFlowTests.cs b/Slon.Tests/Pg/ExclusiveAccessFlowTests.cs index 184353a..23dfdd6 100644 --- a/Slon.Tests/Pg/ExclusiveAccessFlowTests.cs +++ b/Slon.Tests/Pg/ExclusiveAccessFlowTests.cs @@ -25,7 +25,7 @@ public async Task BindingFailure_FaultsFlowWithoutCondemningWire() protocol.SetFlowBindingContext(new BindingProbeContext("wire")); var flow = protocol.Queue(new BindingProbeFlow(fail: true)); - await Assert.ThrowsExactlyAsync(async () => await DrainAsync(flow)); + await Assert.ThrowsExactlyAsync(async () => await DrainBindingProbeAsync(flow)); Assert.AreEqual(1, flow.BindCount); await DrainAsync(protocol.Queue(new CommandFlow(async: true, Command.Create("select 1")))); @@ -37,6 +37,9 @@ static async Task DrainAsync(CommandFlow flow) await e.DisposeAsync(); } + static async Task DrainBindingProbeAsync(BindingProbeFlow flow) + => await flow.WaitForComplete(); + [TestMethod] public async Task Scope_RoundTrip_RunsCommandOnInnerPipeline() { @@ -86,7 +89,7 @@ public async Task LongRunningScope_RejectsOuterAdmissionUntilReleased() Assert.IsTrue(protocol.IsSchedulable); Assert.IsTrue(protocol.TryQueue(bindProbe)); - await DrainAsync(bindProbe); + await DrainBindingProbeAsync(bindProbe); Assert.AreEqual(1, bindProbe.BindCount, "the portable flow must bind exactly once when its later placement reaches dispatch"); } diff --git a/Slon.Tests/Pg/FlowAuthoring/AutonomousFlowAuthoringTests.cs b/Slon.Tests/Pg/FlowAuthoring/AutonomousFlowAuthoringTests.cs new file mode 100644 index 0000000..3ff8d78 --- /dev/null +++ b/Slon.Tests/Pg/FlowAuthoring/AutonomousFlowAuthoringTests.cs @@ -0,0 +1,135 @@ +using Slon.Pg; +using Slon.Pg.Protocol; +using Slon.Pg.Protocol.Flows; + +namespace Slon.Tests.Pg.FlowAuthoring; + +public interface IAutonomousFlowContract +{ + string Name { get; } + PgClientFlow Create(bool async, params Command[] commands); + bool WasReleased(PgClientFlow flow); +} + +sealed class AutonomousCommandFlowContract : IAutonomousFlowContract +{ + public static AutonomousCommandFlowContract Instance { get; } = new(); + public string Name => nameof(AutonomousCommandFlow); + + public PgClientFlow Create(bool async, params Command[] commands) + => new AutonomousCommandFlow(async, commands); + + public bool WasReleased(PgClientFlow flow) + => ((AutonomousCommandFlow)flow).Released; + + public override string ToString() => Name; +} + +sealed class AutonomousCommandFlow : PgClientFlow +{ + readonly CommandList _commands; + + internal AutonomousCommandFlow(bool async, params Command[] commands) + : base(supportsDeferredFlush: true) + { + IsAsync = async; + _commands = new(commands); + } + + internal bool Released { get; private set; } + + protected override ValueTask ExecuteAuto(Context context) + { + var write = _commands.WriteCommandsAsync(context.GetEncoder(), appendSync: true); + return new(new FlowTasks(write, DrainAsync(context))); + + static async ValueTask DrainAsync(Context context) + { + var decoder = await context.GetDecoderAsync().ConfigureAwait(false); + while (context.OutstandingRfqCount is not 0) + _ = await decoder.GetNextAsync().ConfigureAwait(false); + } + } + + protected override void OnReleasing(Exception? exception) => Released = true; +} + +[TestClass] +public class AutonomousFlowAuthoringTests : ConnectionCreatingTest +{ + public static IEnumerable Implementations + { + get + { + yield return [AutonomousCommandFlowContract.Instance]; + } + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task NaturalCompletion_ReleasesAfterPipelineTask( + IAutonomousFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(contract.Create(async: true, + Command.Create("select 1"), Command.Create("select 2"))); + + await flow.WaitForComplete(); + + Assert.IsTrue(contract.WasReleased(flow)); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task SynchronousMode_DoesNotRequireCallerHandoff( + IAutonomousFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(contract.Create(async: false, Command.Create("select 1"))); + + await flow.WaitForComplete(); + + Assert.IsFalse(flow.NeedsSyncHandoff); + await PgTestPool.RunAsync(protocol, "select 2"); + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task ForcefulAbort_CompletesWithoutExternalDriver( + IAutonomousFlowContract contract) + { + await using var blocker = await PgAdvisoryLock.AcquireAsync(); + var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(contract.Create(async: true, blocker.WaitCommand)); + try + { + await blocker.WaitUntilContendedAsync(protocol.FlowControl.BackendProcessId); + await protocol.CompleteAsync(new IOException("contract abort")); + await Assert.ThrowsExactlyAsync( + async () => await flow.WaitForComplete()); + } + finally + { + await blocker.ReleaseAsync(); + await protocol.DisposeAsync(); + } + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task PipelinedInstances_CompleteAndRelease( + IAutonomousFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var first = protocol.Queue(contract.Create(async: true, Command.Create("select 1"))); + var second = protocol.Queue(contract.Create(async: true, Command.Create("select 2"))); + + await first.WaitForComplete(); + await second.WaitForComplete(); + + Assert.IsTrue(contract.WasReleased(first)); + Assert.IsTrue(contract.WasReleased(second)); + await PgTestPool.RunAsync(protocol, "select 3"); + } +} diff --git a/Slon.Tests/Pg/FlowAuthoring/ConsumerDrivenFlowAuthoringTests.cs b/Slon.Tests/Pg/FlowAuthoring/ConsumerDrivenFlowAuthoringTests.cs new file mode 100644 index 0000000..a421c5c --- /dev/null +++ b/Slon.Tests/Pg/FlowAuthoring/ConsumerDrivenFlowAuthoringTests.cs @@ -0,0 +1,238 @@ +using Draghi.Pipelining; +using Slon.Pg; +using Slon.Pg.Protocol; +using Slon.Pg.Protocol.Flows; + +namespace Slon.Tests.Pg.FlowAuthoring; + +public interface IConsumerDrivenFlowContract +{ + string Name { get; } + PgClientFlow Create(bool async, params Command[] commands); + IEnumerator GetEnumerator(PgClientFlow flow); + IAsyncEnumerator GetAsyncEnumerator(PgClientFlow flow); + ValueTask MoveNextAsync( + IAsyncEnumerator results, CancellationToken cancellationToken); +} + +public interface IReusableConsumerDrivenFlowContract : IConsumerDrivenFlowContract +{ + PgClientFlow CreateReusable(bool async, params Command[] commands); + void Reset(PgClientFlow flow, bool async, params Command[] commands); +} + +sealed class CommandFlowContract : IReusableConsumerDrivenFlowContract +{ + public static CommandFlowContract Instance { get; } = new(); + public string Name => nameof(CommandFlow); + + public PgClientFlow Create(bool async, params Command[] commands) + => new CommandFlow(async, commands); + + public IEnumerator GetEnumerator(PgClientFlow flow) + => ((CommandFlow)flow).GetEnumerator(); + + public IAsyncEnumerator GetAsyncEnumerator(PgClientFlow flow) + => ((CommandFlow)flow).GetAsyncEnumerator(); + + public ValueTask MoveNextAsync( + IAsyncEnumerator results, CancellationToken cancellationToken) + => ((CommandFlow.Enumerator)results).MoveNextAsync(cancellationToken); + + public PgClientFlow CreateReusable(bool async, params Command[] commands) + => new CommandFlow(async, enableActivationTimeout: false, commands); + + public void Reset(PgClientFlow flow, bool async, params Command[] commands) + { + flow.Reset(); + ((CommandFlow)flow).Initialize(async, commands); + } + + public override string ToString() => Name; +} + +[TestClass] +public class ConsumerDrivenFlowAuthoringTests : ConnectionCreatingTest +{ + public static IEnumerable Implementations + { + get + { + yield return [CommandFlowContract.Instance]; + } + } + + public static IEnumerable ReusableImplementations + { + get + { + yield return [CommandFlowContract.Instance]; + } + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task AsyncNaturalExhaustion_ReleasesWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(contract.Create(async: true, + Command.Create("select 1"), Command.Create("select 2"))); + + var results = contract.GetAsyncEnumerator(flow); + while (await results.MoveNextAsync()) + await results.Current.DisposeAsync(); + await results.DisposeAsync(); + + await flow.WaitForComplete(); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task SyncNaturalExhaustion_ReleasesWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(contract.Create(async: false, + Command.Create("select 1"), Command.Create("select 2"))); + + using (var results = contract.GetEnumerator(flow)) + { + while (results.MoveNext()) + results.Current.Dispose(); + } + + await flow.WaitForComplete(); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task AsyncAbandonBeforeRead_DrainsAndReleasesWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(contract.Create(async: true, + Command.Create("select generate_series(1, 20)"), Command.Create("select 2"))); + + await contract.GetAsyncEnumerator(flow).DisposeAsync(); + + await flow.WaitForComplete(); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task AsyncAbandonAfterPublication_DrainsAndReleasesWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(contract.Create(async: true, + Command.Create("select generate_series(1, 20)"), Command.Create("select 2"))); + var results = contract.GetAsyncEnumerator(flow); + + Assert.IsTrue(await results.MoveNextAsync()); + await results.DisposeAsync(); + + await flow.WaitForComplete(); + await PgTestPool.RunAsync(protocol, "select 3"); + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task AbortWithoutConsumer_CompletesFlow(IConsumerDrivenFlowContract contract) + { + await using var blocker = await PgAdvisoryLock.AcquireAsync(); + var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = protocol.Queue(contract.Create(async: true, blocker.WaitCommand)); + try + { + await blocker.WaitUntilContendedAsync(protocol.FlowControl.BackendProcessId); + await protocol.CompleteAsync(new IOException("contract abort")); + await Assert.ThrowsExactlyAsync( + async () => await flow.WaitForComplete()); + } + finally + { + await blocker.ReleaseAsync(); + await protocol.DisposeAsync(); + } + } + + [TestMethod] + [DynamicData(nameof(Implementations))] + public async Task MixedSyncAndAsync_NeverRunsConsumerOnExecutor( + IConsumerDrivenFlowContract contract) + { + var executionScheduler = new TrackingScheduler(); + var activationScheduler = new TrackingScheduler(); + await using var protocol = await PgTestPool.NewIsolatedAsync(o => + { + o.ExecutionScheduler = executionScheduler; + o.ActivationScheduler = activationScheduler; + }); + Exception? failure = null; + void Capture(Exception exception) + => Interlocked.CompareExchange(ref failure, exception, null); + + var asyncLoop = Task.Run(async () => + { + try + { + for (var i = 0; i < StressEnv.Iterations(32, 8_000) + && Volatile.Read(ref failure) is null; i++) + { + var flow = protocol.Queue(contract.Create(async: true, Command.Create("select 1"))); + var results = contract.GetAsyncEnumerator(flow); + var hasResult = await results.MoveNextAsync(); + if (hasResult && executionScheduler.IsExecuting) + Capture(new InvalidOperationException( + $"{contract.Name} resumed its async consumer on the pipeline executor.")); + while (hasResult) + hasResult = await results.MoveNextAsync(); + await results.DisposeAsync(); + } + } + catch (Exception ex) { Capture(ex); } + }); + + var syncThread = new Thread(() => + { + try + { + for (var i = 0; i < StressEnv.Iterations(32, 8_000) + && Volatile.Read(ref failure) is null; i++) + { + var flow = protocol.Queue(contract.Create(async: false, Command.Create("select 1"))); + using var results = contract.GetEnumerator(flow); + while (results.MoveNext()) { } + } + } + catch (Exception ex) { Capture(ex); } + }) { IsBackground = true, Name = $"{contract.Name}-sync-contract" }; + + syncThread.Start(); + await asyncLoop; + await Task.Run(syncThread.Join); + if (failure is not null) + Assert.Fail($"{contract.Name} violated mixed-consumer execution: {failure}"); + } + + sealed class TrackingScheduler : PipelineScheduler + { + [ThreadStatic] + static TrackingScheduler? _executing; + + internal bool IsExecuting => ReferenceEquals(_executing, this); + + public override void SubmitDetached( + Action action, object? state, bool preferLocal = true) + => PipelineScheduler.ThreadPool.SubmitDetached(static state => + { + var work = (Work)state!; + var prior = _executing; + _executing = work.Scheduler; + try { work.Action(work.State); } + finally { _executing = prior; } + }, new Work(this, action, state), preferLocal); + + sealed record Work(TrackingScheduler Scheduler, Action Action, object? State); + } +} diff --git a/Slon.Tests/Pg/FlowAuthoring/ConsumerDrivenFlowSemanticContractTests.cs b/Slon.Tests/Pg/FlowAuthoring/ConsumerDrivenFlowSemanticContractTests.cs new file mode 100644 index 0000000..fe8f054 --- /dev/null +++ b/Slon.Tests/Pg/FlowAuthoring/ConsumerDrivenFlowSemanticContractTests.cs @@ -0,0 +1,431 @@ +using Slon.Pg; +using Slon.Pg.Protocol; +using Slon.Pg.Protocol.Flows; +using Slon.Pg.Serialization; +using Slon.Pg.Types; +using Slon.Text; + +namespace Slon.Tests.Pg.FlowAuthoring; + +// Command-result semantics shared by consumer-driven PostgreSQL flows. +[TestClass] +public class ConsumerDrivenFlowSemanticContractTests +{ + public static IEnumerable Implementations + => ConsumerDrivenFlowAuthoringTests.Implementations; + + public static IEnumerable ReusableImplementations + => ConsumerDrivenFlowAuthoringTests.ReusableImplementations; + + public static IEnumerable PreparedCases + { + get + { + foreach (var implementation in Implementations) + { + yield return [implementation[0], 0]; + yield return [implementation[0], 3]; + } + } + } + static async Task Prepare( + IConsumerDrivenFlowContract contract, + PgClientProtocol protocol, string sql, EncodedCString name) + { + var results = Queue(contract, protocol, + Command.Create(sql, commandName: name) with { DescribeOnly = true }); + CommandDescriptor descriptor = default; + while (await results.MoveNextAsync()) + descriptor = results.Current.GetMetadata().ToPreparedDescriptor(); + await results.DisposeAsync(); + return descriptor; + } + + static Results Queue(IConsumerDrivenFlowContract contract, + PgClientProtocol protocol, in Command command, + CancellationToken cancellationToken = default) + { + var flow = protocol.Queue(contract.Create(async: true, command), cancellationToken); + return new(contract, contract.GetAsyncEnumerator(flow)); + } + + static async Task CountRows(Results results) + { + var count = 0; + while (await results.MoveNextAsync()) + { + var rows = results.Current.GetAsyncEnumerator(); + while (await rows.MoveNextAsync()) + count++; + await rows.DisposeAsync(); + } + await results.DisposeAsync(); + return count; + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(PreparedCases))] + public async Task Prepared_NaturalExhaustion( + IConsumerDrivenFlowContract contract, int rowCount) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var descriptor = await Prepare(contract, protocol, + $"select generate_series(1, {rowCount})", $"contract_rows_{rowCount}"); + + Assert.AreEqual(rowCount, + await CountRows(Queue(contract, protocol, Command.Create(descriptor)))); + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task Unprepared_NaturalExhaustion(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + + Assert.AreEqual(2, await CountRows(Queue(contract, protocol, + Command.Create("select generate_series(1, 2)")))); + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task DisposeBeforeAnyRead_DrainsAndKeepsWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var descriptor = await Prepare(contract, protocol, + "select generate_series(1, 1000)", "contract_unread"); + var results = Queue(contract, protocol, Command.Create(descriptor)); + + await results.DisposeAsync(); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task DisposeAfterOneRow_DrainsAndKeepsWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var descriptor = await Prepare(contract, protocol, + "select generate_series(1, 20000)", "contract_partial"); + var results = Queue(contract, protocol, Command.Create(descriptor)); + + Assert.IsTrue(await results.MoveNextAsync()); + var rows = results.Current.GetAsyncEnumerator(); + Assert.IsTrue(await rows.MoveNextAsync()); + await results.DisposeAsync(); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task CommandError_IsResultAndKeepsWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var descriptor = await Prepare(contract, protocol, "select 1 / 0", "contract_error"); + var results = Queue(contract, protocol, Command.Create(descriptor)); + + Assert.IsTrue(await results.MoveNextAsync()); + var failed = results.Current; + var rows = failed.GetAsyncEnumerator(); + Assert.IsFalse(await rows.MoveNextAsync()); + await rows.DisposeAsync(); + Assert.ThrowsExactly(() => failed.GetCommandComplete()); + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task PreparedMetadataAndCompletion_Agree(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var descriptor = await Prepare(contract, protocol, "select 42::int4", "contract_metadata"); + var results = Queue(contract, protocol, Command.Create(descriptor)); + + Assert.IsTrue(await results.MoveNextAsync()); + var result = results.Current; + var metadata = result.GetMetadata(); + Assert.IsTrue(metadata.IsPrepared); + Assert.AreEqual(descriptor.CommandName, metadata.CommandName); + Assert.IsNotNull(metadata.RowDescription); + var rows = result.GetAsyncEnumerator(); + Assert.IsTrue(await rows.MoveNextAsync()); + Assert.AreEqual(42, rows.Current.GetReader().Read()); + Assert.IsFalse(await rows.MoveNextAsync()); + await rows.DisposeAsync(); + Assert.IsTrue(result.IsComplete); + Assert.AreEqual(StatementType.Select, result.GetCommandComplete().StatementType); + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task CancellationWhileReadPending_DeliversTokenAndKeepsWire(IConsumerDrivenFlowContract contract) + { + await using var blocker = await PgAdvisoryLock.AcquireAsync(); + await using var protocol = await NewCancelableProtocolAsync(); + using var cancellation = new CancellationTokenSource(); + var results = Queue(contract, protocol, blocker.WaitCommand); + + var pending = results.MoveNextAsync(cancellation.Token); + Assert.IsFalse(pending.IsCompleted); + await blocker.WaitUntilContendedAsync(protocol.FlowControl.BackendProcessId); + cancellation.Cancel(); + var exception = await Assert.ThrowsExactlyAsync( + async () => await pending); + Assert.AreEqual(cancellation.Token, exception.CancellationToken); + await Assert.ThrowsExactlyAsync( + async () => await results.MoveNextAsync()); + await results.DisposeAsync(); + await blocker.ReleaseAsync(); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task CancellationAfterRow_DrainsAndKeepsWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await NewCancelableProtocolAsync(); + var descriptor = await Prepare(contract, protocol, + "select generate_series(1, 20000)", "contract_cancel_row"); + using var cancellation = new CancellationTokenSource(); + var results = Queue(contract, protocol, Command.Create(descriptor)); + + Assert.IsTrue(await results.MoveNextAsync(cancellation.Token)); + var rows = results.Current.GetAsyncEnumerator(); + Assert.IsTrue(await rows.MoveNextAsync()); + cancellation.Cancel(); + await Assert.ThrowsExactlyAsync( + async () => await results.MoveNextAsync(cancellation.Token)); + await results.DisposeAsync(); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task PreCancelledRead_ReleasesCallerAndKeepsWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await NewCancelableProtocolAsync(); + var descriptor = await Prepare(contract, protocol, + "select generate_series(1, 1000)", "contract_precancel"); + var results = Queue(contract, protocol, Command.Create(descriptor)); + var cancellationToken = new CancellationToken(canceled: true); + + var exception = await Assert.ThrowsExactlyAsync( + async () => await results.MoveNextAsync(cancellationToken)); + Assert.AreEqual(cancellationToken, exception.CancellationToken); + await Assert.ThrowsExactlyAsync( + async () => await results.MoveNextAsync(cancellationToken)); + await results.DisposeAsync(); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task SuccessorProgressesAfterAbandonment(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var descriptor = await Prepare(contract, protocol, + "select generate_series(1, 20000)", "contract_successor"); + var first = Queue(contract, protocol, Command.Create(descriptor)); + var second = Queue(contract, protocol, Command.Create(descriptor)); + + Assert.IsTrue(await first.MoveNextAsync()); + await first.DisposeAsync(); + Assert.AreEqual(20000, await CountRows(second)); + + await PgTestPool.RunAsync(protocol, "select 1"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task GracefulStopDrainsHeldResultAndFaultsConsumer(IConsumerDrivenFlowContract contract) + { + var protocol = await PgTestPool.NewIsolatedAsync(); + var descriptor = await Prepare(contract, protocol, + "select generate_series(1, 1000)", "contract_graceful"); + var results = Queue(contract, protocol, Command.Create(descriptor)); + Assert.IsTrue(await results.MoveNextAsync()); + + var complete = protocol.CompleteAsync(); + await protocol.Heartbeat(TimeSpan.Zero); + await complete; + await Assert.ThrowsAsync( + async () => await results.MoveNextAsync()); + await results.DisposeAsync(); + } + + [ConnectionCreatingTestMethod(connections: 2)] + [DynamicData(nameof(Implementations))] + public async Task BackendTermination_IsCollateral(IConsumerDrivenFlowContract contract) + { + await using var protocols = await PgTestPool.NewIsolatedProtocolsAsync(2); + var killer = protocols.Items[1]; + + var exception = await Terminate(); + Assert.IsInstanceOfType(exception); + + async Task Terminate() + { + await using var victim = await PgTestPool.NewIsolatedAsync(); + var pid = await ReadBackendPid(contract, victim); + var descriptor = await Prepare(contract, victim, "select pg_sleep(10)", + "contract_terminate_command"); + var results = Queue(contract, victim, Command.Create(descriptor)); + var pending = results.MoveNextAsync(); + Assert.IsFalse(pending.IsCompleted); + await PgTestPool.RunAsync(killer, $"select pg_terminate_backend({pid})"); + var exception = await Assert.ThrowsAsync(async () => await pending); + await Assert.ThrowsAsync(async () => await results.DisposeAsync()); + await victim.Completion; + return exception; + } + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(Implementations))] + public async Task TornTrailingWrite_RecoversWire(IConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var results = Queue(contract, protocol, TornStreamedBind()); + + Exception? observed = null; + try + { + while (await results.MoveNextAsync()) + { + var rows = results.Current.GetAsyncEnumerator(); + while (await rows.MoveNextAsync()) { } + await rows.DisposeAsync(); + results.Current.GetCommandComplete(); + } + } + catch (Exception exception) + { + observed = exception; + } + try + { + await results.DisposeAsync(); + } + catch (Exception exception) + { + observed ??= exception; + } + Assert.IsNotNull(observed); + + await PgTestPool.RunAsync(protocol, "select 42::int4"); + } + + [ConnectionCreatingTestMethod] + [DynamicData(nameof(ReusableImplementations))] + public async Task ResetAfterCompletion_StartsIndependentTenure( + IReusableConsumerDrivenFlowContract contract) + { + await using var protocol = await PgTestPool.NewIsolatedAsync(); + var flow = contract.CreateReusable(async: true, Command.Create("select 1")); + + for (var tenure = 0; tenure < 2; tenure++) + { + if (tenure > 0) + contract.Reset(flow, async: true, Command.Create("select 2")); + protocol.Queue(flow); + var results = contract.GetAsyncEnumerator(flow); + Assert.IsTrue(await results.MoveNextAsync(), + $"tenure {tenure} must publish its result"); + await results.Current.DisposeAsync(); + Assert.IsFalse(await results.MoveNextAsync()); + await results.DisposeAsync(); + await flow.WaitForComplete(); + } + } + + static Task NewCancelableProtocolAsync() + => PgTestPool.NewIsolatedAsync(options => + options.CancelSender = PgTestPool.CreateCancelSender(PgTestPool.NewOptions())); + + static async Task ReadBackendPid( + IConsumerDrivenFlowContract contract, PgClientProtocol protocol) + { + var results = Queue(contract, protocol, Command.Create("select pg_backend_pid()")); + var pid = 0; + while (await results.MoveNextAsync()) + { + var rows = results.Current.GetAsyncEnumerator(); + while (await rows.MoveNextAsync()) + pid = rows.Current.GetReader().Read(); + await rows.DisposeAsync(); + } + await results.DisposeAsync(); + return pid; + } + + internal static Command TornStreamedBind() + { + var serializerOptions = new PgSerializerOptions(PgTypeCatalog.Default); + var value = new SlonParameter(new ThrowingReadStream(256 * 1024, 64 * 1024)); + var parameters = new SlonParameters { value }; + parameters.GetOrResolveTypeInfo( + 0, serializerOptions, preparedTypeId: null, allowUnspecified: false); + var parameterSource = new ParameterSource(parameters, SerializerParameterWriter.Instance); + return Command.Create( + "select octet_length($1::bytea)", new ParameterTypeList(parameterSource)) with + { + Parameters = parameterSource + }; + } + + sealed class ThrowingReadStream(int length, int throwAfter) : Stream + { + int _position; + + public override int Read(Span buffer) + { + if (_position >= throwAfter) + throw new IOException("Synthetic parameter read failure."); + var count = Math.Min(buffer.Length, throwAfter - _position); + buffer[..count].Clear(); + _position += count; + return count; + } + + public override ValueTask ReadAsync( + Memory buffer, CancellationToken cancellationToken = default) + => new(Read(buffer.Span)); + public override bool CanRead => true; + public override bool CanSeek => true; + public override bool CanWrite => false; + public override long Length => length; + public override long Position { get => _position; set => throw new NotSupportedException(); } + public override void Flush() { } + public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException(); + public override void SetLength(long value) => throw new NotSupportedException(); + public override void Write(byte[] buffer, int offset, int count) + => throw new NotSupportedException(); + public override int Read(byte[] buffer, int offset, int count) + => Read(buffer.AsSpan(offset, count)); + } + + readonly struct Results( + IConsumerDrivenFlowContract contract, + IAsyncEnumerator inner) : IAsyncDisposable + { + internal CommandResult Current => inner.Current; + internal ValueTask MoveNextAsync(CancellationToken cancellationToken = default) + => contract.MoveNextAsync(inner, cancellationToken); + public ValueTask DisposeAsync() => inner.DisposeAsync(); + } + +} + diff --git a/Slon.Tests/Pg/FlowAuthoring/README.md b/Slon.Tests/Pg/FlowAuthoring/README.md new file mode 100644 index 0000000..05193a7 --- /dev/null +++ b/Slon.Tests/Pg/FlowAuthoring/README.md @@ -0,0 +1,59 @@ +# PostgreSQL flow authoring contracts + +These tests are the executable lifecycle specification for `PgClientFlow` +implementations. They exercise behavior shared by flow types; SQL semantics and +result-shape tests remain with the implementation. + +## Profiles + +- `IAutonomousFlowContract` applies when the flow owns its complete pipeline + task and needs no external consumer to reach RFQ. +- `IConsumerDrivenFlowContract` applies when an external sync or async consumer + owns result advancement and can abandon that ownership. +- `IReusableConsumerDrivenFlowContract` adds the requirement that a completed + instance can be reset for an independent tenure. + +A flow exposing multiple kinds of entry point should register an adapter for every +profile it supports. Profiles are separate interfaces so an implementation +cannot silently opt out of individual required scenarios with capability flags. + +## Registering a flow + +Implement the relevant adapter in this directory and add its singleton to the +profile's `Implementations` data source. For a consumer-driven flow the complete +adapter is: + +```csharp +sealed class MyFlowContract : IConsumerDrivenFlowContract +{ + public string Name => nameof(MyFlow); + public PgClientFlow Create(bool async, params Command[] commands) => + new MyFlow(async, commands); + public IEnumerator GetEnumerator(PgClientFlow flow) => + ((MyFlow)flow).GetEnumerator(); + public IAsyncEnumerator GetAsyncEnumerator(PgClientFlow flow) => + ((MyFlow)flow).GetAsyncEnumerator(); +} +``` + +The interface conversions may box a value-type enumerator. This suite verifies +correctness rather than allocation behavior; implementation-specific benchmarks +continue to call concrete enumerators directly. + +## Contract rules + +The suite treats submission, activation, consumer progress, and protocol +lifecycle as independent axes. In particular: + +- the pipeline task retains wire/decoder tenure through RFQ; +- consumer abandonment transfers read ownership and leaves the wire reusable; +- forceful termination completes a flow even if no consumer attaches; +- sync execution blocks only its caller and does not require a handoff for an + autonomous flow; +- an async consumer never runs on the pipeline executor strand; +- release occurs before completion makes a flow reusable. +- a reset tenure does not retain terminal or consumer state from its predecessor. + +Contract scenarios must use explicit scheduler gates, in-memory transports, or +`FakeTimeProvider` when ordering matters. Do not use wall-clock delays or +timeouts as correctness oracles. diff --git a/Slon.Tests/Pg/HeartbeatTests.cs b/Slon.Tests/Pg/HeartbeatTests.cs index 562e425..e614e01 100644 --- a/Slon.Tests/Pg/HeartbeatTests.cs +++ b/Slon.Tests/Pg/HeartbeatTests.cs @@ -10,6 +10,22 @@ namespace Slon.Tests.Pg; [TestClass] public class HeartbeatTests { + [ConnectionCreatingTestMethod] + [DataRow(PgClientProtocolHeartbeatMode.Automatic, 1)] + [DataRow(PgClientProtocolHeartbeatMode.External, 0)] + public async Task ProtocolHeartbeatModeControlsTimerOwnership( + PgClientProtocolHeartbeatMode mode, int expectedTimerCount) + { + var time = new CountingTimeProvider(); + await using var protocol = await PgTestPool.NewIsolatedAsync(options => + { + options.TimeProvider = time; + options.HeartbeatMode = mode; + }); + + Assert.AreEqual(expectedTimerCount, time.TimerCount); + } + [TestMethod] public async Task BackloggedFlow_ActivationTimeoutAdvancesBeforeDispatch() { @@ -94,4 +110,16 @@ public void Log(LogLevel logLevel, EventId eventId, TState state, } readonly record struct Entry(LogLevel Level, Exception? Exception, string Message); + + sealed class CountingTimeProvider : FakeTimeProvider + { + internal int TimerCount { get; private set; } + + public override ITimer CreateTimer(TimerCallback callback, object? state, + TimeSpan dueTime, TimeSpan period) + { + TimerCount++; + return base.CreateTimer(callback, state, dueTime, period); + } + } } diff --git a/Slon.Tests/Pg/PgEncoderPreparedExecutionTests.cs b/Slon.Tests/Pg/PgEncoderPreparedExecutionTests.cs new file mode 100644 index 0000000..5ec1b5c --- /dev/null +++ b/Slon.Tests/Pg/PgEncoderPreparedExecutionTests.cs @@ -0,0 +1,112 @@ +using System.Buffers.Binary; +using System.Text; +using Slon.Pg.Protocol; +using Slon.Text; + +namespace Slon.Tests.Pg; + +// The fused prepared-execution write must produce exactly the bytes the message-per-message writers +// produce for every Describe, Execute, and Sync combination, then return the writer at a complete +// message boundary. +[TestClass] +public class PgEncoderPreparedExecutionTests +{ + static readonly Encoding Encoding = Encoding.UTF8; + + static (ProtocolDataWriter Writer, BufferOutputWriter Sink) NewWriter() + { + var sink = new BufferOutputWriter(); + var writer = new ProtocolDataWriter(sink, Encoding, static _ => { }, default, null!); + return (writer, sink); + } + + static byte[] Expected(string commandName, bool describe, bool execute, int syncCount) + { + var bytes = new List(); + var name = Encoding.GetBytes(commandName); + + // Bind: unnamed portal, statement name, no parameter format codes, no parameters, one result + // format code, binary. + Message(bytes, (byte)'B', [0, .. name, 0, 0, 0, 0, 0, 0, 1, 0, 1]); + if (describe) + Message(bytes, (byte)'D', [(byte)'P', 0]); + if (execute) + Message(bytes, (byte)'E', [0, 0, 0, 0, 0]); + for (var i = 0; i < syncCount; i++) + Message(bytes, (byte)'S', []); + return bytes.ToArray(); + + static void Message(List bytes, byte type, ReadOnlySpan body) + { + bytes.Add(type); + Span length = stackalloc byte[sizeof(uint)]; + BinaryPrimitives.WriteUInt32BigEndian(length, (uint)(sizeof(uint) + body.Length)); + bytes.AddRange(length); + bytes.AddRange(body); + } + } + + [TestMethod] + [DataRow(false, false, 0)] + [DataRow(false, false, 1)] + [DataRow(false, false, 2)] + [DataRow(false, true, 0)] + [DataRow(false, true, 1)] + [DataRow(false, true, 2)] + [DataRow(true, false, 0)] + [DataRow(true, false, 1)] + [DataRow(true, false, 2)] + [DataRow(true, true, 0)] + [DataRow(true, true, 1)] + [DataRow(true, true, 2)] + public void WritesExactWireBytes(bool describe, bool execute, int syncCount) + { + var (writer, sink) = NewWriter(); + + PgEncoder.WritePreparedExecutionCore(writer, Encoding, new EncodedCString("prepared_probe"), + describe, execute, syncCount); + writer.Flush(); + + CollectionAssert.AreEqual(Expected("prepared_probe", describe, execute, syncCount), sink.ToArray()); + } + + [TestMethod] + public void UnnamedStatement_WritesEmptyName() + { + var (writer, sink) = NewWriter(); + + PgEncoder.WritePreparedExecutionCore(writer, Encoding, default, describe: false, execute: true, syncCount: 1); + writer.Flush(); + + CollectionAssert.AreEqual(Expected("", describe: false, execute: true, syncCount: 1), sink.ToArray()); + } + + [TestMethod] + public void CompleteSequence_LeavesIncrementalTrackerIdle() + { + var (writer, sink) = NewWriter(); + + PgEncoder.WritePreparedExecutionCore(writer, Encoding, new EncodedCString("prepared_probe"), + describe: true, execute: true, syncCount: 2); + Assert.AreEqual(0, writer.CurrentMessagePaddingLength); + // The next incremental message starts from a clean boundary. + writer.StartMessage(totalLength: 5); + writer.WriteRaw(new byte[5]); + writer.Flush(); + + Assert.AreEqual(Expected("prepared_probe", describe: true, execute: true, syncCount: 2).Length + 5, + sink.ToArray().Length); + } + + [TestMethod] + [DataRow(-1)] + [DataRow(3)] + public void SyncCount_OutsideZeroToTwo_Throws(int syncCount) + { + var (writer, _) = NewWriter(); + + Assert.ThrowsExactly(() => + PgEncoder.WritePreparedExecutionCore(writer, Encoding, new EncodedCString("prepared_probe"), + describe: false, execute: true, syncCount)); + } +} diff --git a/Slon.Tests/Pg/PostgreSqlSslTests.cs b/Slon.Tests/Pg/PostgreSqlSslTests.cs index 42e63e0..039245e 100644 --- a/Slon.Tests/Pg/PostgreSqlSslTests.cs +++ b/Slon.Tests/Pg/PostgreSqlSslTests.cs @@ -175,7 +175,7 @@ public async Task VerifyFull_RejectsUntrustedCertificate() await CreateFactory(listener, PostgreSqlSslNegotiation.Direct, PostgreSqlSslMode.VerifyFull).CreateAsync()); try { await server; } - catch (IOException) { } + catch (Exception ex) when (ex is IOException or AuthenticationException) { } } [TestMethod] @@ -190,7 +190,7 @@ public async Task VerifyCA_RejectsUntrustedCertificate() await CreateFactory(listener, PostgreSqlSslNegotiation.Direct, PostgreSqlSslMode.VerifyCA).CreateAsync()); try { await server; } - catch (IOException) { } + catch (Exception ex) when (ex is IOException or AuthenticationException) { } } [TestMethod] diff --git a/Slon.Tests/Pg/ProtocolDataWriterMessageBudgetTests.cs b/Slon.Tests/Pg/ProtocolDataWriterMessageBudgetTests.cs index daa2c79..4274617 100644 --- a/Slon.Tests/Pg/ProtocolDataWriterMessageBudgetTests.cs +++ b/Slon.Tests/Pg/ProtocolDataWriterMessageBudgetTests.cs @@ -58,6 +58,17 @@ public void UnderWrite_FaultsAtNextStartMessage() Assert.ThrowsExactly(() => writer.StartMessage(totalLength: 5)); } + [TestMethod] + public void UnderWrite_FaultsBeforeCompleteMessageReservation() + { + var (writer, sink) = NewWriter(); + writer.StartMessage(totalLength: 5); + writer.WriteRaw(new byte[3]); + + Assert.ThrowsExactly(() => writer.GetCompleteMessagesSpan(5)); + Assert.AreEqual(0, sink.ToArray().Length); + } + [TestMethod] public void ExactWrite_Flushes_AllBytesReachWire() { diff --git a/Slon.Tests/Pg/PublicCommandFlowSurfaceTests.cs b/Slon.Tests/Pg/PublicCommandFlowSurfaceTests.cs new file mode 100644 index 0000000..88e4fd9 --- /dev/null +++ b/Slon.Tests/Pg/PublicCommandFlowSurfaceTests.cs @@ -0,0 +1,51 @@ +using System.Collections; +using System.Diagnostics.CodeAnalysis; +using System.Reflection; +using Slon.Pg; +using Slon.Pg.Protocol.Flows; + +namespace Slon.Tests.Pg; + +[TestClass] +public class PublicCommandFlowSurfaceTests +{ + [TestMethod] + public void ReplacementIsTheSealedPublicCommandFlowSurface() + { + var flow = typeof(Slon.Pg.Protocol.Flows.CommandFlow); + Assert.IsTrue(flow.IsPublic); + Assert.IsTrue(flow.IsSealed); + Assert.IsTrue(typeof(Slon.Pg.Protocol.Flows.CommandFlowObserver).IsPublic); + Assert.IsTrue(typeof(Slon.Pg.Protocol.Flows.CommandFlowOptions).IsPublic); + Assert.IsTrue(typeof(IEnumerator) + .IsAssignableFrom(typeof(Slon.Pg.Protocol.Flows.CommandFlow.Enumerator))); + + Assert.IsNotNull(flow.GetConstructor([ + typeof(bool), typeof(ReadOnlySpan) + ])); + Assert.IsNotNull(flow.GetConstructor([ + typeof(bool), typeof(Slon.Pg.Protocol.Flows.CommandFlowOptions).MakeByRefType() + ])); + } + + [TestMethod] + public void CollectionIsOnTheExperimentalCommandResultSurface() + { + var result = typeof(CommandResult); + var diagnosticId = result.GetCustomAttribute()!.DiagnosticId; + var collect = result.GetMethods().Single(method => + method.Name == nameof(CommandResult.CollectAsync)); + var rowView = result.GetNestedType( + nameof(CommandResult.RowView), BindingFlags.Public)!; + + Assert.IsTrue(collect.IsPublic); + Assert.IsTrue(collect.IsGenericMethodDefinition); + Assert.AreEqual(diagnosticId, + collect.GetCustomAttribute()!.DiagnosticId); + Assert.IsTrue(rowView.IsNestedPublic); + Assert.IsNotNull(rowView.GetMethod( + nameof(CommandResult.RowView.BorrowFieldMemory), BindingFlags.Public | BindingFlags.Instance)); + Assert.AreEqual(diagnosticId, + rowView.GetCustomAttribute()!.DiagnosticId); + } +} diff --git a/Slon.Tests/Pg/RacingDisposeInMemoryTests.cs b/Slon.Tests/Pg/RacingDisposeInMemoryTests.cs index 9b3f911..5e1e3b6 100644 --- a/Slon.Tests/Pg/RacingDisposeInMemoryTests.cs +++ b/Slon.Tests/Pg/RacingDisposeInMemoryTests.cs @@ -143,7 +143,7 @@ public async Task GracefulCloseMidRowEofSurfacesClosedException() for (var i = 0; i < dataRowIndex; i++) transport.ReleaseSegment(messages[i]); transport.ReleaseSegment(messages[dataRowIndex] - .AsSpan(0, BackendMessageBatch.Segmenter.DefaultDataRowStreamingThreshold).ToArray()); + .AsSpan(0, BackendMessageCursor.DefaultDataRowStreamingThreshold).ToArray()); Assert.IsTrue(await resultPending); var rows = flowEnumerator.Current.GetAsyncEnumerator(CommandResult.RowBuffering.Streaming); @@ -399,6 +399,7 @@ public Task RunLoop() => Task.Run(async () => // a forceful abort faults that read -> body's closed catch. Pre-fix: rethrow escaped DisposeAsync. // Now: DisposeAsyncCore swallows it. [TestMethod] + [Ignore("Exercises legacy body-driven throw and caller-gate ordering.")] public async Task Ordering1_BodyDrivenThrow_DisposeConverges() { await using var s = await BuildToFirstResultParked(); @@ -421,6 +422,7 @@ public async Task Ordering1_BodyDrivenThrow_DisposeConverges() // This isolates terminal publication and the continuation handoff without ThreadPool admission, // PostgreSQL timing, or an advisory lock. [TestMethod] + [Ignore("Exercises takeover of the legacy body coroutine during synchronous disposal.")] public async Task SyncDispose_InFlightReadFault_Converges() { var iterations = Math.Clamp(StressEnv.Iterations(fallback: 1, cap: int.MaxValue), 1, 500); @@ -449,6 +451,7 @@ public async Task SyncDispose_InFlightReadFault_Converges() // bounded exactly as production bounds it - advance past CompletionTimeout (30s) so the graceful->abort // escalation faults the parked drain read and DisposeAsync converges, swallowing the close. [TestMethod] + [Ignore("Exercises the legacy caller gate winning before the body throw.")] public async Task Ordering2_GateFirstThrow_DisposeConverges() { await using var s = await BuildToFirstResultParked(); @@ -473,6 +476,7 @@ public async Task Ordering2_GateFirstThrow_DisposeConverges() // consume that progress without receiving a continuation and transfer to AwaitDrainOnDispose. // Escalation must still redrive/fault the held read and complete both teardown participants. [TestMethod] + [Ignore("Exercises the legacy caller gate's progress-before-takeover ordering.")] public async Task SyncDispose_GateProgressBeforeTakeover_Converges() { var iterations = Math.Clamp(StressEnv.Iterations(fallback: 1, cap: int.MaxValue), 1, 500); @@ -500,6 +504,7 @@ public async Task SyncDispose_GateProgressBeforeTakeover_Converges() // progress-only wake; only afterwards does releasing the response let the body publish its sync // handoff continuation. Null progress cannot be mistaken for body termination. [TestMethod] + [Ignore("Exercises late publication of a legacy body handoff continuation.")] public async Task SyncDispose_ProgressWakeBeforeLateHandoff_DrivesBodyToTermination() { var clock = new FakeTimeProvider(); @@ -539,6 +544,7 @@ public async Task SyncDispose_ProgressWakeBeforeLateHandoff_DrivesBodyToTerminat // A sync-at-bind body parked between command results still owns a continuation after close becomes // consumer-terminal. Dispose must drive that continuation before returning or surfacing the close. [TestMethod] + [Ignore("Exercises the legacy body's inter-result caller gate.")] public async Task SyncFlow_CloseAtInterResultPark_DisposeRetainsDriveObligation() { var clock = new FakeTimeProvider(); @@ -583,6 +589,7 @@ public async Task SyncFlow_CloseAtInterResultPark_DisposeRetainsDriveObligation( // published CancelException, so the consumer's next MoveNextAsync self-delivers the close and the // loop converges. Asserts convergence (a hang here would regress that protection). [TestMethod] + [Ignore("Exercises self-delivery after a no-op fault on the legacy caller gate.")] public async Task Ordering3_GateFaultNoOp_SelfDeliverConverges() { await using var s = await BuildToFirstResultParked(); @@ -607,6 +614,7 @@ public async Task Ordering3_GateFaultNoOp_SelfDeliverConverges() // abort lands CompleteEnumerationWithException on the LIVE (Reset) generation, because every body read in this flow // is preceded by a consumer Reset. Documents why the read-fault alone cannot produce the hang. [TestMethod] + [Ignore("Exercises the legacy body's read-fault-to-caller-gate transition.")] public async Task Ordering3_ReadFaultPath_NeverNoOps_Converges() { await using var s = await BuildToFirstResultParked(); @@ -632,6 +640,7 @@ public async Task Ordering3_ReadFaultPath_NeverNoOps_Converges() // inline ping-pong), so this is a convergence regression test, not an isolation of any one waker. The // never-started lost-completion (the actual reported hang) is gated by the stress repro + dump. [TestMethod] + [Ignore("Exercises graceful close while the legacy body is parked at its inter-result gate.")] public async Task MultiCommand_GracefulCloseAtInterResultGate_Converges() { await using var s = await BuildMultiToFirstResultParked(); @@ -664,6 +673,7 @@ public async Task MultiCommand_GracefulCloseAtInterResultGate_Converges() // the stale result 1. Ordering3_GateFaultNoOp asserts the read loop converges; a stale re-yield does // not hang, so this pins the specific next-call outcome the loop cannot catch. [TestMethod] + [Ignore("Exercises stale continuation suppression in the legacy caller gate.")] public async Task GateFaultBeforeNextMoveNext_SelfDeliversClose_NeverReYieldsStale() { await using var s = await BuildToFirstResultParked(); diff --git a/Slon.Tests/Pg/StoppingTokenInMemoryTests.cs b/Slon.Tests/Pg/StoppingTokenInMemoryTests.cs index 2da723d..56ee86c 100644 --- a/Slon.Tests/Pg/StoppingTokenInMemoryTests.cs +++ b/Slon.Tests/Pg/StoppingTokenInMemoryTests.cs @@ -437,8 +437,8 @@ public Task ArmReadPark() return tcs.Task; } - // The protocol's transport read (PipeSegmentEnumerator) funnels through ReadAsync, including - // via the base ReadAtLeastAsync default; signaling here covers both. AdvanceTo is delegated + // The protocol transport funnels through ReadAsync, including via the base + // ReadAtLeastAsync default; signaling here covers both. AdvanceTo is delegated // verbatim (no teeing), so the base default's position tracking stays correct. public override ValueTask ReadAsync(CancellationToken cancellationToken = default) { diff --git a/Slon.Tests/Pg/SyncFlowHandoffTests.cs b/Slon.Tests/Pg/SyncFlowHandoffTests.cs index 89f9993..571d40d 100644 --- a/Slon.Tests/Pg/SyncFlowHandoffTests.cs +++ b/Slon.Tests/Pg/SyncFlowHandoffTests.cs @@ -34,40 +34,10 @@ sealed class WakeHolder internal FlowCallerInteractionCore Core; } - static ref FlowCallerInteractionCore GetWakeCore(WakeHolder holder) => ref holder.Core; - - [ConnectionCreatingTestMethod] - public async Task ConcurrentSyncAndAsync_NoSharedPromiseCollision() + readonly struct WakeRef(WakeHolder holder) + : IFieldRef> { - await using var protocol = await PgTestPool.NewIsolatedAsync(); - Exception? failure = null; - void Capture(Exception ex) => Interlocked.CompareExchange(ref failure, ex, null); - - var asyncLoop = Task.Run(async () => - { - try - { - for (var i = 0; i < StressIterations && Volatile.Read(ref failure) is null; i++) - await PgTestPool.RunAsync(protocol, "select 1"); - } - catch (Exception ex) { Capture(ex); } - }); - - var syncThread = new Thread(() => - { - try - { - for (var i = 0; i < StressIterations && Volatile.Read(ref failure) is null; i++) - PgTestPool.RunSync(protocol, "select 1").GetAwaiter().GetResult(); - } - catch (Exception ex) { Capture(ex); } - }) { IsBackground = true, Name = "sync-flow-handoff" }; - - syncThread.Start(); - await asyncLoop; - syncThread.Join(); - if (failure is not null) - Assert.Fail($"concurrent sync/async raised {failure}"); + public ref FlowCallerInteractionCore GetField() => ref holder.Core; } [ConnectionCreatingTestMethod] @@ -107,12 +77,7 @@ public async Task DedicatedWakeResumesOutsideTheThreadPool() static async Task Suspend(WakeHolder holder, TaskCompletionSource resumedOnThreadPool) { - FieldRef> fieldRef; - unsafe - { - fieldRef = FieldRef>.Create(&GetWakeCore, holder); - } - await holder.Core.YieldToCaller(fieldRef); + await holder.Core.YieldToCaller(new WakeRef(holder)); resumedOnThreadPool.SetResult(Thread.CurrentThread.IsThreadPoolThread); } } diff --git a/Slon.Tests/Pipelines/SegmentChainBuilderTests.cs b/Slon.Tests/Pipelines/SegmentChainBuilderTests.cs index c71d4ac..fb9c658 100644 --- a/Slon.Tests/Pipelines/SegmentChainBuilderTests.cs +++ b/Slon.Tests/Pipelines/SegmentChainBuilderTests.cs @@ -6,6 +6,51 @@ namespace Slon.Tests.Pipelines; [TestClass] public class SegmentChainBuilderTests { + const int ReaderBufferSize = 65536; + const int MinimumReadSize = 512; + + [TestMethod] + public void ReaderReserve_WithConsumedPrefix_ConsolidatesRemainingBytes() + { + using var builder = new SegmentChainBuilder( + MemoryPool.Shared, ReaderBufferSize, MinimumReadSize, + retainBufferOnEmpty: true); + var memory = builder.Reserve(ReaderBufferSize, enforceHint: true); + memory.Span[^1] = 42; + builder.Grow(ReaderBufferSize); + var initial = builder.GetReadOnlySequence(); + var initialHead = builder.HeadInfo.Head; + + builder.AdvanceTo(initial.GetPosition(ReaderBufferSize - 1024)); + builder.Reserve(MinimumReadSize, enforceHint: true); + + var consolidated = builder.GetReadOnlySequence(); + Assert.IsTrue(consolidated.IsSingleSegment); + Assert.AreNotSame(initialHead, builder.HeadInfo.Head); + Assert.AreEqual(0, builder.HeadInfo.Index); + Assert.AreEqual(1024, consolidated.Length); + Assert.AreEqual(42, consolidated.FirstSpan[^1]); + } + + [TestMethod] + public void ReaderReserve_AtUnadvancedHead_CreatesSecondSegment() + { + using var builder = new SegmentChainBuilder( + MemoryPool.Shared, ReaderBufferSize, MinimumReadSize, + retainBufferOnEmpty: true); + builder.Reserve(ReaderBufferSize, enforceHint: true).Span[^1] = 42; + builder.Grow(ReaderBufferSize); + var initialHead = builder.HeadInfo.Head; + + builder.Reserve(MinimumReadSize, enforceHint: true); + + var chained = builder.GetReadOnlySequence(); + Assert.IsFalse(chained.IsSingleSegment); + Assert.AreSame(initialHead, builder.HeadInfo.Head); + Assert.AreEqual(ReaderBufferSize, chained.Length); + Assert.AreEqual(42, chained.FirstSpan[^1]); + } + [TestMethod] public void AdvanceToEmpty_RetainsSegmentAndOwnedMemory_WhenEnabled() { diff --git a/Slon.Tests/Pipelines/StreamPipeReaderTests.cs b/Slon.Tests/Pipelines/StreamPipeReaderTests.cs index cae06d5..8087073 100644 --- a/Slon.Tests/Pipelines/StreamPipeReaderTests.cs +++ b/Slon.Tests/Pipelines/StreamPipeReaderTests.cs @@ -19,6 +19,30 @@ static async Task CreateBufferedReader(byte[] bytes) return reader; } + [TestMethod] + public async Task AdvanceTo_CanUnexamineBufferedData() + { + var reader = new DefaultStreamPipeReader( + new MemoryStream("Hello World"u8.ToArray(), writable: false), + new StreamPipeReaderOptions(bufferSize: 1024, useZeroByteReads: false), + supportCancelPending: false); + var first = await reader.ReadAsync(); + reader.AdvanceTo(first.Buffer.GetPosition(6), first.Buffer.End); + + // Examining the complete first grant forces the stream EOF probe. + var eof = await reader.ReadAsync(); + CollectionAssert.AreEqual("World"u8.ToArray(), eof.Buffer.ToArray()); + Assert.IsTrue(eof.IsCompleted); + + // Moving examined back to consumed must republish the existing suffix without another + // stream read, matching the .NET 10 PipeReader un-examine contract. + reader.AdvanceTo(eof.Buffer.Start, eof.Buffer.Start); + Assert.IsTrue(reader.TryRead(out var unexamined)); + CollectionAssert.AreEqual("World"u8.ToArray(), unexamined.Buffer.ToArray()); + reader.AdvanceTo(unexamined.Buffer.End); + await reader.CompleteAsync(); + } + [TestMethod] public async Task CopyToAsync_Stream_ConsumesSuccessfullyCopiedBufferedData() { diff --git a/Slon.Tests/Slon.Tests.csproj b/Slon.Tests/Slon.Tests.csproj index a1dbcc8..5273122 100644 --- a/Slon.Tests/Slon.Tests.csproj +++ b/Slon.Tests/Slon.Tests.csproj @@ -1,6 +1,7 @@ - net10.0 + net10.0;net11.0 + runtime-async=on preview enable enable diff --git a/Slon.slnx b/Slon.slnx index f4c0908..56b8f3d 100644 --- a/Slon.slnx +++ b/Slon.slnx @@ -3,5 +3,7 @@ + + diff --git a/Slon/Ado/AdoBatchCore.Preparation.cs b/Slon/Ado/AdoBatchCore.Preparation.cs index 76ebe34..cdd7ded 100644 --- a/Slon/Ado/AdoBatchCore.Preparation.cs +++ b/Slon/Ado/AdoBatchCore.Preparation.cs @@ -6,7 +6,9 @@ namespace Slon; -partial struct AdoBatchCore where TCommand : IAdoCommand +partial struct AdoBatchCore + where TCommand : IAdoCommand + where TFieldRef : struct, IAdoBatchCoreRef { public void Prepare(DbParameterCollection? parameters) { @@ -25,7 +27,7 @@ public void Prepare(DbParameterCollection? parameters) void PrepareCore(DbParameterCollection? parameters) { var operation = Preparation.Begin(_fieldRef); - CommandFlow.Enumerator enumerator = default; + AdoCommandExecutionFlow.Enumerator enumerator = default; try { var flow = Enqueue(parameters, CommandBehavior.SchemaOnly, GetDependencies(), @@ -64,10 +66,10 @@ public ValueTask PrepareAsync(DbParameterCollection? parameters, CancellationToken cancellationToken = default) => PrepareAsyncProjected(_fieldRef, parameters, cancellationToken); - static async ValueTask PrepareAsyncProjected(FieldRef> fieldRef, + static async ValueTask PrepareAsyncProjected(TFieldRef fieldRef, DbParameterCollection? parameters, CancellationToken cancellationToken) { - using var activity = fieldRef.Invoke().StartActivity(); + using var activity = fieldRef.GetField().StartActivity(); try { await PrepareAsyncCore(fieldRef, parameters, cancellationToken).ConfigureAwait(false); @@ -80,17 +82,17 @@ static async ValueTask PrepareAsyncProjected(FieldRef> fi } // Async instance methods on structs copy this, so the state machine resolves the live core - // through its stable field reference instead. - static async ValueTask PrepareAsyncCore(FieldRef> fieldRef, + // through its stable fieldRef reference instead. + static async ValueTask PrepareAsyncCore(TFieldRef fieldRef, DbParameterCollection? parameters, CancellationToken cancellationToken) { var operation = Preparation.Begin(fieldRef); - CommandFlow.Enumerator enumerator = default; + AdoCommandExecutionFlow.Enumerator enumerator = default; try { - var dependencies = await fieldRef.Invoke().GetDependenciesAsync(cancellationToken) + var dependencies = await fieldRef.GetField().GetDependenciesAsync(cancellationToken) .ConfigureAwait(false); - var flow = await fieldRef.Invoke().EnqueueAsync(parameters, CommandBehavior.SchemaOnly, + var flow = await fieldRef.GetField().EnqueueAsync(parameters, CommandBehavior.SchemaOnly, dependencies, cancellationToken, preparing: true).ConfigureAwait(false); enumerator = flow.GetAsyncEnumerator(cancellationToken); for (var i = 0; i < operation.CommandCount; i++) @@ -124,12 +126,12 @@ static async ValueTask PrepareAsyncCore(FieldRef> fieldRe struct Preparation { - readonly FieldRef> _fieldRef; + readonly TFieldRef _fieldRef; readonly SlonDataSource? _dataSource; readonly SlonConnection? _connection; List? _exceptions; - Preparation(FieldRef> fieldRef, + Preparation(TFieldRef fieldRef, SlonDataSource? dataSource, SlonConnection? connection) { _fieldRef = fieldRef; @@ -137,16 +139,16 @@ struct Preparation _connection = connection; } - internal static Preparation Begin(FieldRef> fieldRef) + internal static Preparation Begin(TFieldRef fieldRef) { - ref var core = ref fieldRef.Invoke(); + ref var core = ref fieldRef.GetField(); core.ThrowIfDisposedOrReadOnly(); core.TryGetDataSource(out var dataSource, out var connection); core._explicitlyPrepared = true; return new(fieldRef, dataSource, connection); } - internal int CommandCount => _fieldRef.Invoke()._commands.Count; + internal int CommandCount => _fieldRef.GetField()._commands.Count; internal void Observe(CommandResult result) { @@ -169,34 +171,34 @@ internal void ThrowIfFailed() internal void Commit() { - foreach (ref var command in _fieldRef.Invoke()._commands.AsSpan()) + foreach (ref var command in _fieldRef.GetField()._commands.AsSpan()) command.MakeReadOnly(); } internal void Rollback() - => _fieldRef.Invoke()._explicitlyPrepared = false; + => _fieldRef.GetField()._explicitlyPrepared = false; internal void ReleaseFailedPreparation() { - if (_fieldRef.Invoke()._explicitlyPrepared) + if (_fieldRef.GetField()._explicitlyPrepared) return; if (_connection is not null) - _connection.UnprepareOwned(async: false, _fieldRef.Instance).GetAwaiter().GetResult(); + _connection.UnprepareOwned(async: false, _fieldRef.Owner).GetAwaiter().GetResult(); else if (_dataSource is not null) _ = _dataSource.ReleaseOwnedPreparedCommand( - _fieldRef.Instance, awaitable: false); + _fieldRef.Owner, awaitable: false); } internal ValueTask ReleaseFailedPreparationAsync() { - if (_fieldRef.Invoke()._explicitlyPrepared) + if (_fieldRef.GetField()._explicitlyPrepared) return default; if (_connection is not null) - return _connection.UnprepareOwned(async: true, _fieldRef.Instance); + return _connection.UnprepareOwned(async: true, _fieldRef.Owner); return _dataSource?.ReleaseOwnedPreparedCommand( - _fieldRef.Instance, awaitable: true) ?? default; + _fieldRef.Owner, awaitable: true) ?? default; } } } diff --git a/Slon/Ado/AdoBatchCore.cs b/Slon/Ado/AdoBatchCore.cs index 7fac606..a7f41eb 100644 --- a/Slon/Ado/AdoBatchCore.cs +++ b/Slon/Ado/AdoBatchCore.cs @@ -9,10 +9,20 @@ namespace Slon; +interface IAdoBatchCoreRef + : IFieldRef> + where TCommand : IAdoCommand + where TFieldRef : struct, IAdoBatchCoreRef +{ + IAdoCommandExecutionOwner Owner { get; } +} + // Shared between DbBatch and DbCommand -partial struct AdoBatchCore where TCommand : IAdoCommand +partial struct AdoBatchCore + where TCommand : IAdoCommand + where TFieldRef : struct, IAdoBatchCoreRef { - readonly FieldRef> _fieldRef; + readonly TFieldRef _fieldRef; object _dataSourceOrConnection; bool _disposed; bool _explicitlyPrepared; @@ -20,22 +30,22 @@ partial struct AdoBatchCore where TCommand : IAdoCommand TimeSpan _timeout; TimeSpan? _pendingTimeout; bool _enableErrorBarriers; - CommandFlow? _activeFlow; + AdoCommandExecutionFlow? _activeFlow; AdoCommandList _commands; - public AdoBatchCore(FieldRef> fieldRef) + public AdoBatchCore(TFieldRef fieldRef) { _dataSourceOrConnection = null!; _fieldRef = fieldRef; } - public AdoBatchCore(SlonConnection connection, FieldRef> fieldRef) + public AdoBatchCore(TFieldRef fieldRef, SlonConnection connection) { _dataSourceOrConnection = connection; _fieldRef = fieldRef; } - public AdoBatchCore(SlonDataSource dataSource, FieldRef> fieldRef) + public AdoBatchCore(TFieldRef fieldRef, SlonDataSource dataSource) { _dataSourceOrConnection = dataSource; _fieldRef = fieldRef; @@ -113,7 +123,7 @@ public void ThrowIfDisposedOrReadOnly() public void ThrowIfDisposed() { if (_disposed) - Throw(_fieldRef.Instance); + Throw(_fieldRef.Owner); static void Throw(object instance) => throw new ObjectDisposedException(instance.GetType().Name); } @@ -172,7 +182,7 @@ internal AdoCommandFlowOptions CreateAdoCommandFlowOptions( PgConnection? pgConnection = null, TimeSpan? pendingTimeout = null, bool preparing = false) { var factory = new AdoCommandFlowFactory( - _fieldRef.Instance, _commands.AsSpan(), dependencies); + _fieldRef.Owner, _commands.AsSpan(), dependencies); return factory.Create( parametersSpan, behavior, _explicitlyPrepared, _allowAutoPreparation, _enableErrorBarriers, Timeout, @@ -180,21 +190,29 @@ internal AdoCommandFlowOptions CreateAdoCommandFlowOptions( } SlonDataSource.PgDbDependencies GetDependencies() + => GetDependencies(out _); + + SlonDataSource.PgDbDependencies GetDependencies( + out SlonConnection? connection) { - TryGetDataSource(out var dataSource, out var connection); + TryGetDataSource(out var dataSource, out connection); connection ??= dataSource is null ? ThrowConnectionNotInitialized() : null; return (dataSource ?? connection!.DbDataSource).GetDbDependencies(); } ValueTask GetDependenciesAsync( CancellationToken cancellationToken) + => GetDependenciesAsync(cancellationToken, out _); + + ValueTask GetDependenciesAsync( + CancellationToken cancellationToken, out SlonConnection? connection) { - TryGetDataSource(out var dataSource, out var connection); + TryGetDataSource(out var dataSource, out connection); connection ??= dataSource is null ? ThrowConnectionNotInitialized() : null; return (dataSource ?? connection!.DbDataSource).GetDbDependenciesAsync(cancellationToken); } - CommandFlow Enqueue(DbParameterCollection? parameters, CommandBehavior behavior, + AdoCommandExecutionFlow Enqueue(DbParameterCollection? parameters, CommandBehavior behavior, SlonDataSource.PgDbDependencies dependencies, bool preparing = false) { if (TryGetDataSource(out var dataSource, out var connection)) @@ -202,20 +220,22 @@ CommandFlow Enqueue(DbParameterCollection? parameters, CommandBehavior behavior, ThrowIfHasCloseConnection(behavior); var pendingTimeout = PendingTimeout; return dataSource.EnqueueCommands( - new AdoCommandFlow( - async: false, _fieldRef, parameters, behavior, dependencies, + new AdoCommandExecutionFlow( + async: false, _fieldRef.Owner, + parameters, behavior, dependencies, connection: null, pendingTimeout, preparing, _commands.Count, - _explicitlyPrepared && _fieldRef.Instance is SlonCommand ? null : _fieldRef.Instance), + ownsLifetime: !(_explicitlyPrepared && _fieldRef.Owner is SlonCommand)), pendingTimeout); } connection ??= ThrowConnectionNotInitialized(); - return connection.Enqueue(new AdoCommandFlow( - async: false, _fieldRef, parameters, behavior, dependencies, - connection, PendingTimeout, preparing, _commands.Count, _fieldRef.Instance)); + return connection.Enqueue(new AdoCommandExecutionFlow( + async: false, _fieldRef.Owner, + parameters, behavior, dependencies, connection, PendingTimeout, preparing, + _commands.Count, ownsLifetime: true)); } - ValueTask EnqueueAsync(DbParameterCollection? parameters, + ValueTask EnqueueAsync(DbParameterCollection? parameters, CommandBehavior behavior, SlonDataSource.PgDbDependencies dependencies, CancellationToken cancellationToken, bool preparing = false) { @@ -224,17 +244,19 @@ ValueTask EnqueueAsync(DbParameterCollection? parameters, ThrowIfHasCloseConnection(behavior); var pendingTimeout = PendingTimeout; return dataSource.EnqueueCommandsAsync( - new AdoCommandFlow( - async: true, _fieldRef, parameters, behavior, dependencies, + new AdoCommandExecutionFlow( + async: true, _fieldRef.Owner, + parameters, behavior, dependencies, connection: null, pendingTimeout, preparing, _commands.Count, - _explicitlyPrepared && _fieldRef.Instance is SlonCommand ? null : _fieldRef.Instance), + ownsLifetime: !(_explicitlyPrepared && _fieldRef.Owner is SlonCommand)), pendingTimeout, cancellationToken); } connection ??= ThrowConnectionNotInitialized(); - return connection.EnqueueAsync(new AdoCommandFlow( - async: true, _fieldRef, parameters, behavior, dependencies, - connection, PendingTimeout, preparing, _commands.Count, _fieldRef.Instance), cancellationToken); + return connection.EnqueueAsync(new AdoCommandExecutionFlow( + async: true, _fieldRef.Owner, + parameters, behavior, dependencies, connection, PendingTimeout, preparing, + _commands.Count, ownsLifetime: true), cancellationToken); } [DoesNotReturn] @@ -277,9 +299,9 @@ int ExecuteNonQueryCore(DbParameterCollection? parameters) return checked((int)recordsAffected); } - static async ValueTask ExecuteNonQueryAsyncCore(FieldRef> fieldRef, DbParameterCollection? parameters, CancellationToken cancellationToken) + static async ValueTask ExecuteNonQueryAsyncCore(TFieldRef fieldRef, DbParameterCollection? parameters, CancellationToken cancellationToken) { - ref var thisRef = ref fieldRef.Invoke(); + ref var thisRef = ref fieldRef.GetField(); using var activity = thisRef.StartActivity(); try { @@ -287,7 +309,7 @@ static async ValueTask ExecuteNonQueryAsyncCore(FieldRef ExecuteNonQueryAsync(DbParameterCollection? parameters, Ca return null; } - static async ValueTask ExecuteScalarAsyncCore(FieldRef> fieldRef, DbParameterCollection? parameters, CancellationToken cancellationToken) + static async ValueTask ExecuteScalarAsyncCore(TFieldRef fieldRef, DbParameterCollection? parameters, CancellationToken cancellationToken) { - ref var thisRef = ref fieldRef.Invoke(); + ref var thisRef = ref fieldRef.GetField(); using var activity = thisRef.StartActivity(); - CommandFlow.Enumerator enumerator = default; + AdoCommandExecutionFlow.Enumerator enumerator = default; try { thisRef.ThrowIfDisposed(); @@ -354,7 +376,7 @@ public ValueTask ExecuteNonQueryAsync(DbParameterCollection? parameters, Ca var dependencies = await thisRef.GetDependenciesAsync(cancellationToken) .ConfigureAwait(false); var fieldReader = new PgSerializerFieldReader(dependencies.SerializerOptions); - enumerator = (await fieldRef.Invoke().EnqueueAsync(parameters, CommandBehavior.Default, + enumerator = (await fieldRef.GetField().EnqueueAsync(parameters, CommandBehavior.Default, dependencies, cancellationToken).ConfigureAwait(false)) .GetAsyncEnumerator(cancellationToken); while (await enumerator.MoveNextAsync().ConfigureAwait(false)) @@ -426,22 +448,16 @@ public SlonDataReader ExecuteReader(DbParameterCollection? parameters, CommandBe SlonDataReader ExecuteReaderCore(DbParameterCollection? parameters, CommandBehavior behavior) { - var dependencies = GetDependencies(); + var dependencies = GetDependencies(out var connection); return SlonDataReader.Create(behavior, Enqueue(parameters, behavior, dependencies), - dependencies.SerializerOptions, GetConnectionToClose(behavior)); - } - - SlonConnection? GetConnectionToClose(CommandBehavior behavior) - { - if (!HasCloseConnection(behavior) || TryGetDataSource(out _, out var connection)) - return null; - return connection; + dependencies.SerializerOptions, + HasCloseConnection(behavior) ? connection : null); } public ValueTask ExecuteDbReaderAsync(DbParameterCollection? parameters, CommandBehavior behavior, CancellationToken cancellationToken = default) { if (_disposed) - return ValueTask.FromException(new ObjectDisposedException(_fieldRef.Instance.GetType().Name)); + return ValueTask.FromException(new ObjectDisposedException(_fieldRef.Owner.GetType().Name)); if (cancellationToken.IsCancellationRequested) return ValueTask.FromCanceled(cancellationToken); @@ -452,7 +468,7 @@ public ValueTask ExecuteDbReaderAsync(DbParameterCollection? param public ValueTask ExecuteReaderAsync(DbParameterCollection? parameters, CommandBehavior behavior, CancellationToken cancellationToken = default) { if (_disposed) - return ValueTask.FromException(new ObjectDisposedException(_fieldRef.Instance.GetType().Name)); + return ValueTask.FromException(new ObjectDisposedException(_fieldRef.Owner.GetType().Name)); if (cancellationToken.IsCancellationRequested) return ValueTask.FromCanceled(cancellationToken); @@ -461,21 +477,22 @@ public ValueTask ExecuteReaderAsync(DbParameterCollection? param } static ValueTask ExecuteReaderAsyncCore( - FieldRef> fieldRef, DbParameterCollection? parameters, + TFieldRef fieldRef, DbParameterCollection? parameters, CommandBehavior behavior, CancellationToken cancellationToken) where TReader : DbDataReader { - ref var core = ref fieldRef.Invoke(); + ref var core = ref fieldRef.GetField(); var activity = core.StartActivity(); try { - var connectionToClose = core.GetConnectionToClose(behavior); - var dependenciesTask = core.GetDependenciesAsync(cancellationToken); + var closeConnection = core.HasCloseConnection(behavior); + var dependenciesTask = core.GetDependenciesAsync( + cancellationToken, out var connection); return dependenciesTask.IsCompletedSuccessfully ? BeginReaderCreation(fieldRef, parameters, behavior, cancellationToken, - connectionToClose, dependenciesTask.Result, activity) + connection, closeConnection, dependenciesTask.Result, activity) : AwaitDependenciesAndCreateReaderAsync(fieldRef, parameters, behavior, - cancellationToken, connectionToClose, dependenciesTask, activity); + cancellationToken, connection, closeConnection, dependenciesTask, activity); } catch (Exception ex) { @@ -484,28 +501,32 @@ static ValueTask ExecuteReaderAsyncCore( } static ValueTask BeginReaderCreation( - FieldRef> fieldRef, DbParameterCollection? parameters, + TFieldRef fieldRef, DbParameterCollection? parameters, CommandBehavior behavior, CancellationToken cancellationToken, - SlonConnection? connectionToClose, SlonDataSource.PgDbDependencies dependencies, + SlonConnection? connection, bool closeConnection, + SlonDataSource.PgDbDependencies dependencies, Activity? activity) where TReader : DbDataReader { try { return SlonDataReader.CreateAsync(behavior, - fieldRef.Invoke().EnqueueAsync(parameters, behavior, dependencies, cancellationToken), - dependencies.SerializerOptions, cancellationToken, connectionToClose, activity); + fieldRef.GetField().EnqueueAsync(parameters, behavior, dependencies, cancellationToken), + dependencies.SerializerOptions, cancellationToken, + closeConnection ? connection : null, activity); } catch (Exception ex) { + if (closeConnection && connection is not null) + return FailAndCloseReaderCreation(connection, activity, ex); return FailReaderCreation(activity, ex); } } static async ValueTask AwaitDependenciesAndCreateReaderAsync( - FieldRef> fieldRef, DbParameterCollection? parameters, + TFieldRef fieldRef, DbParameterCollection? parameters, CommandBehavior behavior, CancellationToken cancellationToken, - SlonConnection? connectionToClose, + SlonConnection? connection, bool closeConnection, ValueTask dependenciesTask, Activity? activity) where TReader : DbDataReader { @@ -523,7 +544,7 @@ static async ValueTask AwaitDependenciesAndCreateReaderAsync( } return await BeginReaderCreation(fieldRef, parameters, behavior, cancellationToken, - connectionToClose, dependencies, activity).ConfigureAwait(false); + connection, closeConnection, dependencies, activity).ConfigureAwait(false); } static ValueTask FailReaderCreation(Activity? activity, Exception exception) @@ -534,7 +555,35 @@ static ValueTask FailReaderCreation(Activity? activity, Except return ValueTask.FromException(AdoException.Project(exception)); } + static async ValueTask FailAndCloseReaderCreation( + SlonConnection connection, Activity? activity, Exception exception) + where TReader : DbDataReader + { + try + { + await connection.CloseAsync().ConfigureAwait(false); + } + catch (Exception cleanupException) + { + exception = cleanupException; + } + SlonTracing.RecordException(activity, exception); + activity?.Dispose(); + AdoException.Throw(exception); + return default!; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] Activity? StartActivity() + { + if (!SlonTracing.ShouldStart) + return null; + + return StartActivityCore(); + } + + [MethodImpl(MethodImplOptions.NoInlining)] + Activity? StartActivityCore() { TryGetDataSource(out var dataSource, out var connection); return dataSource is null && connection is null @@ -562,12 +611,12 @@ public void Dispose() return; if (TryGetDataSource(out var dataSource, out var connection)) { - _ = dataSource.ReleaseOwnedPreparedCommand(_fieldRef.Instance, awaitable: false); + _ = dataSource.ReleaseOwnedPreparedCommand(_fieldRef.Owner, awaitable: false); return; } if (connection is not null) - connection.UnprepareOwned(async: false, _fieldRef.Instance).GetAwaiter().GetResult(); + connection.UnprepareOwned(async: false, _fieldRef.Owner).GetAwaiter().GetResult(); } public ValueTask DisposeAsync() @@ -579,9 +628,9 @@ public ValueTask DisposeAsync() if (!_explicitlyPrepared) return new(); if (TryGetDataSource(out var dataSource, out var connection)) - return dataSource.ReleaseOwnedPreparedCommand(_fieldRef.Instance, awaitable: true); + return dataSource.ReleaseOwnedPreparedCommand(_fieldRef.Owner, awaitable: true); - return connection?.UnprepareOwned(async: true, _fieldRef.Instance) ?? default; + return connection?.UnprepareOwned(async: true, _fieldRef.Owner) ?? default; } public void Cancel() @@ -607,9 +656,9 @@ static void ThrowPreparedCancellationNotSupported() => throw new NotSupportedException( "A datasource-prepared command can have multiple executions; cancel an execution with its token."); - internal void OnFlowStarted(CommandFlow flow) => Volatile.Write(ref _activeFlow, flow); + internal void OnFlowStarted(AdoCommandExecutionFlow flow) => Volatile.Write(ref _activeFlow, flow); - internal void OnFlowCompleting(CommandFlow flow, Exception? exception) + internal void OnFlowCompleting(AdoCommandExecutionFlow flow, Exception? exception) { // A flow-level fault while holding an ADO connection lease breaks that lease. SQL errors don't // reach here: they surface on CommandResult and the flow completes cleanly. OnCompleting runs diff --git a/Slon/Ado/AdoCommandFactory.cs b/Slon/Ado/AdoCommandFactory.cs index a89bcb3..0d96a2f 100644 --- a/Slon/Ado/AdoCommandFactory.cs +++ b/Slon/Ado/AdoCommandFactory.cs @@ -1,6 +1,7 @@ using System.Data; using System.Data.Common; using System.Diagnostics; +using System.Runtime.CompilerServices; using Slon.Pg; using Slon.Pg.Serialization; using Slon.Pg.Types; @@ -22,6 +23,49 @@ interface IAdoCommand static class AdoCommandFactory { + [MethodImpl(MethodImplOptions.NoInlining)] + internal static Command CreatePreparedCommandWithParameters( + in TCommand command, TrackedCommand tracked, in CommandDescriptor descriptor, + SlonParameters? commandParameters, bool enableErrorBarriers, CommandBehavior behavior, + DbParameterCollection? dbParameters, TimeSpan timeout, + PgSerializerOptions serializerOptions, ParameterWriter parameterWriter) + where TCommand : IAdoCommand + { + if (tracked.Kind is not TrackedCommandKind.Command + && dbParameters is not null && commandParameters is { Count: > 0 }) + { + throw new InvalidOperationException( + "Execution parameters cannot be combined with parameters stored on the command."); + } + + dbParameters ??= commandParameters; + if (dbParameters is not null and not SlonParameters) + { + throw new ArgumentException( + $"Execution parameters must be a {nameof(SlonParameters)} instance.", nameof(dbParameters)); + } + if ((dbParameters?.Count ?? 0) != descriptor.ParameterTypes.Count) + { + throw new InvalidOperationException( + $"Prepared command expects {descriptor.ParameterTypes.Count} parameters, " + + $"received {dbParameters?.Count ?? 0}."); + } + + var parameters = dbParameters is { Count: > 0 } + ? ResolveNonEmptyParameters( + (SlonParameters)dbParameters, serializerOptions, descriptor.ParameterTypes, + allowUnspecified: false, parameterWriter) + : default; + return new Command + { + Descriptor = descriptor, + DescribeOnly = behavior.HasFlag(CommandBehavior.SchemaOnly), + WithSync = enableErrorBarriers || command.AppendErrorBarrier, + Parameters = parameters, + Timeout = timeout + }; + } + public static (Command, TrackerResult) CreateCommand(in TCommand command, bool allowAutoPreparation, bool enableErrorBarriers, CommandBehavior behavior, in TrackerContext trackerContext, DbParameterCollection? dbParameters, TimeSpan timeout, @@ -66,23 +110,9 @@ public static (Command, TrackerResult) CreateCommand(in TCommand comma { if (serializerOptions is null) ThrowHelper.ThrowInvalidOperation("ADO parameter serialization requires serializer options."); - using var preparedTypes = preparedParameterTypes.GetEnumerator(); - var parameterIndex = 0; - foreach (var kv in slonParameters!.GetStructEnumerator()) - { - if (kv.Key != SlonParameters.PositionalName) - { - throw new NotSupportedException( - "Named parameters are not yet supported; they require client-side SQL parsing."); - } - - var currentParameterIndex = parameterIndex++; - var preparedType = preparedTypes.MoveNext() ? preparedTypes.Current : (PgTypeId?)null; - slonParameters.GetOrResolveTypeInfo( - currentParameterIndex, serializerOptions, preparedType, allowUnspecified: preparing); - } - - parameters = new(slonParameters!, + parameters = ResolveNonEmptyParameters( + slonParameters!, serializerOptions, preparedParameterTypes, + allowUnspecified: preparing, parameterWriter ?? throw new InvalidOperationException( "ADO parameter serialization requires a parameter writer.")); parameterTypes = new(parameters); @@ -113,4 +143,28 @@ public static (Command, TrackerResult) CreateCommand(in TCommand comma Timeout = timeout }, trackerResult); } + + static ParameterSource ResolveNonEmptyParameters( + SlonParameters parameters, PgSerializerOptions serializerOptions, + ParameterTypeList preparedParameterTypes, bool allowUnspecified, + ParameterWriter parameterWriter) + { + using var preparedTypes = preparedParameterTypes.GetEnumerator(); + var parameterIndex = 0; + foreach (var parameter in parameters.GetStructEnumerator()) + { + if (parameter.Key != SlonParameters.PositionalName) + { + throw new NotSupportedException( + "Named parameters are not yet supported; they require client-side SQL parsing."); + } + + var currentParameterIndex = parameterIndex++; + var preparedType = preparedTypes.MoveNext() ? preparedTypes.Current : (PgTypeId?)null; + parameters.GetOrResolveTypeInfo( + currentParameterIndex, serializerOptions, preparedType, allowUnspecified); + } + + return new(parameters, parameterWriter); + } } diff --git a/Slon/Ado/AdoCommandFlow.cs b/Slon/Ado/AdoCommandFlow.cs index 39ee129..4afa18b 100644 --- a/Slon/Ado/AdoCommandFlow.cs +++ b/Slon/Ado/AdoCommandFlow.cs @@ -1,5 +1,6 @@ using System.Data; using System.Data.Common; +using System.Threading.Tasks.Sources; using Slon.Pg; using Slon.Pg.Protocol; using Slon.Pg.Protocol.Flows; @@ -15,40 +16,23 @@ readonly struct AdoCommandFlowOptions internal object? ResultObserverState { get; init; } } -sealed class AdoCommandFlowObserver : CommandFlowObserver - where TCommand : IAdoCommand +interface IAdoCommandExecutionOwner { - internal static readonly AdoCommandFlowObserver Instance = new(); - - protected internal override void OnStarted(CommandFlow flow, object? state) - { - switch (((AdoCommandFlow)flow).LifetimeOwner) - { - case SlonCommand command: - command.OnFlowStarted(flow); - break; - case SlonBatch batch: - batch.OnFlowStarted(flow); - break; - } - } - - protected internal override void OnCommandResult(CommandFlow flow, CommandResult result, object? state) - => ((AdoCommandFlow)flow).ObserveResult(result); + AdoCommandFlowOptions CreateExecutionOptions( + DbParameterCollection? parameters, CommandBehavior behavior, + SlonDataSource.PgDbDependencies dependencies, SlonConnection? connection, + PgConnection pgConnection, TimeSpan? pendingTimeout, bool preparing); + void OnFlowStarted(AdoCommandExecutionFlow flow); + void OnFlowCompleting(AdoCommandExecutionFlow flow, Exception? exception); +} - protected internal override void OnCompleting(PgClientFlow flow, Exception? exception, object? state) - { - switch (((AdoCommandFlow)flow).LifetimeOwner) - { - case SlonCommand command: - command.OnFlowCompleting((CommandFlow)flow, exception); - break; - case SlonBatch batch: - batch.OnFlowCompleting((CommandFlow)flow, exception); - break; - } - } +sealed class AdoCommandExecutionObserver : PgClientFlowObserver +{ + internal static readonly AdoCommandExecutionObserver Instance = new(); + protected internal override void OnCompleting( + PgClientFlow flow, Exception? exception, object? state) + => ((AdoCommandExecutionFlow)flow).CompleteLifetime(exception); } static class AdoCommandResultObserver @@ -123,65 +107,162 @@ internal static void DispatchIndexed(CommandResult result, object? state) } } -sealed class AdoCommandFlow : CommandFlow - where TCommand : IAdoCommand +sealed class AdoCommandExecutionFlow : PgClientFlow, IValueTaskSource, IValueTaskSource { - readonly FieldRef> _core; + readonly IAdoCommandExecutionOwner _bindingOwner; readonly DbParameterCollection? _parameters; readonly CommandBehavior _behavior; readonly SlonDataSource.PgDbDependencies _dependencies; readonly SlonConnection? _connection; readonly bool _preparing; readonly int _commandCount; - object? _lifetimeOwner; + int _lifetimePending; Action? _resultObserver; object? _resultObserverState; + CommandExecutionState _state; - internal AdoCommandFlow( - bool async, FieldRef> core, + internal AdoCommandExecutionFlow( + bool async, IAdoCommandExecutionOwner bindingOwner, DbParameterCollection? parameters, CommandBehavior behavior, SlonDataSource.PgDbDependencies dependencies, SlonConnection? connection, - TimeSpan? pendingTimeout, bool preparing, int commandCount, object? lifetimeOwner) - : base(async, pendingTimeout) + TimeSpan? pendingTimeout, bool preparing, int commandCount, + bool ownsLifetime) + : base(supportsDeferredFlush: true) { - _core = core; + _bindingOwner = bindingOwner; _parameters = parameters; _behavior = behavior; _dependencies = dependencies; _connection = connection; _preparing = preparing; _commandCount = commandCount; - _lifetimeOwner = lifetimeOwner; - SetObserver(AdoCommandFlowObserver.Instance, null); - AdoCommandFlowObserver.Instance.OnStarted(this, null); + _lifetimePending = ownsLifetime ? 1 : 0; + _state.CommandIndex = -1; + _state.EnableActivationTimeout = true; + _state.WaitForDrainOnDispose = true; + _state.PendingTimeout = pendingTimeout; + IsAsync = async; + if (!async) + _state.HandoffEvent = new(false); + SetObserver(AdoCommandExecutionObserver.Instance, null); + if (ownsLifetime) + bindingOwner.OnFlowStarted(this); + } + + internal override bool DefersSyncHandoff => true; + private protected override FlowHandoffEvent? HandoffEvent => _state.HandoffEvent; + protected override bool EnableActivationTimeout => true; + protected override TimeSpan? PendingTimeout => _state.PendingTimeout; + internal override TimeSpan? BackendCancellationGracePeriod + => Volatile.Read(ref _state.ConsumerDetached) ? TimeSpan.FromSeconds(1) : null; + + internal override void BindCallerToken(CancellationToken cancellationToken) + => _state.FlowToken = cancellationToken; + internal override CancellationToken MigrationCancellationToken => _state.FlowToken; + + internal int VisibleCommandCount => _commandCount; + internal CommandResult? CurrentResult => _state.Current; + internal bool IsResultReady => Core.IsResultReady; + internal bool HasCancellationState => _state.ColdState is not null; + + public Enumerator GetEnumerator() => new(this, default); + + public Enumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) + { + if (cancellationToken.CanBeCanceled) + _state.FlowToken = cancellationToken; + return new(this, cancellationToken); } - internal override int VisibleCommandCount => _commandCount; - internal object? LifetimeOwner => _lifetimeOwner; + CommandFlowCore Core => new(new(this)); + + internal ValueTask ConsumeNonQueryAsync(CancellationToken cancellationToken = default) + => Core.ConsumeNonQueryAsync(cancellationToken); + + protected override ValueTask ExecuteAuto(Context context) => Core.ExecuteAuto(context); + internal Task CancelAsync() => Core.CancelAsync(); + internal override bool ResetsSharedReadStateBeforeRelease => true; + protected override void OnStopping(Exception exception) => Core.OnStopping(exception); + protected override void OnAbort(Exception exception) => Core.OnAbort(exception); + internal override void Fail(Exception exception) => Core.Fail(exception); + protected override void OnReleasing(Exception? exception) => Core.OnReleasing(exception); + protected override void OnDiscarded() => Core.OnDiscarded(); + protected override void OnReset() => Core.OnReset(); - internal void ObserveResult(CommandResult result) + void ObserveResult(CommandResult result) => _resultObserver?.Invoke(result, _resultObserverState); + internal void CompleteLifetime(Exception? exception) + { + if (Interlocked.Exchange(ref _lifetimePending, 0) is not 0) + _bindingOwner.OnFlowCompleting(this, exception); + } + internal override void Bind(PgClientFlowBindingContext? context) { var pgConnection = (context as PgConnection.FlowBindingContext)?.Connection ?? throw new InvalidOperationException( "An ADO command requires a PgConnection flow binding context."); - ref var core = ref _core.Invoke(); - InitializeAdo(IsAsync, core.CreateAdoCommandFlowOptions( - [_parameters], _behavior, _dependencies, _connection, pgConnection, - pendingTimeout: PendingTimeout, preparing: _preparing)); + var options = _bindingOwner.CreateExecutionOptions( + _parameters, _behavior, _dependencies, _connection, pgConnection, + PendingTimeout, _preparing); + _resultObserver = options.ResultObserver; + _resultObserverState = options.ResultObserverState; + _state.Commands = options.Commands; + _state.PendingTimeout = options.PendingTimeout; } - void InitializeAdo(bool async, in AdoCommandFlowOptions options) + readonly struct Ops(AdoCommandExecutionFlow owner) : ICommandExecutionFlowOps { - _resultObserver = options.ResultObserver; - _resultObserverState = options.ResultObserverState; - Initialize(async, new CommandFlowOptions + readonly AdoCommandExecutionFlow _owner = owner; + + public static Ops Create(PgClientFlow flow) => new((AdoCommandExecutionFlow)flow); + public PgClientFlow Flow => _owner; + public ref CommandExecutionState GetField() => ref _owner._state; + public bool IsAsync { - Commands = options.Commands, - PendingTimeout = options.PendingTimeout - }); + get => _owner.IsAsync; + set => _owner.IsAsync = value; + } + public bool IsAsyncAtDispatch => _owner.IsAsyncAtDispatch; + public bool HasSuccessfulActivation => _owner.HasSuccessfulActivation; + public void WaitForSyncHandoff() => _owner.WaitForSyncHandoff(); + public void OnCommandResult(CommandResult result) => _owner.ObserveResult(result); + public void OnDrainStarted() { } + public void OnDiscarded() => _owner.CompleteLifetime(null); + } + + bool IValueTaskSource.GetResult(short token) => _state.ReadySource.GetResult(token); + ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) + => _state.ReadySource.GetStatus(token); + void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, + ValueTaskSourceOnCompletedFlags flags) + => _state.ReadySource.OnCompleted(continuation, state, token, flags); + + void IValueTaskSource.GetResult(short token) => _state.PipelineTaskSource.GetResult(token); + ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) + => _state.PipelineTaskSource.GetStatus(token); + void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, + ValueTaskSourceOnCompletedFlags flags) + => _state.PipelineTaskSource.OnCompleted(continuation, state, token, flags); + + public readonly struct Enumerator : IAsyncEnumerator, IDisposable + { + readonly AdoCommandExecutionFlow? _flow; + readonly CancellationToken _cancellationToken; + + internal Enumerator(AdoCommandExecutionFlow flow, CancellationToken cancellationToken) + => (_flow, _cancellationToken) = (flow, cancellationToken); + + internal bool IsDefault => _flow is null; + + public bool MoveNext() => _flow?.Core.MoveNext() ?? false; + public ValueTask MoveNextAsync() => MoveNextAsync(_cancellationToken); + public ValueTask MoveNextAsync(CancellationToken cancellationToken) + => _flow is null ? new(false) : _flow.Core.MoveNextAsync(cancellationToken); + public CommandResult Current => _flow?._state.Current ?? default!; + public ValueTask DisposeAsync() => _flow is null ? default : _flow.Core.DisposeAsync(); + public void Dispose() => _flow?.Core.Dispose(); } } diff --git a/Slon/Ado/AdoCommandFlowFactory.cs b/Slon/Ado/AdoCommandFlowFactory.cs index c7ccc12..a8e2636 100644 --- a/Slon/Ado/AdoCommandFlowFactory.cs +++ b/Slon/Ado/AdoCommandFlowFactory.cs @@ -1,6 +1,7 @@ using System.Buffers; using System.Data; using System.Data.Common; +using System.Runtime.CompilerServices; using Slon.Pg; using Slon.Pg.Protocol.Flows; @@ -33,6 +34,26 @@ public AdoCommandFlowOptions Create( ThrowHelper.ThrowArgumentException(nameof(parametersSpan), "The number of parameter collections must match the number of commands."); var pendingPrefix = connection?.TakePendingTransactionStatement(); + if (commands.Length is 1 && explicitlyPrepared && !preparing && pendingPrefix is null) + return CreatePreparedSingle( + parametersSpan.IsEmpty ? null : parametersSpan[0], + behavior, enableErrorBarriers, + timeout, pgConnection, pendingTimeout); + + return CreateGeneral( + parametersSpan, behavior, explicitlyPrepared, allowAutoPreparation, enableErrorBarriers, + timeout, connection, pgConnection, pendingTimeout, preparing, + pendingPrefix, indexParameters); + } + + AdoCommandFlowOptions CreateGeneral( + ReadOnlySpan parametersSpan, CommandBehavior behavior, + bool explicitlyPrepared, bool allowAutoPreparation, bool enableErrorBarriers, TimeSpan timeout, + SlonConnection? connection, PgConnection? pgConnection, + TimeSpan? pendingTimeout, bool preparing, string? pendingPrefix, + bool indexParameters) + { + var commands = _commands; var commandOffset = pendingPrefix is null ? 0 : 1; var commandCount = commands.Length + commandOffset; var commandArray = commandCount > 1 ? ArrayPool.Shared.Rent(commandCount) : null; @@ -149,6 +170,96 @@ public AdoCommandFlowOptions Create( } } + [MethodImpl(MethodImplOptions.NoInlining)] + AdoCommandFlowOptions CreatePreparedSingle( + DbParameterCollection? parameters, CommandBehavior behavior, + bool enableErrorBarriers, TimeSpan timeout, + PgConnection? pgConnection, TimeSpan? pendingTimeout) + { + ref var adoCommand = ref _commands[0]; + var tracked = adoCommand.Tracked + ?? throw new InvalidOperationException( + "The explicitly prepared command has no tracked command."); + if (!tracked.TryGetPreparedDescriptor(out var descriptor)) + throw new InvalidOperationException( + "The explicitly prepared command has no prepared descriptor."); + var commandParameters = adoCommand.Parameters; + var command = parameters is null && commandParameters is null + && descriptor.ParameterTypes.Count is 0 + ? new Command + { + Descriptor = descriptor, + DescribeOnly = behavior.HasFlag(CommandBehavior.SchemaOnly), + WithSync = enableErrorBarriers || adoCommand.AppendErrorBarrier, + Parameters = default, + Timeout = timeout + } + : AdoCommandFactory.CreatePreparedCommandWithParameters( + adoCommand, tracked, descriptor, commandParameters, enableErrorBarriers, + behavior, parameters, timeout, + dependencies.SerializerOptions, dependencies.ParameterWriter); + + Action? resultAction = null; + object? resultActionState = null; + if (pgConnection is not null) + { + var status = pgConnection.GetTrackedStatus(tracked); + if (status is TrackedStatus.Tracked) + { + resultAction = AdoCommandResultObserver.AttachPrepared; + resultActionState = pgConnection; + } + else + { + var preparation = ResolvePreparedSingleSlow( + adoCommand, command, tracked, pgConnection, status); + command = preparation.Command; + resultAction = preparation.ResultAction; + resultActionState = preparation.ResultActionState; + } + } + + return new() + { + ResultObserver = resultAction, + ResultObserverState = resultActionState, + Commands = new(command), + PendingTimeout = pendingTimeout + }; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + static PreparationResolution ResolvePreparedSingleSlow( + TCommand adoCommand, Command command, TrackedCommand tracked, + PgConnection connection, TrackedStatus status) + { + var parameterTypes = command.Descriptor.ParameterTypes; + if (status is TrackedStatus.Preparing || !connection.TryBeginPreparing(tracked)) + { + return new(command with + { + Descriptor = CommandDescriptor.Create( + adoCommand.CommandText, parameterTypes, default) + }, null, null); + } + + try + { + command = command with + { + Descriptor = CommandDescriptor.Create( + adoCommand.CommandText, parameterTypes, tracked.CommandName) + }; + return new(command, AdoCommandResultObserver.ObservePreparing, + (connection, tracked, (SlonBatchCommand?)null)); + } + catch + { + connection.RemoveTracked(tracked); + throw; + } + } + static PreparationResolution ResolvePreparation( TCommand adoCommand, Command command, TrackedCommand tracked, SlonBatchCommand? batchCommand, PgConnection connection, diff --git a/Slon/Ado/AdoConnectionProxy.cs b/Slon/Ado/AdoConnectionProxy.cs index 36bc790..2b39c18 100644 --- a/Slon/Ado/AdoConnectionProxy.cs +++ b/Slon/Ado/AdoConnectionProxy.cs @@ -48,7 +48,7 @@ internal ConnectionState State return ConnectionState.Open; var activatedFlow = scope.ActivatedFlow; - if (activatedFlow is CommandFlow { IsResultReady: true }) + if (activatedFlow is AdoCommandExecutionFlow { IsResultReady: true }) return ConnectionState.Fetching; return activatedFlow is not null || scope.ExecutingFlow is not null diff --git a/Slon/Ado/TrackedCommand.cs b/Slon/Ado/TrackedCommand.cs index f4b6dca..28ff087 100644 --- a/Slon/Ado/TrackedCommand.cs +++ b/Slon/Ado/TrackedCommand.cs @@ -65,7 +65,8 @@ internal bool TryGetPreparedDescriptor(out CommandDescriptor descriptor) // Sample across this thread's lookups: hot commands remain recent without a clock read and // shared timestamp write on every execution. - if ((++accessSampleCounter & AccessSampleMask) is 0) + if (Kind is TrackedCommandKind.Auto + && (++accessSampleCounter & AccessSampleMask) is 0) Volatile.Write(ref _lastAccessedTicks, Environment.TickCount64); descriptor = state.Descriptor; return true; diff --git a/Slon/Pg/CommandDescriptor.cs b/Slon/Pg/CommandDescriptor.cs index d7a22af..c4dc02c 100644 --- a/Slon/Pg/CommandDescriptor.cs +++ b/Slon/Pg/CommandDescriptor.cs @@ -76,4 +76,13 @@ public static CommandDescriptor CreatePrepared(EncodedCString commandName, Param public static CommandDescriptor Create(string commandText, ParameterTypeList parameterTypes = default, EncodedCString commandName = default) => new(commandText, parameterTypes, commandName); + + // Fieldwise equivalent of assignment, avoiding stores for reference components already present. + internal static void Assign(ref CommandDescriptor destination, in CommandDescriptor value) + { + if (!ReferenceEquals(destination._rowDescriptionOrCommandText, value._rowDescriptionOrCommandText)) + Unsafe.AsRef(in destination._rowDescriptionOrCommandText) = value._rowDescriptionOrCommandText; + EncodedCString.Assign(ref Unsafe.AsRef(in destination._commandName), in value._commandName); + ParameterTypeList.Assign(ref Unsafe.AsRef(in destination._parameterTypes), in value._parameterTypes); + } } diff --git a/Slon/Pg/CommandResult.cs b/Slon/Pg/CommandResult.cs index f4ba2f8..dd34669 100644 --- a/Slon/Pg/CommandResult.cs +++ b/Slon/Pg/CommandResult.cs @@ -1,9 +1,11 @@ +using System.Buffers.Binary; using System.Collections; using System.Diagnostics; using System.Diagnostics.CodeAnalysis; using System.Runtime.CompilerServices; using Slon.Pg.Protocol; using Slon.Pg.Protocol.Flows; +using Slon.Pg.Types; namespace Slon.Pg; @@ -34,16 +36,16 @@ public enum RowBuffering : byte PgError? _errorMessage; Action? _completionAction; object? _completionActionState; - CommandFlow _flow = null!; + PgClientFlow _flow = null!; // The requested row description is what was returned for this exact command (i.e. commands that requested a describe). - internal void Initialize(CommandFlow flow, int index, CommandDescriptor descriptor, + internal void Initialize(PgClientFlow flow, int index, CommandDescriptor descriptor, RowDescription? requestedRowDescription, bool requestedExecution, bool simpleProtocol, PgError? error = null) { if (!ReferenceEquals(_flow, flow)) _flow = flow; _index = index; - _descriptor = descriptor; + CommandDescriptor.Assign(ref _descriptor, in descriptor); // If the command wasn't redescribed, and the prepared description is valid use it instead. var rowDescription = requestedRowDescription; @@ -101,6 +103,21 @@ public RowEnumerator GetAsyncEnumerator(RowBuffering buffering, CancellationToke return new(this, buffering); } + /// + /// Retains backend-message memory for this command result until it is released. + /// + /// + /// Retention may cause subsequent messages to be buffered. Memory returned by + /// remains valid until this command result is released. + /// + public void EnableResultBuffering() + { + if (_firstRowEnumerated) + ThrowHelper.ThrowInvalidOperation( + "Result buffering must be enabled before row enumeration begins."); + _messageEnumerator.EnableResultBuffering(); + } + public bool TryGetCommandComplete([NotNullWhen(true)]out CommandCompleteMessage? value) { // For commands without rows we enumerate once ourselves. @@ -192,6 +209,9 @@ public long RecordsAffected } } + internal long BatchRecordsAffected + => _commandCompleteMessage?.BatchRecordsAffected ?? -1; + internal void CompleteNonQuery(BackendMessage message) { if (message.Header.Type is PgTypes.BackendType.DataRow) @@ -221,12 +241,19 @@ internal void Complete() EnsureComplete(); } - [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] - internal async ValueTask CompleteAsync() + internal ValueTask CompleteAsync() { if (IsComplete) - return; + return default; + return CompleteAsyncCore(); + } + + [MethodImpl(MethodImplOptions.NoInlining)] + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] + async ValueTask CompleteAsyncCore() + { await _row.RevokeColumnLeaseAsync().ConfigureAwait(false); while (await MoveNextMessageAsync().ConfigureAwait(false)) { @@ -241,6 +268,110 @@ internal async ValueTask CompleteAsync() EnsureComplete(); } + /// A fully buffered PostgreSQL data row supplied to a collection callback. + /// + /// The view normally remains valid only for the duration of its callback. Calling + /// before retains + /// collected views until this command result is released. + /// + [Experimental(ExperimentalDiagnostics.PostgreSqlLowerLayer)] + public readonly struct RowView + { + readonly ReadOnlyMemory _memory; + + internal RowView(ReadOnlyMemory memory) + => _memory = memory; + + /// Decodes the field at as . + public T GetValue(int ordinal) + => BootstrapFieldDecoder.Read(GetFieldSpan(ordinal)); + + /// Borrows the field's raw PostgreSQL representation as contiguous memory. + /// + /// The memory normally remains valid only for the duration of the collection callback. + /// Calling before collection retains it until this + /// command result is released. + /// + public ReadOnlyMemory BorrowFieldMemory(int ordinal) + => GetFieldMemory(ordinal); + + /// Decodes the Int32 field at . + public int GetInt32(int ordinal) + { + if (ordinal == 0) + { + var row = _memory.Span; + const int valueOffset = sizeof(short) + sizeof(int); + if (row.Length >= valueOffset + sizeof(int) + && BinaryPrimitives.ReadInt32BigEndian(row[sizeof(short)..]) == sizeof(int)) + return BinaryPrimitives.ReadInt32BigEndian(row[valueOffset..]); + ThrowHelper.ThrowInvalidOperation("The first DataRow field is not a non-null Int32."); + } + return BinaryPrimitives.ReadInt32BigEndian(GetFieldSpan(ordinal)); + } + + ReadOnlySpan GetFieldSpan(int ordinal) => GetFieldMemory(ordinal).Span; + + ReadOnlyMemory GetFieldMemory(int ordinal) + { + ArgumentOutOfRangeException.ThrowIfNegative(ordinal); + var fields = _memory.Span; + if (fields.Length >= sizeof(short)) + { + var offset = sizeof(short); + for (var index = 0; ; index++) + { + if (fields.Length - offset < sizeof(int)) + ThrowHelper.ThrowInvalidOperation("The DataRow field length is truncated."); + var length = BinaryPrimitives.ReadInt32BigEndian(fields[offset..]); + offset += sizeof(int); + if (length < 0) + { + if (index == ordinal) + ThrowHelper.ThrowInvalidOperation("The requested field is null."); + continue; + } + if ((uint)length > (uint)(fields.Length - offset)) + ThrowHelper.ThrowInvalidOperation("The DataRow field is truncated."); + if (index == ordinal) + return _memory.Slice(offset, length); + offset += length; + } + } + + ThrowHelper.ThrowInvalidOperation("The DataRow field count is truncated."); + return default; + } + } + + /// Collects every row through a synchronous callback. + /// + /// Each complete DataRow body is buffered before is invoked. If the + /// collector throws, the result is drained before that exception is rethrown. Call + /// first when collected values must + /// remain usable after their callback returns. + /// + /// State passed to every collector invocation. + /// The synchronous callback invoked once for each row. + /// A token for cancelling collection. + [Experimental(ExperimentalDiagnostics.PostgreSqlLowerLayer)] + public async ValueTask CollectAsync( + TState state, Action collector, + CancellationToken cancellationToken = default) + { + ArgumentNullException.ThrowIfNull(collector); + cancellationToken.ThrowIfCancellationRequested(); + if (_firstRowEnumerated) + ThrowHelper.ThrowInvalidOperation( + "Rows have already been consumed from this command result."); + _firstRowEnumerated = true; + var terminal = await _messageEnumerator + .CollectRowsAsync(state, collector, cancellationToken) + .ConfigureAwait(false); + CompleteTerminal(terminal); + _messageEnumerator.ThrowCollectorException(); + } + void EnsureComplete() { if (_requestedExecution && !IsComplete) @@ -289,21 +420,24 @@ void CompleteCommand(BackendMessage message) break; } - InvokeCompletionAction(); + if (_completionAction is not null) + InvokeCompletionAction(); } + [MethodImpl(MethodImplOptions.NoInlining)] void InvokeCompletionAction() { - if (_completionAction is { } action) + var action = _completionAction!; + var state = _completionActionState; + _completionAction = null; + _completionActionState = null; + try { - var state = _completionActionState; - _completionAction = null; - _completionActionState = null; - try { action(this, state); } - catch (Exception ex) - { - _flow.Fail(ex); - } + action(this, state); + } + catch (Exception ex) + { + _flow.Fail(ex); } } @@ -314,8 +448,9 @@ Row GetRow() } BackendMessage GetCurrentMessage() => _messageEnumerator.Current; + BackendMessage.Accessor GetCurrentMessageAccessor() => _messageEnumerator.CurrentAccessor; bool MoveNextMessage() => _messageEnumerator.MoveNext(); - CommandFlow.MoveNextStatus TryMoveNextMessage() => _messageEnumerator.TryMoveNext(); + CommandFlow.MoveNextStatus TryMoveNextMessage() => _messageEnumerator.TryMoveNextRow(); ValueTask MoveNextMessageAsync() => _messageEnumerator.MoveNextAsync(); public struct RowEnumerator : IEnumerator, IAsyncEnumerator @@ -327,17 +462,17 @@ public struct RowEnumerator : IEnumerator, IAsyncEnumerator internal RowEnumerator(CommandResult instance, RowBuffering buffering) => (_instance, _buffering) = (instance, buffering); - BackendMessage PrepareRow(BackendMessage message) + BackendMessage.Accessor PrepareRow(BackendMessage.Accessor message) { if (_buffering is RowBuffering.Buffered && !message.Buffered) { message.BufferBody(); - message = _instance!.GetCurrentMessage(); + message = _instance!.GetCurrentMessageAccessor(); } return message; } - bool PublishRow(in BackendMessage message) + bool PublishRow(in BackendMessage.Accessor message) { (_row ??= _instance!.GetRow()).InitializeRow(message); return true; @@ -362,8 +497,8 @@ public bool MoveNext() // https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY // "Therefore, an Execute phase is always terminated by the appearance of exactly one of these messages: // CommandComplete, EmptyQueryResponse (if the portal was created from an empty query string), ErrorResponse, or PortalSuspended" - var current = instance.GetCurrentMessage(); - if (current.Header.Type is PgTypes.BackendType.DataRow) + var current = instance.GetCurrentMessageAccessor(); + if (current.Type is PgTypes.BackendType.DataRow) return PublishRow(PrepareRow(current)); return HandleUncommon(current); @@ -393,21 +528,14 @@ public ValueTask MoveNextAsync() return MoveNextAfterRevokeAsync(leasedRow); var status = instance.TryMoveNextMessage(); - if (status is CommandFlow.MoveNextStatus.RequiresInput) - return MoveNextAsyncCore(instance.MoveNextMessageAsync()); - - if (status is CommandFlow.MoveNextStatus.EndOfSequence) - { - if (instance._requestedExecution && instance._commandCompleteMessage is null && instance._errorMessage is null) - ThrowHelper.ThrowInvalidOperation("Underlying message enumerator completed before CommandComplete was returned."); - return new(false); - } + if (status is not CommandFlow.MoveNextStatus.Moved) + return HandleNonMoved(instance, status); // https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY // "Therefore, an Execute phase is always terminated by the appearance of exactly one of these messages: // CommandComplete, EmptyQueryResponse (if the portal was created from an empty query string), ErrorResponse, or PortalSuspended" - var current = instance.GetCurrentMessage(); - if (current.Header.Type is PgTypes.BackendType.DataRow) + var current = instance.GetCurrentMessageAccessor(); + if (current.Type is PgTypes.BackendType.DataRow) { if (_buffering is RowBuffering.Buffered && !current.Buffered) return BufferCurrentRow(in current); @@ -418,13 +546,24 @@ public ValueTask MoveNextAsync() } [MethodImpl(MethodImplOptions.NoInlining)] - ValueTask BufferCurrentRow(in BackendMessage current) + ValueTask HandleNonMoved(CommandResult instance, CommandFlow.MoveNextStatus status) + { + if (status is CommandFlow.MoveNextStatus.RequiresInput) + return MoveNextAsyncCore(instance.MoveNextMessageAsync()); + + if (instance._requestedExecution && instance._commandCompleteMessage is null && instance._errorMessage is null) + ThrowHelper.ThrowInvalidOperation("Underlying message enumerator completed before CommandComplete was returned."); + return new(false); + } + + [MethodImpl(MethodImplOptions.NoInlining)] + ValueTask BufferCurrentRow(in BackendMessage.Accessor current) { var instance = _instance!; var bufferTask = current.BufferBodyAsync(default); if (!bufferTask.IsCompletedSuccessfully) return BufferRowAsync(bufferTask, _row ??= instance.GetRow()); - return new(PublishRow(instance.GetCurrentMessage())); + return new(PublishRow(instance.GetCurrentMessageAccessor())); } [MethodImpl(MethodImplOptions.NoInlining)] @@ -435,16 +574,17 @@ async ValueTask MoveNextAfterRevokeAsync(Row row) } [MethodImpl(MethodImplOptions.NoInlining)] - bool HandleUncommon(in BackendMessage current) + bool HandleUncommon(in BackendMessage.Accessor current) { var instance = _instance!; - var type = current.Header.Type; + var type = current.Type; switch (type) { case PgTypes.BackendType.EmptyQueryResponse: case PgTypes.BackendType.CommandComplete: case PgTypes.BackendType.ErrorResponse: - instance.CompleteCommand(current); + instance._messageEnumerator.MarkCurrentTerminal(); + instance.CompleteCommand(current.Message); return false; case PgTypes.BackendType.PortalSuspended when !instance._simpleProtocol: default: @@ -454,6 +594,7 @@ bool HandleUncommon(in BackendMessage current) } [MethodImpl(MethodImplOptions.NoInlining)] + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] async ValueTask MoveNextAsyncCore(ValueTask task) { @@ -468,23 +609,27 @@ async ValueTask MoveNextAsyncCore(ValueTask task) // https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY // "Therefore, an Execute phase is always terminated by the appearance of exactly one of these messages: // CommandComplete, EmptyQueryResponse (if the portal was created from an empty query string), ErrorResponse, or PortalSuspended" - var current = instance.GetCurrentMessage(); - if (current.Header.Type is PgTypes.BackendType.DataRow) + var current = instance.GetCurrentMessageAccessor(); + if (current.Type is PgTypes.BackendType.DataRow) { if (_buffering is RowBuffering.Buffered && !current.Buffered) + { await current.BufferBodyAsync(default).ConfigureAwait(false); - return PublishRow(instance.GetCurrentMessage()); + current = instance.GetCurrentMessageAccessor(); + } + return PublishRow(current); } return HandleUncommon(current); } [MethodImpl(MethodImplOptions.NoInlining)] + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] async ValueTask BufferRowAsync(ValueTask task, Row row) { await task.ConfigureAwait(false); - row.InitializeRow(_instance!.GetCurrentMessage()); + row.InitializeRow(_instance!.GetCurrentMessageAccessor()); return true; } diff --git a/Slon/Pg/ParameterTypeList.cs b/Slon/Pg/ParameterTypeList.cs index 0c8af02..19748d2 100644 --- a/Slon/Pg/ParameterTypeList.cs +++ b/Slon/Pg/ParameterTypeList.cs @@ -2,6 +2,7 @@ using System.Collections.Immutable; using System.Diagnostics; using System.Diagnostics.CodeAnalysis; +using System.Runtime.CompilerServices; using System.Runtime.InteropServices; using Slon.Pg.Types; @@ -167,4 +168,13 @@ public bool Equals(ParameterTypeList other) public override int GetHashCode() => HashCode.Combine(_source, _writer, _count); public static bool operator ==(ParameterTypeList left, ParameterTypeList right) => left.Equals(right); public static bool operator !=(ParameterTypeList left, ParameterTypeList right) => !left.Equals(right); + + internal static void Assign(ref ParameterTypeList destination, in ParameterTypeList value) + { + if (!ReferenceEquals(destination._source, value._source)) + Unsafe.AsRef(in destination._source) = value._source; + if (!ReferenceEquals(destination._writer, value._writer)) + Unsafe.AsRef(in destination._writer) = value._writer; + Unsafe.AsRef(in destination._count) = value._count; + } } diff --git a/Slon/Pg/PgClientOptions.cs b/Slon/Pg/PgClientOptions.cs index c73cf1a..eb2b022 100644 --- a/Slon/Pg/PgClientOptions.cs +++ b/Slon/Pg/PgClientOptions.cs @@ -43,7 +43,7 @@ public PostgreSqlSslOptions Ssl public TimeSpan ConnectionTimeout { get; init; } = Timeout.InfiniteTimeSpan; internal PgSessionResetOptions SessionReset { get; init; } = new(); - internal int DataRowStreamingThreshold { get; init; } = BackendMessageBatch.Segmenter.DefaultDataRowStreamingThreshold; + internal int DataRowStreamingThreshold { get; init; } = BackendMessageCursor.DefaultDataRowStreamingThreshold; internal int MaxInFlightFlowsPerWire { get; init; } internal PipelineScheduler? ExecutionScheduler { get; init; } diff --git a/Slon/Pg/Protocol/BackendMessage.cs b/Slon/Pg/Protocol/BackendMessage.cs index ac6a3f2..f679c79 100644 --- a/Slon/Pg/Protocol/BackendMessage.cs +++ b/Slon/Pg/Protocol/BackendMessage.cs @@ -13,83 +13,185 @@ public readonly struct BackendMessage { string DebuggerDisplay => $"Type = {Header.Type}, Length = {Header.MessageLength}"; - readonly ReadOnlySequence _buffer; - readonly BackendMessageContext _context; + readonly object? _firstObject; + readonly object? _contextOrEndObject; - // Packed to avoid another 8 bytes. - readonly bool _buffered; - readonly BackendType _type; - readonly short _token; + // Buffered, peeked, type, and token fit in one word so the message remains 32 bytes. + readonly uint _state; readonly int _length; + readonly int _startIndex; + readonly int _endIndexOrBufferedLength; - BackendMessage(BackendHeader header, ReadOnlySequence buffer, BackendMessageContext context, short token, bool buffered) + BackendMessage(BackendHeader header, ReadOnlySequence buffer, + BackendMessageContext? context, short token, bool buffered, + bool independent = false) { - _buffer = buffer; - _context = context; - _buffered = buffered; - _type = header.Type; - _token = token; + _firstObject = buffer.Start.GetObject(); + _contextOrEndObject = independent ? buffer.End.GetObject() : context; + _startIndex = buffer.Start.GetInteger() & int.MaxValue; + _endIndexOrBufferedLength = independent + ? buffer.End.GetInteger() & int.MaxValue + : checked((int)buffer.Length); + _state = (buffered ? 1u : 0) + | (independent ? 2u : 0) + | ((uint)(byte)header.Type << 2) + | ((uint)(ushort)token << 10); _length = header.Length; + if (context is not null) + context.SetCurrentFallbackBuffer(in buffer, + _firstObject is ReadOnlySequenceSegment); } internal BackendMessage(BackendHeader header, ReadOnlySequence buffer, BackendMessageContext context, short token) : this(header, buffer, context, token, buffer.Length >= header.MessageLength) {} + internal static BackendMessage CreateIndependent( + BackendHeader header, ReadOnlySequence buffer) + { + if (buffer.Length < header.MessageLength) + ThrowHelper.ThrowInvalidOperation( + "An independent backend message must be fully buffered."); + return new(header, buffer, context: null, token: 0, + buffered: true, independent: true); + } + + internal static void InitializeIndependent(ref BackendMessage destination, + BackendHeader header, ReadOnlySequence buffer) + { + var value = CreateIndependent(header, buffer); + Assign(ref destination, in value); + } + internal static void Initialize(ref BackendMessage destination, BackendHeader header, ReadOnlySequence buffer, BackendMessageContext context, short token, bool buffered) { var value = new BackendMessage(header, buffer, context, token, buffered); - WriteGranularly(ref destination, in value, destinationIsZero: false); + Assign(ref destination, in value, destinationIsZero: false); + } + + internal static void Initialize(ref BackendMessage destination, BackendHeader header, + in BackendMessageCursor.FastReadOnlySequence buffer, + BackendMessageContext context, short token, bool buffered) + { + var firstObject = buffer.StartObject; + if (!ReferenceEquals(destination._contextOrEndObject, context)) + Unsafe.AsRef(in destination._contextOrEndObject) = context; + if (!ReferenceEquals(destination._firstObject, firstObject)) + Unsafe.AsRef(in destination._firstObject) = firstObject; + + Unsafe.AsRef(in destination._state) = (buffered ? 1u : 0) + | ((uint)(byte)header.Type << 2) + | ((uint)(ushort)token << 10); + Unsafe.AsRef(in destination._length) = header.Length; + Unsafe.AsRef(in destination._startIndex) = buffer.StartIndex; + Unsafe.AsRef(in destination._endIndexOrBufferedLength) = checked((int)buffer.Length); + context.SetCurrentFallbackBuffer(in buffer); } // The JIT should have a phase for picking granular writes (and write barriers) over full struct assignments. // This translation is entirely mechanical (even though these implementations need to deviate for external types). [MethodImpl(MethodImplOptions.AggressiveInlining)] - static void WriteGranularly(ref BackendMessage destination, in BackendMessage value, bool destinationIsZero = false) + static void Assign(ref BackendMessage destination, in BackendMessage value, bool destinationIsZero = false) { - if ((destinationIsZero && value._context is not null) || !ReferenceEquals(destination._context, value._context)) - Unsafe.AsRef(in destination._context) = value._context!; - - WriteGranularly(ref Unsafe.AsRef(in destination._buffer), in value._buffer); + if ((destinationIsZero && value._contextOrEndObject is not null) + || !ReferenceEquals(destination._contextOrEndObject, value._contextOrEndObject)) + Unsafe.AsRef(in destination._contextOrEndObject) = value._contextOrEndObject; + if (!ReferenceEquals(destination._firstObject, value._firstObject)) + Unsafe.AsRef(in destination._firstObject) = value._firstObject; - Unsafe.AsRef(in destination._buffered) = value._buffered; - Unsafe.AsRef(in destination._type) = value._type; - Unsafe.AsRef(in destination._token) = value._token; + Unsafe.AsRef(in destination._state) = value._state; Unsafe.AsRef(in destination._length) = value._length; + Unsafe.AsRef(in destination._startIndex) = value._startIndex; + Unsafe.AsRef(in destination._endIndexOrBufferedLength) = value._endIndexOrBufferedLength; } - [MethodImpl(MethodImplOptions.AggressiveInlining)] - static void WriteGranularly(ref ReadOnlySequence destination, in ReadOnlySequence value) - => destination = value; + BackendType Type => (BackendType)((_state >> 2) & byte.MaxValue); + short Token => (short)(_state >> 10); + bool IsIndependent => (_state & 2) != 0; + internal bool IsDefault => Type == default; + BackendMessageContext Context + => _contextOrEndObject as BackendMessageContext + ?? throw new InvalidOperationException( + "The independent backend message has no decoder context."); [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal static void SetSequence(ref ReadOnlySequence destination, in ReadOnlySequence value) - => WriteGranularly(ref destination, in value); + ReadOnlySequence GetBuffer() + { + if (!IsIndependent) + { + if (_firstObject is byte[] array) + return new(array, _startIndex, _endIndexOrBufferedLength); + if (_firstObject is MemoryManager manager) + return new(manager.Memory.Slice(_startIndex, _endIndexOrBufferedLength)); + return Context.GetFallbackBuffer(Token); + } - BackendType Type => _type; - internal bool IsDefault => _type == default; + if (_firstObject is byte[] independentArray + && ReferenceEquals(_firstObject, _contextOrEndObject)) + return new(independentArray, _startIndex, + _endIndexOrBufferedLength - _startIndex); + if (_firstObject is MemoryManager independentManager + && ReferenceEquals(_firstObject, _contextOrEndObject)) + return new(independentManager.Memory.Slice( + _startIndex, _endIndexOrBufferedLength - _startIndex)); + return new((ReadOnlySequenceSegment)_firstObject!, _startIndex, + (ReadOnlySequenceSegment)_contextOrEndObject!, + _endIndexOrBufferedLength); + } public BackendHeader Header { [MethodImpl(MethodImplOptions.AggressiveInlining)] - get => BackendHeader.CreateUnchecked(_type, _length); + get => BackendHeader.CreateUnchecked(Type, _length); } public ReadOnlySequence GetSequence(SequencePosition start) { EnsureBodyWindowAvailable(); - return _buffer.Slice(start); + return GetBuffer().Slice(start); } public ReadOnlySequence GetSequence(long offset) { EnsureBodyWindowAvailable(); - return _buffer.Slice(BackendHeader.ByteCount + offset); + ArgumentOutOfRangeException.ThrowIfNegative(offset); + var start = checked(BackendHeader.ByteCount + offset); + var length = BufferedLength - start; + if (length < 0 || length > int.MaxValue) + throw new ArgumentOutOfRangeException(nameof(offset)); + + if (!IsIndependent || ReferenceEquals(_firstObject, _contextOrEndObject)) + { + if (_firstObject is byte[] array) + return new(array, checked(_startIndex + (int)start), (int)length); + if (_firstObject is MemoryManager manager) + return new(manager.Memory.Slice( + checked(_startIndex + (int)start), (int)length)); + } + return GetBuffer().Slice(start); } public ReadOnlySequence GetSequence() => GetSequence(0); + internal ReadOnlyMemory GetContiguousMemory( + ReadOnlyMemory source) + { + EnsureBodyWindowAvailable(); + return IsIndependent + ? source + : Context.GetContiguousMemory(Token, source); + } + + internal ReadOnlyMemory GetContiguousMemory( + in ReadOnlySequence source) + { + EnsureBodyWindowAvailable(); + if (IsIndependent) + return source.IsSingleSegment ? source.First : source.ToArray(); + return Context.GetContiguousMemory(Token, source); + } + [MethodImpl(MethodImplOptions.AggressiveInlining)] public bool TryGetFirstSpan(int offset, out ReadOnlySpan span) { @@ -101,17 +203,7 @@ public bool TryGetFirstSpan(int offset, out ReadOnlySpan span) internal bool TryGetFirstSpanUnchecked(int offset, out ReadOnlySpan span) { offset += BackendHeader.ByteCount; - ref var buffer = ref Unsafe.AsRef(in _buffer); - ReadOnlySpan firstSpan; - if (SequenceMarshal.TryGetArray(buffer, out var array)) - { - Debug.Assert(buffer.IsSingleSegment); - firstSpan = array.AsSpan(); - } - else - { - firstSpan = buffer.FirstSpan; - } + var firstSpan = GetFirstMemory().Span; if ((uint)offset <= (uint)firstSpan.Length) { span = firstSpan.Slice(offset); @@ -123,21 +215,45 @@ internal bool TryGetFirstSpanUnchecked(int offset, out ReadOnlySpan span) } [MethodImpl(MethodImplOptions.AggressiveInlining)] - internal bool TryGetBufferedFirstMemory(int offset, out ReadOnlyMemory memory) + internal bool TryGetFirstByte(int offset, out byte value) { - Debug.Assert(_buffered); + Debug.Assert(Buffered); offset += BackendHeader.ByteCount; - ref var buffer = ref Unsafe.AsRef(in _buffer); - ReadOnlyMemory firstMemory; - if (SequenceMarshal.TryGetArray(buffer, out var array)) + var firstLength = IsIndependent + ? _endIndexOrBufferedLength - _startIndex + : _endIndexOrBufferedLength; + if (_firstObject is byte[] array && (uint)offset < (uint)firstLength) + { + value = array[_startIndex + offset]; + return true; + } + + var firstMemory = GetFirstMemory(); + if ((uint)offset < (uint)firstMemory.Length) { - Debug.Assert(buffer.IsSingleSegment); - firstMemory = array.AsMemory(); + value = firstMemory.Span[offset]; + return true; } - else + + value = 0; + return false; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal bool TryGetBufferedFirstMemory(int offset, out ReadOnlyMemory memory) + { + Debug.Assert(Buffered); + offset += BackendHeader.ByteCount; + var firstLength = IsIndependent + ? _endIndexOrBufferedLength - _startIndex + : _endIndexOrBufferedLength; + if (_firstObject is byte[] array && (uint)offset <= (uint)firstLength) { - firstMemory = buffer.First; + memory = array.AsMemory(_startIndex + offset, firstLength - offset); + return true; } + + var firstMemory = GetFirstMemory(); if ((uint)offset <= (uint)firstMemory.Length) { memory = firstMemory.Slice(offset); @@ -148,13 +264,69 @@ internal bool TryGetBufferedFirstMemory(int offset, out ReadOnlyMemory mem return false; } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal bool TryGetBufferedArrayMemory(int offset, out ReadOnlyMemory memory) + { + Debug.Assert(Buffered); + offset += BackendHeader.ByteCount; + var firstLength = IsIndependent + ? _endIndexOrBufferedLength - _startIndex + : _endIndexOrBufferedLength; + if (_firstObject is byte[] array && (uint)offset <= (uint)firstLength) + { + memory = array.AsMemory(_startIndex + offset, firstLength - offset); + return true; + } + + memory = default; + return false; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal bool TryGetBufferedArray(int offset, [NotNullWhen(true)] out byte[]? array, + out int arrayOffset, out int length) + { + Debug.Assert(Buffered); + offset += BackendHeader.ByteCount; + var firstLength = IsIndependent + ? _endIndexOrBufferedLength - _startIndex + : _endIndexOrBufferedLength; + if (_firstObject is byte[] value && (uint)offset <= (uint)firstLength) + { + array = value; + arrayOffset = _startIndex + offset; + length = firstLength - offset; + return true; + } + + array = null; + arrayOffset = 0; + length = 0; + return false; + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + ReadOnlyMemory GetFirstMemory() + { + var length = IsIndependent + ? _endIndexOrBufferedLength - _startIndex + : _endIndexOrBufferedLength; + return _firstObject switch + { + byte[] array => array.AsMemory(_startIndex, length), + MemoryManager manager => manager.Memory.Slice(_startIndex, length), + ReadOnlySequenceSegment segment => segment.Memory.Slice(_startIndex), + _ => GetBuffer().First + }; + } + public SequenceReader BodyReader => new(GetSequence()); [MethodImpl(MethodImplOptions.AggressiveInlining)] void EnsureBodyWindowAvailable() { - if (!_buffered) - _context.EnsureBodyWindowAvailable(_token); + if (!Buffered) + Context.EnsureBodyWindowAvailable(Token); } public (PgError? Error, BackendType Type) EnsureExpectedOrError(params ReadOnlySpan expected) @@ -181,41 +353,61 @@ static void Throw(BackendType actual, ReadOnlySpan expected) $"Unexpected backend message: {actual}, expected: {string.Join(" or ", expected.ToArray())}."); } - public Accessor GetAccessor() => new(_context, _token); + public Accessor GetAccessor() => new(Context, Token, Type, Buffered); internal BackendMessageBodyReader OpenBodyReader() - => new(_context, _token, GetSequence(), Buffered); + => new(Context, Token, GetSequence(), Buffered); internal void BufferBody() { if (!Buffered) - _context.BufferCurrentMessage(_token); + Context.BufferCurrentMessage(Token); } internal ValueTask BufferBodyAsync(CancellationToken cancellationToken) - => Buffered ? default : _context.BufferCurrentMessageAsync(_token, cancellationToken); + => Buffered ? default : Context.BufferCurrentMessageAsync(Token, cancellationToken); public readonly struct Accessor { readonly BackendMessageContext _context; readonly short _token; + readonly PgTypes.BackendType _type; + readonly bool _buffered; - internal Accessor(BackendMessageContext context, short token) + internal Accessor(BackendMessageContext context, short token, + PgTypes.BackendType type, bool buffered) { _context = context; _token = token; + _type = type; + _buffered = buffered; } public BackendMessage Message => _context.GetCurrent(_token); + internal PgTypes.BackendType Type => _type; + internal bool Buffered => _buffered; + internal BackendMessageBodyReader OpenBodyReader() + => _context.OpenCurrentBodyReader(_token); + internal bool TryGetBufferedFirstMemory(int offset, out ReadOnlyMemory memory) + => _context.TryGetCurrentBufferedFirstMemory(_token, offset, out memory); + internal bool TryGetBufferedArray(int offset, [NotNullWhen(true)] out byte[]? array, + out int arrayOffset, out int length) + => _context.TryGetCurrentBufferedArray(_token, offset, out array, out arrayOffset, out length); + internal void BufferBody() => _context.BufferCurrentMessage(_token); + internal ValueTask BufferBodyAsync(CancellationToken cancellationToken) + => _context.BufferCurrentMessageAsync(_token, cancellationToken); + // The JIT should have a phase for picking granular writes (and write barriers) over full struct assignments. // This translation is entirely mechanical (even though these implementations need to deviate for external types). - internal static void WriteGranularly(ref Accessor destination, in Accessor value, bool destinationIsZero = false) + internal static void Assign(ref Accessor destination, in Accessor value, bool destinationIsZero = false) { if ((destinationIsZero && value._context is not null) || !ReferenceEquals(destination._context, value._context)) Unsafe.AsRef(in destination._context) = value._context!; Unsafe.AsRef(in destination._token) = value._token; + Unsafe.AsRef(in destination._type) = value._type; + Unsafe.AsRef(in destination._buffered) = value._buffered; } } @@ -258,6 +450,18 @@ static void Throw(BackendType actual, BackendType expected) => throw new PgProtocolException($"Unexpected backend message: {actual}, expected: {expected}."); } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal void EnsureExpected(BackendType expected1, BackendType expected2) + { + var actual = Type; + if (actual != expected1 && actual != expected2) + Throw(actual, expected1, expected2); + + static void Throw(BackendType actual, BackendType expected1, BackendType expected2) + => throw new PgProtocolException( + $"Unexpected backend message: {actual}, expected: {expected1} or {expected2}."); + } + // Inlining helps as it's usually run over a few RVA items at most. [MethodImpl(MethodImplOptions.AggressiveInlining)] public BackendType EnsureExpected(params ReadOnlySpan expected) @@ -325,23 +529,26 @@ PgError CreateError(ReadOnlySpan expected, bool unhandled = true) } internal void MarkPriorCancellationExposure() - => _context.MarkPriorCancellationExposure(_token); + => Context.MarkPriorCancellationExposure(Token); internal bool HasPriorCancellationExposure - => _context.HasPriorCancellationExposure(_token); + => Context.HasPriorCancellationExposure(Token); internal void MarkBackendTermination() - => _context.MarkBackendTermination(_token); + => Context.MarkBackendTermination(Token); internal bool IsBackendTermination - => _context.IsBackendTermination(_token); + => Context.IsBackendTermination(Token); internal bool TryObserveError() - => _context.TryObserveError(_token); + => Context.TryObserveError(Token); // We have no buffer for header only messages. - public bool Buffered => _buffered; - internal long BufferedLength => _buffer.Length; + public bool Buffered => (_state & 1) != 0; + internal long BufferedLength + => IsDefault ? 0 + : IsIndependent ? Header.MessageLength + : _endIndexOrBufferedLength; } [Experimental(ExperimentalDiagnostics.PostgreSqlLowerLayer)] diff --git a/Slon/Pg/Protocol/BackendMessageBatch.cs b/Slon/Pg/Protocol/BackendMessageBatch.cs deleted file mode 100644 index da1742a..0000000 --- a/Slon/Pg/Protocol/BackendMessageBatch.cs +++ /dev/null @@ -1,220 +0,0 @@ -using System.Buffers; -using System.Diagnostics; -using System.Runtime.CompilerServices; -using System.Runtime.InteropServices; -using Slon.Pipelines; -using static Slon.Pg.Protocol.PgTypes; - -namespace Slon.Pg.Protocol; - -// Note: both the batch and the segmenter are perf sensitive. -struct BackendMessageBatch(ReadOnlySequence buffer) -{ - FastReadOnlySequence _buffer = new(buffer); - long _consumedLength; - - BackendMessageBatch(ReadOnlySequence buffer, long consumedLength) : this(buffer) - => _consumedLength = consumedLength; - - public readonly long ConsumedLength => _consumedLength; - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public bool TryPeekType(out BackendType type) - { - var span = _buffer.FirstSpan; - if (!span.IsEmpty) - { - type = (BackendType)span[0]; - return true; - } - return TryPeekTypeMultiSegment(_buffer.Sequence, out type); - } - - [MethodImpl(MethodImplOptions.NoInlining)] - static bool TryPeekTypeMultiSegment(ReadOnlySequence buffer, out BackendType type) - { - var reader = new SequenceReader(buffer); - if (reader.TryPeek(out var tag)) - { - type = (BackendType)tag; - return true; - } - type = default; - return false; - } - - public bool TryReadNextInPlace(out BackendHeader header, out ReadOnlySequence buffer, out uint bufferLength) - { - if (!Header.TryParse(_buffer.FirstSpan, out var protoHeader) && !Header.TryParseMultiSegment(_buffer.Sequence, out protoHeader)) - { - // We use default(ROSeq) - which is fully supported - as ROSeq.Empty weirdly enough wraps an empty array. - _buffer = default; - buffer = default; - bufferLength = default; - header = default; - return false; - } - - var fastSeq = _buffer.SplitInPlace(Math.Min(_buffer.Length, protoHeader.MessageLength)); - _consumedLength += fastSeq.Length; - buffer = fastSeq.Sequence; - Debug.Assert(fastSeq.Length <= uint.MaxValue); - bufferLength = unchecked((uint)fastSeq.Length); - header = BackendHeader.FromHeader(protoHeader); - return true; - } - - public readonly bool TryReadNext(out BackendHeader header, out ReadOnlySequence buffer, out uint bufferLength, out BackendMessageBatch remaining) - { - var thisCopy = this; - var success = thisCopy.TryReadNextInPlace(out header, out buffer, out bufferLength); - remaining = success ? new(thisCopy._buffer.Sequence, thisCopy._consumedLength) : default; - return success; - } - - // Segmenter parses messages and ensures relevant messages are fully buffered before being returned. - internal struct Segmenter : IPipeSegmenter - { - public const int DefaultDataRowStreamingThreshold = 16 * 1024; - const uint MaxMessageLength = 0x3FFF_FFFF; - - readonly int _dataRowStreamingThreshold; - int _minimumSize; - public int MinimumSize => _minimumSize; - - public Segmenter() : this(DefaultDataRowStreamingThreshold) {} - - public Segmenter(int dataRowStreamingThreshold) - { - ArgumentOutOfRangeException.ThrowIfNegative(dataRowStreamingThreshold); - _dataRowStreamingThreshold = dataRowStreamingThreshold; - } - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public OperationStatus CreateSegment(in ReadOnlySequence buffer, out long segmentLength, out BackendMessageBatch segment) - { - _minimumSize = Header.ByteCount; - var reader = new SequenceReader(buffer); - var messages = 0; - var needMoreData = false; - segmentLength = 0; - - // Try span first before accessing the sequence. - while (Header.TryParse(reader.UnreadSpan, out var header) || Header.TryParseMultiSegment(reader.UnreadSequence, out header)) - { - var backendType = (BackendType)header.Tag; - if (header.MessageLength > MaxMessageLength) - throw new PgFramingException($"PostgreSQL backend message length {header.MessageLength} exceeds the maximum supported length."); - - if (reader.Remaining < header.MessageLength) - { - var required = RequiredBufferedLength(backendType, header.MessageLength); - if (reader.Remaining < required) - { - // MinimumSize is relative to the entire unconsumed pipe buffer, including messages - // already framed before this one. - _minimumSize = int.CreateSaturating(segmentLength + required); - needMoreData = true; - break; - } - - reader.Advance(reader.Remaining); - } - else - { - reader.Advance(header.MessageLength); - } - - messages++; - segmentLength += header.MessageLength; - } - - if (messages is 0) - { - segment = default; - return OperationStatus.NeedMoreData; - } - - segment = new(reader.Length == segmentLength ? buffer : buffer.Slice(0, reader.Position)); - return needMoreData ? OperationStatus.NeedMoreData : OperationStatus.Done; - } - - uint RequiredBufferedLength(BackendType backendType, uint messageLength) => backendType switch - { - BackendType.DataRow => Math.Min(messageLength, (uint)_dataRowStreamingThreshold), - // BackendType.RowDescription or - // BackendType.CopyData or - // BackendType.FunctionCallResponse or - // BackendType.NotificationResponse or - // BackendType.ParameterDescription => false, - _ => messageLength, - }; - } - - // TODO faster firstspan and splitting should be able to be upstreamed. - // Optimizes for faster splitting and length checks. - struct FastReadOnlySequence - { - ReadOnlySequence _sequence; - long _length; - - FastReadOnlySequence(ReadOnlySequence sequence, long length) - { - Debug.Assert(Unsafe.SizeOf>() is 32); - _sequence = sequence; - _length = length; - } - - public FastReadOnlySequence(ReadOnlySequence sequence) - { - Debug.Assert(Unsafe.SizeOf>() is 32); - _sequence = sequence; - _length = sequence.Length; - } - - public ReadOnlySequence Sequence => _sequence; - public long Length => _length; - - public ReadOnlySpan FirstSpan => GetFirstSpan(out _); - - // Returns the sequence before the index, stores the sequence after it in place. - public FastReadOnlySequence SplitInPlace(long offset) - { - FastReadOnlySequence prev; - - // If it's out-of-range of the first, has to resolve another segment, or is not backed by - // one array, let Slice handle it. - if (!SequenceMarshal.TryGetArray(_sequence, out var array) || array.Count <= offset) - { - prev = new(_sequence.Slice(0, offset), offset); - _sequence = _sequence.Slice(offset); - } - else - { - Debug.Assert(offset <= int.MaxValue); - var arrayInstance = array.Array!; - prev = new(new(arrayInstance, array.Offset, (int)offset), offset); - _sequence = new(arrayInstance, array.Offset + (int)offset, array.Count - (int)offset); - } - - _length -= offset; - return prev; - } - - // TODO arrays should not go down the slow path for First and FirstSpan, the SequenceReader variant doesn't either. - // Inline to remove the write barriers. - [MethodImpl(MethodImplOptions.AggressiveInlining)] - ReadOnlySpan GetFirstSpan(out ArraySegment array) - { - if (SequenceMarshal.TryGetArray(_sequence, out array)) - { - Debug.Assert(_sequence.IsSingleSegment); - return array.AsSpan(); - } - - array = default; - return _sequence.FirstSpan; - } - - } -} diff --git a/Slon/Pg/Protocol/BackendMessageBodyReader.cs b/Slon/Pg/Protocol/BackendMessageBodyReader.cs index d293293..7cc40f2 100644 --- a/Slon/Pg/Protocol/BackendMessageBodyReader.cs +++ b/Slon/Pg/Protocol/BackendMessageBodyReader.cs @@ -56,7 +56,7 @@ public void AdvanceTo(SequencePosition consumed, long consumedLength) public bool TryRead() { EnsureAdvanced(); - if (!_context.TryContinue(_token, _consumed, _consumedLength, out var result)) + if (!_context.TrySlide(_token, _consumed, _consumedLength, out var result)) return false; Publish(result); return true; @@ -65,7 +65,8 @@ public bool TryRead() public ValueTask ReadAsync(CancellationToken cancellationToken = default) { EnsureAdvanced(); - var task = _context.ContinueAsync(_token, _consumed, _consumedLength, cancellationToken); + var task = _context.SlideAsync( + _token, _consumed, _consumedLength, cancellationToken); if (task.IsCompletedSuccessfully) { Publish(task.Result); @@ -73,14 +74,16 @@ public ValueTask ReadAsync(CancellationToken cancellationToken = default) } return Core(task); - async ValueTask Core(ValueTask task) + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] + async ValueTask Core(ValueTask task) => Publish(await task.ConfigureAwait(false)); } public void Read() { EnsureAdvanced(); - Publish(_context.Continue(_token, _consumed, _consumedLength)); + Publish(_context.Slide(_token, _consumed, _consumedLength)); } public bool TryExtend() @@ -95,16 +98,19 @@ public bool TryExtend() public ValueTask ExtendAsync(CancellationToken cancellationToken = default) { EnsureCanExtend(); - var task = _context.ExtendAsync(_token, cancellationToken); + var task = _context.BeginExtendAsync(_token, cancellationToken); if (task.IsCompletedSuccessfully) { - Publish(task.Result, retained: true); + Publish(_context.CompleteExtend(_token, task.Result), retained: true); return default; } return Core(task); - async ValueTask Core(ValueTask task) - => Publish(await task.ConfigureAwait(false), retained: true); + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] + async ValueTask Core(ValueTask task) + => Publish(_context.CompleteExtend( + _token, await task.ConfigureAwait(false)), retained: true); } public void Extend() @@ -172,10 +178,19 @@ public ValueTask BufferAllAsync(CancellationToken cancellationToken = default) return default; return Core(cancellationToken); + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] async ValueTask Core(CancellationToken cancellationToken) { while (!IsComplete) - await ExtendAsync(cancellationToken).ConfigureAwait(false); + { + EnsureCanExtend(); + var task = _context.BeginBufferAsync(_token, cancellationToken); + var result = task.IsCompletedSuccessfully + ? task.Result + : await task.ConfigureAwait(false); + Publish(_context.CompleteExtend(_token, result), retained: true); + } } } @@ -191,7 +206,7 @@ void EnsureAdvanced() ThrowHelper.ThrowInvalidOperation("AdvanceTo must be called before reading more message data."); } - void Publish(CurrentSegmentBuffer result, bool retained = false) + void Publish(CurrentMessageBuffer result, bool retained = false) { _buffer = result.Buffer; IsComplete = result.IsComplete; diff --git a/Slon/Pg/Protocol/BackendMessageContext.cs b/Slon/Pg/Protocol/BackendMessageContext.cs index 43fde2c..47e23d5 100644 --- a/Slon/Pg/Protocol/BackendMessageContext.cs +++ b/Slon/Pg/Protocol/BackendMessageContext.cs @@ -1,7 +1,9 @@ using System.Buffers; using System.Diagnostics; +using System.Diagnostics.CodeAnalysis; using System.Runtime.CompilerServices; using Slon.Pipelines; +using Slon.Runtime.CompilerServices; namespace Slon.Pg.Protocol; @@ -9,8 +11,10 @@ namespace Slon.Pg.Protocol; sealed class BackendMessageContext { PgDecoder _decoder = null!; - BackendMessageBatch _remainingBatch; + BackendMessageCursor _cursor; + bool _hasCursor; BackendMessage _current; + FallbackBuffer _currentFallbackBuffer; short _version; const byte PriorCancellationExposure = 1 << 0; const byte BackendTermination = 1 << 1; @@ -19,14 +23,60 @@ sealed class BackendMessageContext const byte MessageOffsetCaptured = 1 << 4; byte _messageState; - // Peek slot: TryPeekNext advances the real batch cursor into here, so the header parse - // happens at peek time and a follow-up TryMoveNext can publish without re-parsing. _hasPeeked - // alone owns validity; leaving the inactive buffer populated avoids a redundant clear and lets - // the next peek usually reuse the same backing objects without write barriers. - bool _hasPeeked; - BackendHeader _peekedHeader; - ReadOnlySequence _peekedBuffer; + enum PublicationState : byte { None, Current, Peeked } + PublicationState _publicationState; long _currentMessageOffset; + ContiguousProjection? _contiguousProjections; + struct FallbackBuffer + { + ReadOnlySequenceSegment? _start; + ReadOnlySequenceSegment? _end; + int _startIndex; + int _endIndex; + + public readonly bool IsEmpty => _start is null; + + public void Set(in ReadOnlySequence buffer) + { + var start = (ReadOnlySequenceSegment)buffer.Start.GetObject()!; + var end = (ReadOnlySequenceSegment)buffer.End.GetObject()!; + Set(start, buffer.Start.GetInteger() & int.MaxValue, + end, buffer.End.GetInteger() & int.MaxValue); + } + + public void Set(ReadOnlySequenceSegment start, int startIndex, + ReadOnlySequenceSegment end, int endIndex) + { + if (!ReferenceEquals(_start, start)) + _start = start; + if (!ReferenceEquals(_end, end)) + _end = end; + _startIndex = startIndex; + _endIndex = endIndex; + } + + public void Clear() + { + if (_start is not null) + _start = null; + if (_end is not null) + _end = null; + _startIndex = 0; + _endIndex = 0; + } + + public readonly ReadOnlySequence Sequence + => new(_start!, _startIndex, _end!, _endIndex); + } + + sealed class ContiguousProjection + { + public required byte[] Buffer { get; init; } + public required SequencePosition Start { get; init; } + public required int Length { get; init; } + public ContiguousProjection? Next { get; init; } + } + public BackendMessage Current { @@ -34,73 +84,262 @@ public BackendMessage Current get { var current = _current; - if (current.IsDefault) + if (_publicationState is not PublicationState.Current) ThrowHelper.ThrowInvalidOperation("The decoder has no current backend message."); return current; } } + public BackendMessage.Accessor CurrentAccessor + { + [MethodImpl(MethodImplOptions.AggressiveInlining)] + get + { + if (_publicationState is not PublicationState.Current) + ThrowHelper.ThrowInvalidOperation("The decoder has no current backend message."); + return new(this, _version, _current.Header.Type, _current.Buffered); + } + } + public bool CurrentIsError { [MethodImpl(MethodImplOptions.AggressiveInlining)] - get => _current.Header.Type is PgTypes.BackendType.ErrorResponse; + get + { + Debug.Assert(_publicationState is PublicationState.Current); + return _current.Header.Type is PgTypes.BackendType.ErrorResponse; + } + } + + public PgTypes.BackendType CurrentType + { + [MethodImpl(MethodImplOptions.AggressiveInlining)] + get + { + Debug.Assert(_publicationState is PublicationState.Current); + return _current.Header.Type; + } + } + + public bool CurrentBuffered + { + [MethodImpl(MethodImplOptions.AggressiveInlining)] + get + { + Debug.Assert(_publicationState is PublicationState.Current); + return _current.Buffered; + } + } + + public ReadOnlyMemory CurrentBufferedBody + { + [MethodImpl(MethodImplOptions.AggressiveInlining)] + get + { + Debug.Assert(_publicationState is PublicationState.Current); + if (_current.TryGetBufferedArrayMemory(0, out var body) + && body.Length == _current.Header.BodyLength) + return body; + return GetCurrentBufferedBodySlow(); + } + } + + [MethodImpl(MethodImplOptions.NoInlining)] + ReadOnlyMemory GetCurrentBufferedBodySlow() + { + if (_current.TryGetBufferedFirstMemory(0, out var body) + && body.Length == _current.Header.BodyLength) + return body; + + var sequence = _current.GetSequence(); + return _current.GetContiguousMemory(sequence); } [MethodImpl(MethodImplOptions.AggressiveInlining)] public bool TryGetCurrent(out BackendMessage current) { current = _current; - return !current.IsDefault; + return _publicationState is PublicationState.Current; } public BackendMessage GetCurrent(short token) { - if (_version != token) + if (_publicationState is not PublicationState.Current || _version != token) ThrowHelper.ThrowInvalidOperation("Backend message has been invalidated by moving to the next message."); return _current; } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal BackendMessageBodyReader OpenCurrentBodyReader(short token) + { + Validate(token); + return _current.OpenBodyReader(); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal bool TryGetCurrentBufferedFirstMemory(short token, int offset, + out ReadOnlyMemory memory) + { + Validate(token); + return _current.TryGetBufferedFirstMemory(offset, out memory); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal bool TryGetCurrentBufferedArray(short token, int offset, + [NotNullWhen(true)] out byte[]? array, + out int arrayOffset, out int length) + { + Validate(token); + return _current.TryGetBufferedArray(offset, out array, out arrayOffset, out length); + } + + internal void SetCurrentFallbackBuffer( + in ReadOnlySequence buffer, bool required) + { + if (required) + _currentFallbackBuffer.Set(in buffer); + else if (!_currentFallbackBuffer.IsEmpty) + _currentFallbackBuffer.Clear(); + } + + internal void SetCurrentFallbackBuffer( + in BackendMessageCursor.FastReadOnlySequence buffer) + { + if (buffer.StartObject is ReadOnlySequenceSegment start) + { + _currentFallbackBuffer.Set(start, buffer.StartIndex, + (ReadOnlySequenceSegment)buffer.EndObject!, buffer.EndIndex); + } + else if (!_currentFallbackBuffer.IsEmpty) + { + _currentFallbackBuffer.Clear(); + } + } + + internal ReadOnlySequence GetFallbackBuffer(short token) + { + Validate(token); + return _currentFallbackBuffer.Sequence; + } + public long GetCurrentMessageOffset(short token) { Validate(token); if ((_messageState & MessageOffsetCaptured) == 0) { - // Fully buffered messages never need their batch-relative offset. Capture it only - // before a streaming operation can replace the segment used to derive it. - Debug.Assert(!_current.Buffered && !_hasPeeked); - _currentMessageOffset = _remainingBatch.ConsumedLength - _current.BufferedLength; + _currentMessageOffset = _cursor.GetCurrentMessageOffset( + _current.BufferedLength); _messageState |= MessageOffsetCaptured; } return _currentMessageOffset; } + internal long CaptureCurrentMessageOffset() + => GetCurrentMessageOffset(_version); + + internal void RebaseCurrentMessageOffset() + { + Debug.Assert((_messageState & MessageOffsetCaptured) != 0); + _currentMessageOffset = 0; + } + + internal bool TryGetCursorConsumedLength(out long consumedLength) + { + if (!_hasCursor) + { + consumedLength = 0; + return false; + } + consumedLength = _cursor.ConsumedLength; + return true; + } + + public ReadOnlyMemory GetContiguousMemory( + short token, ReadOnlyMemory source) + { + Validate(token); + return source; + } + + public ReadOnlyMemory GetContiguousMemory( + short token, in ReadOnlySequence source) + { + Validate(token); + if (source.IsSingleSegment) + return source.First; + if (source.Length > int.MaxValue) + throw new ArgumentOutOfRangeException(nameof(source)); + + var length = (int)source.Length; + for (var projection = _contiguousProjections; + projection is not null; + projection = projection.Next) + { + if (projection.Start.Equals(source.Start) + && projection.Length == length) + return projection.Buffer.AsMemory(0, length); + } + + var buffer = ArrayPool.Shared.Rent(length); + source.CopyTo(buffer); + _contiguousProjections = new() + { + Buffer = buffer, + Start = source.Start, + Length = length, + Next = _contiguousProjections + }; + return buffer.AsMemory(0, length); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void ReleaseContiguousProjections() + { + var projection = _contiguousProjections; + if (projection is null) + return; + ReleaseContiguousProjectionsCore(projection); + } + + [MethodImpl(MethodImplOptions.NoInlining)] + void ReleaseContiguousProjectionsCore(ContiguousProjection projection) + { + _contiguousProjections = null; + do + { + ArrayPool.Shared.Return(projection.Buffer); + projection = projection.Next!; + } + while (projection is not null); + } + public void BindDecoder(PgDecoder decoder) { if (!ReferenceEquals(_decoder, decoder)) _decoder = decoder; } - public bool TryContinue(short token, SequencePosition consumed, long consumedLength, - out CurrentSegmentBuffer result) + public bool TrySlide(short token, SequencePosition consumed, long consumedLength, + out CurrentMessageBuffer result) { MarkBodyWindowAdvanced(token); - return _decoder.TryContinueCurrentMessage(consumed, consumedLength, out result); + return _decoder.TrySlideCurrentMessage(consumed, consumedLength, out result); } - public ValueTask ContinueAsync(short token, SequencePosition consumed, + public ValueTask SlideAsync(short token, SequencePosition consumed, long consumedLength, CancellationToken cancellationToken) { MarkBodyWindowAdvanced(token); - return _decoder.ContinueCurrentMessageAsync(consumed, consumedLength, cancellationToken); + return _decoder.SlideCurrentMessageAsync(consumed, consumedLength, cancellationToken); } - public CurrentSegmentBuffer Continue(short token, SequencePosition consumed, long consumedLength) + public CurrentMessageBuffer Slide(short token, SequencePosition consumed, long consumedLength) { MarkBodyWindowAdvanced(token); - return _decoder.ContinueCurrentMessage(consumed, consumedLength); + return _decoder.SlideCurrentMessage(consumed, consumedLength); } - public bool TryExtend(short token, out CurrentSegmentBuffer result) + public bool TryExtend(short token, out CurrentMessageBuffer result) { EnsureBodyWindowAvailable(token); if (!_decoder.TryExtendCurrentMessage(out result)) @@ -109,19 +348,28 @@ public bool TryExtend(short token, out CurrentSegmentBuffer result) return true; } - public async ValueTask ExtendAsync(short token, CancellationToken cancellationToken) + public ValueTask BeginExtendAsync(short token, CancellationToken cancellationToken) + { + EnsureBodyWindowAvailable(token); + return _decoder.ExtendCurrentMessageAsync(cancellationToken); + } + + public ValueTask BeginBufferAsync(short token, CancellationToken cancellationToken) { EnsureBodyWindowAvailable(token); - return GetBodyBuffer(token, await _decoder.ExtendCurrentMessageAsync(cancellationToken).ConfigureAwait(false)); + return _decoder.BufferCurrentMessageAsync(cancellationToken); } - public CurrentSegmentBuffer Extend(short token) + public CurrentMessageBuffer CompleteExtend(short token, CurrentMessageBuffer result) + => GetBodyBuffer(token, result); + + public CurrentMessageBuffer Extend(short token) { EnsureBodyWindowAvailable(token); return GetBodyBuffer(token, _decoder.ExtendCurrentMessage()); } - CurrentSegmentBuffer GetBodyBuffer(short token, CurrentSegmentBuffer result) + CurrentMessageBuffer GetBodyBuffer(short token, CurrentMessageBuffer result) { Validate(token); var bodyOffset = _currentMessageOffset + BackendHeader.ByteCount; @@ -129,7 +377,15 @@ CurrentSegmentBuffer GetBodyBuffer(short token, CurrentSegmentBuffer result) var bufferedLength = Math.Min(bodyLength, result.Buffer.Length - bodyOffset); var body = result.Buffer.Slice(bodyOffset, bufferedLength); if (result.IsComplete) - SetCurrentFromSegment(token, result.Buffer); + { + _decoder.CompleteCurrentMessage(); + var messageLength = _current.Header.MessageLength; + var message = result.Buffer.Slice(_currentMessageOffset, messageLength); + BackendMessage.Initialize( + ref _current, _current.Header, message, this, token, buffered: true); + var messageEnd = _currentMessageOffset + messageLength; + _cursor = new BackendMessageCursor(result.Buffer).Slice(messageEnd); + } return new(body, result.IsComplete); } @@ -148,16 +404,10 @@ public void EnsureBodyWindowAvailable(short token) _ = GetCurrentMessageOffset(token); } - public void SetBuffered(short token, ReadOnlySequence buffer) - { - Validate(token); - BackendMessage.Initialize(ref _current, _current.Header, buffer, this, token, buffered: true); - } - public void BufferCurrentMessage(short token) { EnsureBodyWindowAvailable(token); - CurrentSegmentBuffer result; + CurrentMessageBuffer result; do result = Extend(token); while (!result.IsComplete); } @@ -169,21 +419,19 @@ public ValueTask BufferCurrentMessageAsync(short token, CancellationToken cancel async ValueTask Core(short token, CancellationToken cancellationToken) { - CurrentSegmentBuffer result; - do result = await ExtendAsync(token, cancellationToken).ConfigureAwait(false); + CurrentMessageBuffer result; + do + { + result = CompleteExtend(token, + await BeginExtendAsync(token, cancellationToken).ConfigureAwait(false)); + } while (!result.IsComplete); } } - void SetCurrentFromSegment(short token, ReadOnlySequence segment) - { - var message = segment.Slice(_currentMessageOffset, _current.Header.MessageLength); - SetBuffered(token, message); - } - void Validate(short token) { - if (_version != token) + if (_publicationState is PublicationState.Peeked || _version != token) ThrowHelper.ThrowInvalidOperation("Backend message has been invalidated by moving to the next message."); } @@ -224,19 +472,24 @@ public bool TryObserveError(short token) public bool TryMoveNext() { - if (_hasPeeked) + if (_publicationState is PublicationState.Peeked) { - _hasPeeked = false; - ResetMessageState(); - BackendMessage.Initialize(ref _current, _peekedHeader, _peekedBuffer, this, ++_version, - _peekedBuffer.Length >= _peekedHeader.MessageLength); + PublishPeeked(); return true; } - if (!_remainingBatch.TryReadNextInPlace(out var header, out var buffer, out var bufferLength)) + var bufferSlot = default(StackValue>); + if (!_cursor.TryReadNextBuffer( + out var header, ref bufferSlot, out var bufferLength)) return false; + var buffer = bufferSlot.Value; ResetMessageState(); + if (bufferLength < header.MessageLength) + _decoder.SetCurrentMessageLength( + _cursor.ConsumedLength - bufferLength + + header.MessageLength); BackendMessage.Initialize(ref _current, header, buffer, this, ++_version, bufferLength >= header.MessageLength); + _publicationState = PublicationState.Current; return true; void ResetMessageState() @@ -245,53 +498,77 @@ void ResetMessageState() } } - public void RetireCurrentBatch() + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public void PublishPeeked() { - // Moving the batch enumerator may return or refill the memory backing every view held here. + Debug.Assert(_publicationState is PublicationState.Peeked); + _publicationState = PublicationState.Current; + } + + public void RetireCursor(bool retainProjections = false) + { + if (!retainProjections) + ReleaseContiguousProjections(); + // Advancing the read grant may return or refill the memory backing every view held here. // A failed message poll preserves Current, but crossing this ownership boundary cannot. - var invalidateToken = !_current.IsDefault || _hasPeeked; + var invalidateToken = _publicationState is not PublicationState.None; _current = default; - _hasPeeked = false; - _peekedHeader = default; - _peekedBuffer = default; - _remainingBatch = default; + _currentFallbackBuffer.Clear(); + _publicationState = PublicationState.None; + _cursor = default; + _hasCursor = false; _currentMessageOffset = 0; _messageState = 0; if (invalidateToken) _version++; } - // Reads the next message WITHOUT publishing it as Current. The remaining batch cursor - // really advances past the header, but the parsed (header, buffer) lands in the peek - // slot and the follow-up TryMoveNext picks it up without re-parsing. The returned - // BackendMessage is valid until the next TryMoveNext (which bumps the version token); - // use it immediately, don't store it. - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public bool TryPeekNextType(out PgTypes.BackendType type) + public bool TryGetReadRequirement( + out long consumedLength, out long requiredLength) { - if (_hasPeeked) + if (!_hasCursor || _cursor.RequiredBufferedLength <= 0) { - type = _peekedHeader.Type; - return true; + consumedLength = 0; + requiredLength = 0; + return false; } - return _remainingBatch.TryPeekType(out type); + + consumedLength = _cursor.ConsumedLength; + requiredLength = _cursor.RequiredBufferedLength + - consumedLength; + return true; } + // Reads the next message WITHOUT publishing it as Current. The message cursor + // really advances past the header, but the parsed (header, buffer) lands in the peek + // slot and the follow-up TryMoveNext picks it up without re-parsing. The returned + // BackendMessage is valid until the next TryMoveNext (which bumps the version token); + // use it immediately, don't store it. + [MethodImpl(MethodImplOptions.NoInlining)] public bool TryPeekNext(out BackendHeader header) { - if (_hasPeeked) + if (_publicationState is PublicationState.Peeked) { - header = _peekedHeader; + header = _current.Header; return true; } - if (!_remainingBatch.TryReadNextInPlace(out _peekedHeader, out var buffer, out _)) + var bufferSlot = default(StackValue>); + if (!_cursor.TryReadNextBuffer( + out header, ref bufferSlot, out var bufferLength)) { - header = default; return false; } - BackendMessage.SetSequence(ref _peekedBuffer, in buffer); - _hasPeeked = true; - header = _peekedHeader; + var buffer = bufferSlot.Value; + var messageLength = header.MessageLength; + var buffered = bufferLength >= messageLength; + if (!buffered) + _decoder.SetCurrentMessageLength( + _cursor.ConsumedLength - bufferLength + + messageLength); + _messageState = 0; + BackendMessage.Initialize(ref _current, header, buffer, this, ++_version, + buffered); + _publicationState = PublicationState.Peeked; return true; } @@ -299,20 +576,19 @@ public BackendMessage Peeked { get { - Debug.Assert(_hasPeeked); - return new(_peekedHeader, _peekedBuffer, this, _version); + Debug.Assert(_publicationState is PublicationState.Peeked); + return _current; } } [MethodImpl(MethodImplOptions.AggressiveInlining)] - public void SetBatch(BackendMessageBatch batch) - { - Debug.Assert(_current.IsDefault && !_hasPeeked, - "The prior batch must be retired before publishing replacement storage."); - // Keep release behavior defensive. The inactive buffer may stay populated because - // _hasPeeked owns validity and the next peek overwrites it. - _hasPeeked = false; - _remainingBatch = batch; + public void SetCursor(BackendMessageCursor cursor) + { + Debug.Assert(_publicationState is PublicationState.None, + "The prior cursor must be retired before publishing replacement storage."); + _publicationState = PublicationState.None; + _cursor = cursor; + _hasCursor = true; } } diff --git a/Slon/Pg/Protocol/BackendMessageCursor.cs b/Slon/Pg/Protocol/BackendMessageCursor.cs new file mode 100644 index 0000000..665dc35 --- /dev/null +++ b/Slon/Pg/Protocol/BackendMessageCursor.cs @@ -0,0 +1,296 @@ +using System.Buffers; +using System.Diagnostics; +using System.Diagnostics.CodeAnalysis; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using Slon.Pipelines; +using Slon.Runtime.CompilerServices; +using static Slon.Pg.Protocol.PgTypes; + +namespace Slon.Pg.Protocol; + +// The cursor is perf sensitive. +struct BackendMessageCursor(ReadOnlySequence buffer) +{ + public const int DefaultDataRowStreamingThreshold = 16 * 1024; + const uint MaxMessageLength = 0x3FFF_FFFF; + + FastReadOnlySequence _buffer = new(buffer); + long _initialLength = buffer.Length; + readonly int _dataRowStreamingThreshold = DefaultDataRowStreamingThreshold; + long _requiredBufferedLength; + + internal BackendMessageCursor( + ReadOnlySequence buffer, int dataRowStreamingThreshold) : this(buffer) + => _dataRowStreamingThreshold = dataRowStreamingThreshold; + + BackendMessageCursor(ReadOnlySequence buffer, + int dataRowStreamingThreshold, long initialLength) + : this(buffer, dataRowStreamingThreshold) + => _initialLength = initialLength; + + public readonly long ConsumedLength => _initialLength - _buffer.Length; + public readonly long RequiredBufferedLength => _requiredBufferedLength; + public readonly SequencePosition UnreadStart => _buffer.Sequence.Start; + + public readonly long GetCurrentMessageOffset(long currentBufferedLength) + => _initialLength - _buffer.Length - currentBufferedLength; + + public readonly BackendMessageCursor Slice(long offset) + { + return new(_buffer.Sequence.Slice(offset), + _dataRowStreamingThreshold, _initialLength); + } + + [MethodImpl(MethodImplOptions.NoInlining)] + public bool TryReadNextInPlace(out BackendHeader header, out ReadOnlySequence buffer, out uint bufferLength) + { + var bufferSlot = default(StackValue>); + if (!TryReadNextBuffer(out header, ref bufferSlot, out bufferLength)) + { + buffer = default; + return false; + } + buffer = bufferSlot.Value.Sequence; + return true; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + internal bool TryReadNextBuffer(out BackendHeader header, + ref StackValue> buffer, out uint bufferLength) + { + var bufferedLength = _buffer.Length; + if (!Header.TryParse(_buffer.FirstSpan, out var protoHeader) + && (bufferedLength < Header.ByteCount + || !Header.TryParseMultiSegment(_buffer.Sequence, out protoHeader))) + { + _requiredBufferedLength = _initialLength - bufferedLength + Header.ByteCount; + bufferLength = default; + header = default; + return false; + } + + var backendType = (BackendType)protoHeader.Tag; + var messageLength = protoHeader.MessageLength; + if (messageLength > MaxMessageLength) + ThrowMessageTooLong(messageLength); + var required = backendType is BackendType.DataRow + ? Math.Min(messageLength, (uint)_dataRowStreamingThreshold) + : messageLength; + if (bufferedLength < required) + { + _requiredBufferedLength = _initialLength - bufferedLength + required; + bufferLength = default; + header = default; + return false; + } + + var result = _buffer.SplitInPlace(Math.Min(bufferedLength, messageLength)); + _requiredBufferedLength = 0; + Debug.Assert(result.Length <= uint.MaxValue); + bufferLength = unchecked((uint)result.Length); + header = BackendHeader.FromHeader(protoHeader); + buffer.Value = result; + return true; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + [DoesNotReturn] + static void ThrowMessageTooLong(uint messageLength) + => throw new PgFramingException( + $"PostgreSQL backend message length {messageLength} exceeds the maximum supported length."); + + public readonly bool TryReadNext(out BackendHeader header, out ReadOnlySequence buffer, out uint bufferLength, out BackendMessageCursor remaining) + { + var thisCopy = this; + var success = thisCopy.TryReadNextInPlace(out header, out buffer, out bufferLength); + remaining = success ? thisCopy : default; + return success; + } + + // Keeps the public position components scalar so consuming the first segment does not repeatedly + // reconstruct and rediscover the backing of a ReadOnlySequence. Materialize one only at API seams. + internal struct FastReadOnlySequence + { + const int SegmentFlag = int.MinValue; + const int IndexMask = int.MaxValue; + + object? _startObject; + object? _endObject; + int _startIndex; + int _endIndex; + long _length; + + FastReadOnlySequence(object? startObject, int startIndex, + object? endObject, int endIndex, long length, bool segmentBacked) + { + Debug.Assert(Unsafe.SizeOf>() is 32); + _startObject = startObject; + _endObject = endObject; + _startIndex = EncodeIndex(startIndex, segmentBacked); + _endIndex = EncodeIndex(endIndex, segmentBacked); + _length = length; + } + + public FastReadOnlySequence(ReadOnlySequence sequence) + { + Debug.Assert(Unsafe.SizeOf>() is 32); + _startObject = sequence.Start.GetObject(); + _endObject = sequence.End.GetObject(); + var segmentBacked = _startObject is ReadOnlySequenceSegment; + _startIndex = EncodeIndex( + sequence.Start.GetInteger() & IndexMask, segmentBacked); + _endIndex = EncodeIndex( + sequence.End.GetInteger() & IndexMask, segmentBacked); + _length = sequence.Length; + } + + static int EncodeIndex(int index, bool segmentBacked) + => segmentBacked ? index | SegmentFlag : index; + + readonly bool IsSegmentBacked => _startIndex < 0; + + public ReadOnlySequence Sequence + { + get + { + if (_startObject is null) + return default; + var startIndex = StartIndex; + var endIndex = EndIndex; + if (IsSegmentBacked) + { + return new((ReadOnlySequenceSegment)_startObject, startIndex, + (ReadOnlySequenceSegment)_endObject!, endIndex); + } + if (_startObject is T[] array) + { + Debug.Assert(ReferenceEquals(_startObject, _endObject)); + return new(array, startIndex, endIndex - startIndex); + } + var manager = (MemoryManager)_startObject; + Debug.Assert(ReferenceEquals(_startObject, _endObject)); + return new(manager.Memory.Slice(startIndex, endIndex - startIndex)); + } + } + public long Length => _length; + public object? StartObject => _startObject; + public object? EndObject => _endObject; + public int StartIndex => _startIndex & IndexMask; + public int EndIndex => _endIndex & IndexMask; + + public ReadOnlySpan FirstSpan + { + get + { + if (_startObject is null) + return default; + var startIndex = StartIndex; + var endIndex = EndIndex; + if (IsSegmentBacked) + { + var memory = ((ReadOnlySequenceSegment)_startObject).Memory; + var end = ReferenceEquals(_startObject, _endObject) + ? endIndex + : memory.Length; + return memory.Span.Slice(startIndex, end - startIndex); + } + if (_startObject is T[] array) + { + Debug.Assert(ReferenceEquals(_startObject, _endObject)); + return array.AsSpan(startIndex, endIndex - startIndex); + } + return ((MemoryManager)_startObject).Memory.Span + .Slice(startIndex, endIndex - startIndex); + } + } + + ReadOnlyMemory FirstMemory + => IsSegmentBacked + ? ((ReadOnlySequenceSegment)_startObject!).Memory + : _startObject switch + { + T[] array => array, + MemoryManager manager => manager.Memory, + _ => throw new UnreachableException() + }; + + // Returns the sequence before the index, stores the sequence after it in place. + [MethodImpl(MethodImplOptions.AggressiveInlining)] + public FastReadOnlySequence SplitInPlace(long offset) + { + var startIndex = StartIndex; + var firstEnd = ReferenceEquals(_startObject, _endObject) + ? EndIndex + : FirstMemory.Length; + var firstLength = firstEnd - startIndex; + if (offset == _length) + { + var exhausted = this; + _startObject = _endObject; + _startIndex = _endIndex; + _length = 0; + return exhausted; + } + if (offset == firstLength + && IsSegmentBacked + && ((ReadOnlySequenceSegment)_startObject!).Next is { } next) + { + var boundaryPrefix = new FastReadOnlySequence( + _startObject, startIndex, _startObject, firstEnd, offset, + segmentBacked: true); + _startObject = next; + _startIndex = SegmentFlag; + _length -= offset; + NormalizeSingleSegmentArray(); + return boundaryPrefix; + } + if ((ulong)offset < (uint)firstLength) + { + var splitIndex = startIndex + (int)offset; + var prev = new FastReadOnlySequence( + _startObject, startIndex, _startObject, splitIndex, offset, + IsSegmentBacked); + _startIndex = EncodeIndex(splitIndex, IsSegmentBacked); + _length -= offset; + return prev; + } + + return SplitSlow(offset); + } + + [MethodImpl(MethodImplOptions.NoInlining)] + FastReadOnlySequence SplitSlow(long offset) + { + var sequence = Sequence; + var prefix = sequence.Slice(0, offset); + var remaining = sequence.Slice(offset); + var result = new FastReadOnlySequence(prefix); + this = new(remaining); + NormalizeSingleSegmentArray(); + return result; + } + + // ReadOnlySequence.Slice retains segment positions after a multi-segment sequence has + // advanced wholly into its final segment. PipeReader-style consumers should not continue + // paying that historical topology: when the remaining segment exposes array memory, carry + // the actual array and absolute indices from this point forward. + [MethodImpl(MethodImplOptions.NoInlining)] + void NormalizeSingleSegmentArray() + { + if (!IsSegmentBacked || !ReferenceEquals(_startObject, _endObject)) + return; + + var memory = ((ReadOnlySequenceSegment)_startObject!).Memory; + if (!MemoryMarshal.TryGetArray(memory, out ArraySegment array)) + return; + + var startIndex = StartIndex; + var endIndex = EndIndex; + _startObject = _endObject = array.Array; + _startIndex = array.Offset + startIndex; + _endIndex = array.Offset + endIndex; + } + + } +} diff --git a/Slon/Pg/Protocol/CancellationCoordinator.cs b/Slon/Pg/Protocol/CancellationCoordinator.cs index 9f7dad1..7568fe9 100644 --- a/Slon/Pg/Protocol/CancellationCoordinator.cs +++ b/Slon/Pg/Protocol/CancellationCoordinator.cs @@ -335,8 +335,14 @@ void StartDispatch(DispatchLease lease) // Sender delegates may execute arbitrary synchronous work before returning their ValueTask. // The queued callback claims physical start immediately before invocation, so neither a // dormant work item nor blocked sender can occupy the coordinator lock or deadline thread. - if (!ThreadPool.UnsafeQueueUserWorkItem(static state => state.Coordinator.InvokeDispatch(state.Lease), - (Coordinator: this, Lease: lease), preferLocal: false)) + if (!Slon.Threading.SchedulingContext.TrySubmitDetached( + static state => + { + var dispatch = ((CancellationCoordinator Coordinator, DispatchLease Lease))state!; + dispatch.Coordinator.InvokeDispatch(dispatch.Lease); + }, + (this, lease), + preferLocal: false)) _ = ObserveDispatchAsync(lease, new(CancelRequestState.NotSent)); } diff --git a/Slon/Pg/Protocol/CommandCompleteMessage.cs b/Slon/Pg/Protocol/CommandCompleteMessage.cs index da150a0..9039344 100644 --- a/Slon/Pg/Protocol/CommandCompleteMessage.cs +++ b/Slon/Pg/Protocol/CommandCompleteMessage.cs @@ -56,13 +56,14 @@ or StatementType.Copy or StatementType.Move or StatementType.Fetch or StatementT internal static CommandCompleteMessage Create(in BackendMessage message) { + var header = message.Header; message.EnsureExpected(PgTypes.BackendType.EmptyQueryResponse, PgTypes.BackendType.CommandComplete); message.EnsureBuffered(); - if (message.Header.Type is PgTypes.BackendType.EmptyQueryResponse) + if (header.Type is PgTypes.BackendType.EmptyQueryResponse) return new(StatementType.Empty, 0, 0); Span scratch = stackalloc byte[64]; - var bodyLength = message.Header.BodyLength; + var bodyLength = header.BodyLength; var bytes = message.TryGetFirstSpan(0, out var first) && first.Length >= bodyLength ? first[..bodyLength] : CopyToScratch(message.GetSequence(), scratch); diff --git a/Slon/Pg/Protocol/CurrentMessageBuffer.cs b/Slon/Pg/Protocol/CurrentMessageBuffer.cs new file mode 100644 index 0000000..8ed975e --- /dev/null +++ b/Slon/Pg/Protocol/CurrentMessageBuffer.cs @@ -0,0 +1,10 @@ +using System.Buffers; + +namespace Slon.Pg.Protocol; + +readonly struct CurrentMessageBuffer( + ReadOnlySequence buffer, bool isComplete) +{ + public ReadOnlySequence Buffer { get; } = buffer; + public bool IsComplete { get; } = isComplete; +} diff --git a/Slon/Pg/Protocol/Flows/CommandExtensions.cs b/Slon/Pg/Protocol/Flows/CommandExtensions.cs index 564294e..7f0f2b4 100644 --- a/Slon/Pg/Protocol/Flows/CommandExtensions.cs +++ b/Slon/Pg/Protocol/Flows/CommandExtensions.cs @@ -12,6 +12,22 @@ public static class CommandExtensions public static bool IsSimple(this in Command command) => command.PreferSimple && command.WithSync && !command.DescribeOnly && !command.Descriptor.IsPrepared && command.Descriptor.ParameterTypes.Count is 0; + [MethodImpl(MethodImplOptions.AggressiveInlining)] + static bool CanWritePreparedExecution(in Command command, in CommandDescriptor descriptor) + => descriptor.IsPrepared && command.Parameters.Count is 0 + && descriptor.ParameterTypes.Count is 0 && command.ResultFormats.Length is 0; + [MethodImpl(MethodImplOptions.AggressiveInlining)] + static ValueTask WritePreparedExecutionAsync( + in Command command, in CommandDescriptor descriptor, PgEncoder encoder, bool appendSync, + CancellationToken cancellationToken) + { + encoder.WritePreparedExecution(descriptor.CommandName, + describe: command.DescribeOnly || descriptor.PreparedRowDescription is null, + execute: !command.DescribeOnly, + syncCount: (command.WithSync ? 1 : 0) + (appendSync ? 1 : 0)); + return encoder.FlushAsync(cancellationToken); + } + // Sync/async pair at the command-list (full composition) level. No *Auto wrapper here. // Callers picking sync vs async make that choice once at this level rather than threading // a mode flag through every encoder helper underneath. Keeping the list loop in this state @@ -19,9 +35,18 @@ public static bool IsSimple(this in Command command) => public static ValueTask WriteCommandsAsync(this CommandList commands, PgEncoder encoder, bool appendSync, CancellationToken cancellationToken = default) { + if (commands.Count is 1) + { + ref readonly var command = ref commands.ItemRef(0); + var descriptor = command.Descriptor; + if (CanWritePreparedExecution(command, descriptor)) + return WritePreparedExecutionAsync( + command, descriptor, encoder, appendSync, cancellationToken); + } + for (var i = 0; i < commands.Count; i++) { - var command = commands[i]; + ref readonly var command = ref commands.ItemRef(i); var descriptor = command.Descriptor; if (!descriptor.IsPrepared || command.Parameters.Count is not 0 || descriptor.ParameterTypes.Count is not 0 || command.ResultFormats.Length is not 0) @@ -533,9 +558,14 @@ public static (PgError?, ParameterTypeList, RowDescription?) ReadPreparationDesc /// If more commands before the next Sync are expected these would be discarded and absent from the message stream. /// In case of an error inside an explicit transaction block all commands until rollback are affected. public static (PgError, TransactionStatus)? Complete(this in Command command, PgDecoder decoder) + => Complete(command.DescribeOnly, command.WithSync, decoder); + + // Completion depends only on whether an Execute was written and whether a Sync follows it, so a + // reader can retain those two facts instead of the command. + internal static (PgError, TransactionStatus)? Complete(bool describeOnly, bool withSync, PgDecoder decoder) { PgError? errorMessage = null; - if (!command.DescribeOnly) + if (!describeOnly) { // https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY // "Therefore, an Execute phase is always terminated by the appearance of exactly one of these messages: @@ -544,7 +574,7 @@ public static (PgError, TransactionStatus)? Complete(this in Command command, Pg PgTypes.BackendType.CommandComplete, PgTypes.BackendType.EmptyQueryResponse, PgTypes.BackendType.PortalSuspended); } - if (!command.WithSync) + if (!withSync) return errorMessage is not null ? (errorMessage, TransactionStatus.Unknown) : null; // Reading the following RFQ may retire the batch which owns the ErrorResponse body. @@ -565,9 +595,12 @@ public static (PgError, TransactionStatus)? Complete(this in Command command, Pg /// If more commands before the next Sync are expected these would be discarded and absent from the message stream. /// In case of an error inside an explicit transaction block all commands until rollback are affected. public static ValueTask<(PgError, TransactionStatus)?> CompleteAsync(this in Command command, PgDecoder decoder) + => CompleteAsync(command.DescribeOnly, command.WithSync, decoder); + + internal static ValueTask<(PgError, TransactionStatus)?> CompleteAsync(bool describeOnly, bool withSync, PgDecoder decoder) { PgError? errorMessage = null; - if (!command.DescribeOnly) + if (!describeOnly) { // https://www.postgresql.org/docs/current/protocol-flow.html#PROTOCOL-FLOW-EXT-QUERY // "Therefore, an Execute phase is always terminated by the appearance of exactly one of these messages: @@ -576,7 +609,7 @@ public static (PgError, TransactionStatus)? Complete(this in Command command, Pg PgTypes.BackendType.CommandComplete, PgTypes.BackendType.EmptyQueryResponse, PgTypes.BackendType.PortalSuspended); } - if (!command.WithSync) + if (!withSync) return errorMessage is not null ? new((errorMessage, TransactionStatus.Unknown)) : new(result: null); // Reading the following RFQ may retire the batch which owns the ErrorResponse body. diff --git a/Slon/Pg/Protocol/Flows/CommandFlow.Enumerator.cs b/Slon/Pg/Protocol/Flows/CommandFlow.Enumerator.cs deleted file mode 100644 index 7e2faf1..0000000 --- a/Slon/Pg/Protocol/Flows/CommandFlow.Enumerator.cs +++ /dev/null @@ -1,424 +0,0 @@ -using System.Collections; -using System.Threading.Tasks.Sources; - -namespace Slon.Pg.Protocol.Flows; - -partial class CommandFlow -{ - Slon.Threading.Tasks.Sources.ManualResetValueTaskSourceCore _enumeratorMoveNextTaskSource; - // Serializes move-next rearming against body termination. Otherwise Reset can replace the generation - // just before terminal completion and strand the consumer (see MoveNextRearm.tla). Never hold it while - // dispatching the gate, which may run the body inline. - Slon.Threading.SpinLock _rearmLock; - CommandResult? _enumeratorCurrent; - bool _enumeratorCompleted; - bool _isResultReady; - // Consumer-thread-only. The first call uses the initial source generation; later calls rearm it. - // Body start is not a substitute because an executor-driven body may finish before first consumption. - bool _consumerAdvanced; - - bool IsEnumerationCompleted => Volatile.Read(ref _enumeratorCompleted); - void PublishEnumerationCompleted() => Volatile.Write(ref _enumeratorCompleted, true); - - ValueTask EnumeratorMoveNextTask => new(this, _enumeratorMoveNextTaskSource.Version); - - bool IValueTaskSource.GetResult(short token) => _enumeratorMoveNextTaskSource.GetResult(token); - ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _enumeratorMoveNextTaskSource.GetStatus(token); - void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) - => _enumeratorMoveNextTaskSource.OnCompleted(continuation, state, token, flags); - - // Consumer completion must dispatch asynchronously because it may run while the pipeline still owns - // the current execution frame. - void CompleteEnumeration() - { - // Drain errors outrank cancellation and clean completion. Preserve every error across a batch. - if (_drainErrors is { Count: > 0 } errors) - { - Exception fault = errors.Count == 1 ? errors[0] : new AggregateException(errors); - _enumeratorMoveNextTaskSource.TrySetException(fault, runContinuationsAsynchronously: true); - } - else if (Volatile.Read(ref _cancellationState) is { DeliverOce: true } cancellation - && !_consumerDisposed) - _enumeratorMoveNextTaskSource.TrySetException( - new OperationCanceledException(cancellation.DeliverToken), runContinuationsAsynchronously: true); - else - _enumeratorMoveNextTaskSource.TrySetResult(false, runContinuationsAsynchronously: true); - // _enumeratorCompleted was set by the caller (SetResult's completed branch) before this runs. - SignalPumpProgress(); - } - - // Cancellation stops waiting; it does not stop the autonomous drain or escape from disposal. - async ValueTask AwaitDrainOnDispose() - { - var cancellationToken = Volatile.Read(ref _cancellationState)?.FlowToken ?? default; - try - { - await WaitForComplete(cancellationToken).ConfigureAwait(false); - } - catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) - { - // Caller cancelled the wait; unwind. The body drains autonomously in the background. - } - catch (PgClientClosedException) - { - // The flow completed via a protocol close; that is a clean terminal for a disposing consumer. - return; - } - // Flow completion is independent of errors accumulated while draining. - if (_drainErrors is { Count: > 0 } errors) - throw errors.Count == 1 ? errors[0] : new AggregateException(errors); - } - - void AwaitDrainOnDisposeSynchronously() - { - var cancellationToken = Volatile.Read(ref _cancellationState)?.FlowToken ?? default; - try - { - WaitForCompleteSynchronously(cancellationToken); - } - catch (OperationCanceledException) when (cancellationToken.IsCancellationRequested) - { - // Caller cancelled the wait; the body continues draining autonomously. - } - catch (PgClientClosedException) - { - return; - } - - if (_drainErrors is { Count: > 0 } errors) - throw errors.Count == 1 ? errors[0] : new AggregateException(errors); - } - - void CompleteEnumerationWithClose(Exception closeException) - { - // A result may already own this generation. Preserve the close in the latch, but only publish - // terminal enumeration state when the close wins the generation; otherwise the consumer must - // observe the result once, rearm, and self-deliver the latched close on its next move. - _callerInteractionCore.SetCloseLatch(closeException); - if (_enumeratorMoveNextTaskSource.TrySetException(closeException, runContinuationsAsynchronously: true)) - PublishEnumerationCompleted(); - SignalPumpProgress(); - } - - // Consumer terminality can precede body terminality under shutdown. Transfer a live body to the - // autonomous driver without re-entering MoveNext: its gate may already be faulted while its handoff - // continuation is pending. - bool TransferLiveBodyToDrain() - { - if (IsBodyTerminated) - return false; - MarkTerminalConsumerGone(); - _callerInteractionCore.ResumeBody(runContinuationsAsynchronously: true); - _callerInteractionCore.WakeBody(); - return true; - } - - // Terminal state outlives task-source generations. Close takes precedence over caller cancellation; - // absent either, a terminal enumeration completes cleanly. - void EnsureEnumerationCompleted() - { - if (_enumeratorMoveNextTaskSource.GetStatus(_enumeratorMoveNextTaskSource.Version) is not ValueTaskSourceStatus.Pending) - return; - if (_callerInteractionCore.CloseException is { } latched) - _enumeratorMoveNextTaskSource.TrySetException(latched, runContinuationsAsynchronously: true); - else - { - var cancellation = Volatile.Read(ref _cancellationState); - var effectiveCancellationToken = EffectiveCancellationToken; - if (cancellation is { } && Volatile.Read(ref cancellation.Requested) - || effectiveCancellationToken.IsCancellationRequested) - _enumeratorMoveNextTaskSource.TrySetException( - new OperationCanceledException(effectiveCancellationToken.IsCancellationRequested - ? effectiveCancellationToken - : cancellation!.DeliverToken), - runContinuationsAsynchronously: true); - // Rearming after a clean terminal still needs to complete the new generation. - else if (IsEnumerationCompleted) - _enumeratorMoveNextTaskSource.TrySetResult(false, runContinuationsAsynchronously: true); - } - } - - public readonly struct Enumerator(CommandFlow flow) : IEnumerator, IAsyncEnumerator - { - // Here so we can pass the cancellation token and enumerate without boxing the struct (which WithCancellation must do). - /// - public Enumerator GetAsyncEnumerator() => this; - - // Dispose always calls MoveNext to confirm the enumerator is done without tracking additional state. - // So this method should be resilient to multiple fetches of *at least* the final result. - /// - public bool MoveNext() - { - if (flow is null) - return false; - - // Queueing only established FIFO position. This is the first point at which the caller - // is ready to take the source pump and drive the synchronous body. - if (!flow._consumerAdvanced && !flow.IsAsyncAtDispatch) - flow.WaitForSyncHandoff(); - - var takeOverAsyncGate = false; - - // Guard-decide-rearm serialized against the body's terminal (see _rearmLock). The using scope - // ends before the WaitForContinuation drive below (which runs the body inline and would - // re-enter this non-reentrant lock); try/finally release keeps it lock-safe on any throw. - using (flow._rearmLock.EnterScope()) - { - // See MoveNextAsync: terminal enumeration state can outlive a completed source generation. - // Ensure the current generation carries that terminal before returning it. - if (flow.IsEnumerationCompleted) - { - flow.EnsureEnumerationCompleted(); - return flow.EnumeratorMoveNextTask.Result; - } - - if (flow.IsAsync) - { - if (flow._enumeratorCurrent is null) - ThrowHelper.ThrowInvalidOperation("No immediate sync/async mixing is allowed, the first MoveNext{Async} call has to match the async argument passed during initialize."); - flow.IsAsync = false; - takeOverAsyncGate = true; - } - - // See MoveNextAsync: rearm only on a non-first call; the first-call source is fresh and the - // body's first delivery lands on it. - if (flow._consumerAdvanced) - flow._enumeratorMoveNextTaskSource.Reset(); - flow._consumerAdvanced = true; - } - // Close-latch self-deliver (sync): under close this call completes the generation it just - // armed, on its own thread. - if (flow._callerInteractionCore.CloseException is { } syncClosed) - { - flow.CompleteEnumerationWithClose(syncClosed); - return flow.EnumeratorMoveNextTask.Result; - } - // The body may already be parked on the async inter-result gate. Once this caller changes - // the flow to synchronous driving, open that gate inline so the body can hand its continuation - // to the rendezvous below. The edge also covers an in-flight body that has not parked yet. - if (takeOverAsyncGate) - { - flow._callerInteractionCore.ResumeBody(runContinuationsAsynchronously: false); - var delivered = flow.EnumeratorMoveNextTask; - if (delivered.IsCompleted) - return delivered.Result; - } - // A progress wake may precede both result completion and publication of the body's next - // handoff continuation (close faults the gate before a resumed body reaches YieldToCaller). - // Keep rendezvousing until either the result owns this turn or there is body work to drive. - while (true) - { - var continuation = flow._callerInteractionCore.WaitForContinuation(); - var task = flow.EnumeratorMoveNextTask; - if (task.IsCompleted) - { - if (continuation is not null) - flow._callerInteractionCore.DeferContinuation(continuation); - return task.Result; - } - continuation ??= flow._callerInteractionCore.TryTakePendingContinuation(); - if (continuation is null) - continue; - continuation.Invoke(); - } - } - - // DisposeAsync always calls MoveNextAsync to confirm the enumerator is done without tracking additional state. - // So this method should be resilient to multiple fetches of the final result. - /// - public ValueTask MoveNextAsync() => MoveNextAsync(default); - - /// Advances the enumerator asynchronously to the next element of the collection. - /// A that may be used to cancel the asynchronous operation. - /// A that will complete with a result of if the enumerator was successfully advanced to the next element, or if the enumerator has passed the end of the collection. - public ValueTask MoveNextAsync(CancellationToken cancellationToken) - { - if (flow is null) - return new(false); - - if (cancellationToken.IsCancellationRequested) - { - flow.GetOrCreateCancellationState().CallerToken = cancellationToken; - if (flow.RequestCancel(cancellationToken, CancellationScope.CurrentWindow)) - { - flow._callerInteractionCore.ResumeBody(runContinuationsAsynchronously: false); - flow._callerInteractionCore.WakeBody(); - } - return ValueTask.FromException(new OperationCanceledException(cancellationToken)); - } - - // The guard-decide-rearm is serialized against the body's terminal (see _rearmLock): the using - // scope covers the _enumeratorCompleted read through the move-next Reset, and ends before the - // gate dispatch below (which runs the body inline and would re-enter this non-reentrant lock). - // try/finally release keeps it lock-safe on any throw. - using (flow._rearmLock.EnterScope()) - { - // Publish the per-read token before the terminal repair below: a flow may already have - // completed before its first consumer call, and EnsureEnumerationCompleted needs this token to - // distinguish a pre-fired cancellation from a clean end. - if (cancellationToken.CanBeCanceled) - { - flow.GetOrCreateCancellationState().CallerToken = cancellationToken; - flow._enumeratorMoveNextTaskSource.CanCompleteConcurrently = true; - } - - // Terminal enumeration state may outlive its completed source generation. Complete the - // newly armed generation with the same close, cancellation, or clean-end outcome. - if (flow.IsEnumerationCompleted) - { - flow.EnsureEnumerationCompleted(); - return flow.EnumeratorMoveNextTask; - } - - if (!flow.IsAsync) - { - if (flow._enumeratorCurrent is null) - ThrowHelper.ThrowInvalidOperation("No immediate sync/async mixing is allowed, the first MoveNext{Async} call has to match the async argument passed during initialize."); - flow.IsAsync = true; - } - - // The first delivery targets the initial generation. After consumer disposal, keep the - // current generation for the body's one-shot terminal. Body-initiated drain retains a live - // consumer and therefore continues rearming. - if (flow._consumerAdvanced && !Volatile.Read(ref flow._consumerDisposed)) - flow._enumeratorMoveNextTaskSource.Reset(); - flow._consumerAdvanced = true; - } - // Drive the body; teardown may already have faulted the gate. - flow._callerInteractionCore.ResumeBody(runContinuationsAsynchronously: false); - // Read the close latch after Reset and complete the generation just armed. - if (flow._callerInteractionCore.CloseException is { } closed) - flow.CompleteEnumerationWithClose(closed); - return flow.EnumeratorMoveNextTask; - } - - public CommandResult Current => flow?._enumeratorCurrent ?? default!; - - /// - public void Dispose() - { - if (flow is null) - return; - - // Consumer terminality can precede body terminality under shutdown. A terminal consumer must - // still transfer a live body to autonomous drain, but must not re-enter the ordinary MoveNext - // pump: its gate may already be faulted while its handoff continuation is pending. - if (flow.IsEnumerationCompleted) - { - FinishCompletedDisposal(flow); - return; - } - - // The final result's terminal message has already been consumed. Only the outer - // enumeration's false publication remains, so finish through the ordinary consumer - // path rather than reclassifying structural terminality as autonomous abandonment. - if (flow.IsFullyConsumedFinalResult) - { - if (MoveNext()) - ThrowHelper.ThrowInvalidOperation( - "A fully consumed physical final result produced another result."); - FinishCompletedDisposal(flow); - return; - } - - // Synchronous disposal takes over an async body through a two-way rendezvous. A gate-parked - // body resumes inline; an in-flight body later hands over its continuation. Waiting on this - // rendezvous rather than MoveNext avoids blocking on the task the body itself must complete. - if (flow.IsAsync) - { - flow.IsAsync = false; - if (flow.WaitForDrainOnDispose) - flow.MarkConsumerWaitForDrain(); - else - flow.MarkConsumerGone(); - // Resume a gate-parked body inline; otherwise buffer the edge. - flow._callerInteractionCore.ResumeBody(runContinuationsAsynchronously: false); - if (flow.WaitForDrainOnDispose) - { - DriveBodyToTermination(flow); - // Drain ran on this thread; surface accumulated drain errors (completes immediately). - flow.AwaitDrainOnDisposeSynchronously(); - } - else - { - // Fast-return: wake the body to drain autonomously in the background, then return. - flow._callerInteractionCore.WakeBody(); - } - return; - } - - // Cancellation graduates a synchronous body to the dedicated driver. The disposer no - // longer pumps the same body concurrently; it only waits for tenure release when requested. - flow.MarkSyncConsumerGone(); - if (flow.WaitForDrainOnDispose) - flow.AwaitDrainOnDisposeSynchronously(); - - static void DriveBodyToTermination(CommandFlow flow) - { - while (!flow.IsBodyTerminated) - { - var continuation = flow._callerInteractionCore.WaitForContinuation(); - if (continuation is null) - { - if (flow.IsBodyTerminated) - break; - continuation = flow._callerInteractionCore.TryTakePendingContinuation(); - if (continuation is null) - continue; - } - continuation.Invoke(); - } - } - - static void FinishCompletedDisposal(CommandFlow flow) - { - if (flow.TransferLiveBodyToDrain() && flow.WaitForDrainOnDispose) - flow.AwaitDrainOnDisposeSynchronously(); - } - } - - /// - public ValueTask DisposeAsync() - { - if (flow is null) - return new(); - - if (flow.IsEnumerationCompleted) - return FinishCompletedDisposalAsync(flow); - - if (flow.IsFullyConsumedFinalResult) - return FinishFinalResultAsync(this, flow); - - // Mark autonomous drain, open the async result gate, and wake any synchronous rendezvous. - // Pipeline tenure keeps successors behind the body until it reaches RFQ. - if (flow.WaitForDrainOnDispose) flow.MarkConsumerWaitForDrain(); else flow.MarkConsumerGone(); - flow._callerInteractionCore.ResumeBody(runContinuationsAsynchronously: false); - flow._callerInteractionCore.WakeBody(); - // Optionally await the body's bounded completion; cancellation stops waiting, not draining. - if (flow.WaitForDrainOnDispose) - return flow.AwaitDrainOnDispose(); - return new(); - - static async ValueTask FinishFinalResultAsync(Enumerator enumerator, CommandFlow flow) - { - if (await enumerator.MoveNextAsync().ConfigureAwait(false)) - ThrowHelper.ThrowInvalidOperation( - "A fully consumed physical final result produced another result."); - await FinishCompletedDisposalAsync(flow).ConfigureAwait(false); - } - - static ValueTask FinishCompletedDisposalAsync(CommandFlow flow) - { - // A completed awaited drain may still have errors to surface. - flow.TransferLiveBodyToDrain(); - return flow.WaitForDrainOnDispose ? flow.AwaitDrainOnDispose() : new(); - } - } - - /// - void IEnumerator.Reset() => throw new NotSupportedException(); - /// - object? IEnumerator.Current => Current; - } - -} diff --git a/Slon/Pg/Protocol/Flows/CommandFlow.MessageEnumerator.cs b/Slon/Pg/Protocol/Flows/CommandFlow.MessageEnumerator.cs index 971fa22..3f6fdfa 100644 --- a/Slon/Pg/Protocol/Flows/CommandFlow.MessageEnumerator.cs +++ b/Slon/Pg/Protocol/Flows/CommandFlow.MessageEnumerator.cs @@ -6,7 +6,7 @@ namespace Slon.Pg.Protocol.Flows; -partial class CommandFlow +public sealed partial class CommandFlow { internal enum MoveNextStatus : byte { @@ -19,14 +19,12 @@ internal struct ReadState { public ResultMessageEnumerator ResultMessageEnumerator { get; } public CommandResult CommandResult { get; } - public ValueTaskSourcePromise ReadPromise { get; } public RowDescription RowDescription { get; } public ReadState() { ResultMessageEnumerator = new(); CommandResult = new(ResultMessageEnumerator); - ReadPromise = new(); RowDescription = new(); } @@ -38,6 +36,11 @@ public void Reset() } } + internal readonly struct ReadPromiseState() + { + public ValueTaskSourcePromise Promise { get; } = new(); + } + // The value wrapper lets ReadState and CommandResult share one MessageEnumerator instance // without an interface or another adapter allocation. internal readonly struct ResultMessageEnumerator() : IEnumerator, IAsyncEnumerator @@ -46,7 +49,16 @@ internal readonly struct ResultMessageEnumerator() : IEnumerator public bool MoveNext() => _messageEnumerator.MoveNext(); public ValueTask MoveNextAsync() => _messageEnumerator.MoveNextAsync(); public BackendMessage Current => _messageEnumerator.Current; + internal BackendMessage.Accessor CurrentAccessor => _messageEnumerator.CurrentAccessor; internal MoveNextStatus TryMoveNext() => _messageEnumerator.TryMoveNext(); + // RowEnumerator classifies the publication itself, avoiding a second Current copy here. + internal MoveNextStatus TryMoveNextRow() => _messageEnumerator.TryMoveNextRow(); + internal void MarkCurrentTerminal() => _messageEnumerator.MarkCurrentTerminal(); + internal ValueTask CollectRowsAsync( + TState state, Action collector, + CancellationToken cancellationToken) + => _messageEnumerator.CollectRowsAsync(state, collector, cancellationToken); + internal void ThrowCollectorException() => _messageEnumerator.ThrowCollectorException(); public void Dispose() => _messageEnumerator.Dispose(); public ValueTask DisposeAsync() => _messageEnumerator.DisposeAsync(); @@ -56,26 +68,31 @@ internal readonly struct ResultMessageEnumerator() : IEnumerator BackendMessage IEnumerator.Current => _messageEnumerator.Current; object? IEnumerator.Current => ((IEnumerator)_messageEnumerator).Current; - public void Initialize(CommandFlow flow, PgDecoder decoder) - => _messageEnumerator.Initialize(flow, decoder); + public void Initialize(in Command command, PgDecoder decoder) + => _messageEnumerator.Initialize(command, decoder); public void Reset() => _messageEnumerator.Reset(); + public void EnableResultBuffering() + => _messageEnumerator.EnableResultBuffering(); + public (PgError Error, TransactionStatus TransactionStatus)? CompleteError => _messageEnumerator.CompleteError; sealed class MessageEnumerator : IEnumerator, IAsyncEnumerator { - CommandFlow _flow = null!; + // Completion needs only these two command facts. Retaining them keeps the protocol-static + // enumerator independent of the flow and avoids holding a reference-bearing command copy. + bool _describeOnly; + bool _withSync; PgDecoder _decoder = null!; bool _disposed; bool _first; bool _done; ExceptionDispatchInfo? _exceptionDispatchInfo; + ExceptionDispatchInfo? _collectorException; (PgError, TransactionStatus)? _completeError; - Command Command => _flow._commands[_flow._commandIndex]; - // An Execute response consists of DataRow messages followed by one terminal message. [Conditional("DEBUG")] static void DebugEnsureExpected(BackendMessage message) @@ -135,6 +152,8 @@ public ValueTask MoveNextAsync() return Core(); + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] async ValueTask Core() { try @@ -153,6 +172,144 @@ async ValueTask Core() } } + enum CollectRowsStatus : byte + { + RequiresInput, + RequiresBuffer, + Complete + } + + public ValueTask CollectRowsAsync( + TState state, Action collector, + CancellationToken cancellationToken) + { + try + { + var status = CollectAvailableRows( + state, collector, currentReady: false, out var pending, out var terminal); + return status is CollectRowsStatus.Complete + ? new(terminal) + : Core(this, state, collector, status, pending, cancellationToken); + } + catch (Exception ex) + { + _exceptionDispatchInfo = ExceptionDispatchInfo.Capture(ex); + return ValueTask.FromException(ex); + } + + static async ValueTask Core( + MessageEnumerator enumerator, + TState state, Action collector, + CollectRowsStatus status, BackendMessage.Accessor pending, + CancellationToken cancellationToken) + { + try + { + while (true) + { + if (status is CollectRowsStatus.RequiresInput) + { + cancellationToken.ThrowIfCancellationRequested(); + _ = await enumerator._decoder.GetNextAsync().ConfigureAwait(false); + } + else + { + await pending.BufferBodyAsync(cancellationToken).ConfigureAwait(false); + } + + status = enumerator.CollectAvailableRows( + state, collector, currentReady: true, out pending, out var terminal); + if (status is CollectRowsStatus.Complete) + return terminal; + } + } + catch (Exception ex) + { + enumerator._exceptionDispatchInfo = ExceptionDispatchInfo.Capture(ex); + throw; + } + } + } + + CollectRowsStatus CollectAvailableRows( + TState state, Action collector, + bool currentReady, + out BackendMessage.Accessor pending, + out BackendMessage terminal) + { + var status = CollectAvailableRowsCore(state, collector, currentReady); + if (status is CollectRowsStatus.RequiresInput) + { + pending = default; + terminal = default; + } + else if (status is CollectRowsStatus.RequiresBuffer) + { + pending = _decoder.CurrentAccessor; + terminal = default; + } + else + { + _done = true; + pending = default; + terminal = _decoder.Current; + } + return status; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + CollectRowsStatus CollectAvailableRowsCore( + TState state, Action collector, + bool currentReady) + { + var decoder = _decoder; + var collect = _collectorException is null; + if (currentReady) + goto ProcessCurrent; + if (_first) + { + _first = false; + goto ProcessCurrent; + } + + _exceptionDispatchInfo?.Throw(); + if (_done) + ThrowHelper.ThrowInvalidOperation( + "Underlying message enumerator completed before a terminal message was returned."); + + MoveNext: + if (!decoder.TryMoveNext()) + return CollectRowsStatus.RequiresInput; + + ProcessCurrent: + DebugEnsureExpected(decoder.Current); + if (decoder.CurrentType is not PgTypes.BackendType.DataRow) + return CollectRowsStatus.Complete; + if (!decoder.CurrentBuffered) + return CollectRowsStatus.RequiresBuffer; + + if (collect) + { + try + { + collector(state, new CommandResult.RowView(decoder.CurrentBufferedBody)); + } + catch (Exception ex) + { + _collectorException = ExceptionDispatchInfo.Capture(ex); + collect = false; + } + } + goto MoveNext; + } + + public void ThrowCollectorException() + { + var exception = _collectorException; + _collectorException = null; + exception?.Throw(); + } + public MoveNextStatus TryMoveNext() { if (_first) @@ -177,12 +334,32 @@ public MoveNextStatus TryMoveNext() return MoveNextStatus.RequiresInput; } + public MoveNextStatus TryMoveNextRow() + { + if (_first) + { + _first = false; + return MoveNextStatus.Moved; + } + + _exceptionDispatchInfo?.Throw(); + if (_done) + return MoveNextStatus.EndOfSequence; + + return _decoder.TryMoveNext() + ? MoveNextStatus.Moved + : MoveNextStatus.RequiresInput; + } + + public void MarkCurrentTerminal() => _done = true; public BackendMessage Current { [MethodImpl(MethodImplOptions.AggressiveInlining)] get => _decoder.Current; } + internal BackendMessage.Accessor CurrentAccessor => _decoder.CurrentAccessor; + public void Dispose() { _exceptionDispatchInfo?.Throw(); @@ -197,7 +374,7 @@ public void Dispose() { while (decoder.GetNext().Header.Type is PgTypes.BackendType.DataRow) {} } - _completeError = Command.Complete(_decoder); + _completeError = CommandExtensions.Complete(_describeOnly, _withSync, _decoder); } catch (Exception ex) { @@ -224,7 +401,7 @@ ValueTask DisposeAsyncCore() if (decoder.TryGetCurrent(out var current) && current.Header.Type is not PgTypes.BackendType.DataRow) { - var completion = Command.CompleteAsync(decoder); + var completion = CommandExtensions.CompleteAsync(_describeOnly, _withSync, decoder); if (completion.IsCompletedSuccessfully) { _completeError = completion.Result; @@ -270,7 +447,7 @@ async ValueTask DrainRowsAndComplete(PgDecoder decoder) if (message.Header.Type is not PgTypes.BackendType.DataRow) break; } - _completeError = await Command.CompleteAsync(decoder).ConfigureAwait(false); + _completeError = await CommandExtensions.CompleteAsync(_describeOnly, _withSync, decoder).ConfigureAwait(false); } catch (Exception ex) { @@ -280,33 +457,50 @@ async ValueTask DrainRowsAndComplete(PgDecoder decoder) } } - public void Initialize(CommandFlow flow, PgDecoder decoder) + public void Initialize(in Command command, PgDecoder decoder) { - if (!ReferenceEquals(_flow, flow)) - _flow = flow; + if (_decoder is not null) + _decoder.ResultBuffering = false; + _describeOnly = command.DescribeOnly; + _withSync = command.WithSync; if (!ReferenceEquals(_decoder, decoder)) _decoder = decoder; _exceptionDispatchInfo = null; + if (_collectorException is not null) + _collectorException = null; _disposed = false; _completeError = null; // A command is immediately done if we haven't submitted an execute. - _done = Command.DescribeOnly; + _done = _describeOnly; _first = !_done; } public void Reset() { - _flow = null!; + if (_decoder is not null) + _decoder.ResultBuffering = false; + _describeOnly = false; + _withSync = false; _decoder = null!; _exceptionDispatchInfo = null; + if (_collectorException is not null) + _collectorException = null; _completeError = null; _disposed = true; _first = false; _done = true; } + public void EnableResultBuffering() + { + if (_disposed) + ThrowHelper.ThrowInvalidOperation( + "The command result has already been released."); + _decoder.ResultBuffering = true; + } + public (PgError Error, TransactionStatus TransactionStatus)? CompleteError { get diff --git a/Slon/Pg/Protocol/Flows/CommandFlow.cs b/Slon/Pg/Protocol/Flows/CommandFlow.cs index 66a4870..cbcf656 100644 --- a/Slon/Pg/Protocol/Flows/CommandFlow.cs +++ b/Slon/Pg/Protocol/Flows/CommandFlow.cs @@ -1,9 +1,10 @@ +using System.Collections; using System.Diagnostics; using System.Runtime.CompilerServices; +using System.Runtime.ExceptionServices; +using System.Runtime.InteropServices; using System.Threading.Tasks.Sources; using Slon.Runtime.CompilerServices; -// A unique result type distinguishes the caller gate from this flow's other IValueTaskSource instantiations. -using FlowCallerInteractionCoreResult = System.ValueTuple; namespace Slon.Pg.Protocol.Flows; @@ -11,7 +12,8 @@ namespace Slon.Pg.Protocol.Flows; public abstract class CommandFlowObserver : PgClientFlowObserver { protected internal virtual void OnStarted(CommandFlow flow, object? state) { } - protected internal virtual void OnCommandResult(CommandFlow flow, CommandResult result, object? state) { } + protected internal virtual void OnCommandResult( + CommandFlow flow, CommandResult result, object? state) { } protected internal virtual void OnDrainStarted(CommandFlow flow, object? state) { } } @@ -25,1187 +27,1635 @@ public readonly struct CommandFlowOptions public TimeSpan? PendingTimeout { get; init; } } -[Experimental(ExperimentalDiagnostics.PostgreSqlLowerLayer)] -public partial class CommandFlow : PgClientFlow, IValueTaskSource, IValueTaskSource, IValueTaskSource +internal sealed class CommandExecutionColdState { - static readonly TimeSpan ConsumerDrainCancellationGracePeriod = TimeSpan.FromSeconds(1); + internal bool CancelRequested; + internal CancellationToken CallerToken; + internal CancellationTokenRegistration CallerRegistration; + internal CancellationToken DeliverToken; + internal int Scope; + internal int Timing; + internal int SubsequentTiming; + internal TaskCompletionSource? Delivery; + internal object? EpisodeKey; + internal Exception? CloseException; + // Replayed by later consumer calls once the flow reached its terminal. + internal Exception? TerminalException; + // Command errors observed while draining without a consumer. Multiple Sync windows may each + // produce an independent ErrorResponse, all of which belong to the waiting disposer. + internal List? DrainErrors; +} - internal override bool DefersSyncHandoff => true; +internal enum CommandExecutionCancellationScope : byte +{ + CurrentWindow = 1, + RemainingFlow = 2 +} - internal enum CancellationScope : byte - { - None, - CurrentWindow = 1, - RemainingFlow = 2 - } - - // Flow state - CommandList _commands; - TimeSpan? _pendingTimeout; - FlowCallerInteractionCore _callerInteractionCore; - // Cancellation is cold; keep its tokens, registrations and attribution state off ordinary flows. - CancellationState? _cancellationState; - // Errors encountered while draining are surfaced by a waiting DisposeAsync. Live consumers observe - // their errors directly, and the list remains unallocated on the successful path. - List? _drainErrors; - // Set only by ConsumeNonQueryAsync, after enqueue but before any consumer-side gate release. - // The body cannot reach first publication until such a release, and every release publishes - // this write, so plain accesses suffice and the body observes the mode at first wake. - bool _consumeNonQuery; - bool IsConsumingNonQuery => _consumeNonQuery; - bool IsConsumingAutonomously => IsDraining || IsConsumingNonQuery; - long _nonQueryRecordsAffected; - - // Once draining, the body bypasses result handoffs and reads autonomously to RFQ. This is state, not - // an I/O cancellation token: canceling the I/O would prevent restoration of a clean wire boundary. - bool _draining; - internal bool IsDraining => Volatile.Read(ref _draining); - // Body-thread-only guard: later commands must not change the drive mode chosen on drain entry. - bool _drainModeEntered; - - - // Distinguishes consumer disposal from a body-initiated drain; disposal suppresses terminal OCE delivery. - bool _consumerDisposed; - // When true, DisposeAsync awaits the body's drain to RFQ. Otherwise it returns while the body drains; - // pipeline retirement still prevents the next flow from observing a dirty wire. - internal bool WaitForDrainOnDispose { get; set; } = true; - - // Consumer disposal without waiting for the autonomous drain. - void MarkConsumerGone() - { - Volatile.Write(ref _consumerDisposed, true); - RequestCancel(default, CancellationScope.RemainingFlow, BackendCancellationTiming.AfterGrace, - BackendCancellationTiming.AtReadFrontier); - } +// Mutable execution state is stored inline by each concrete host. Core algorithms re-enter this +// field through their host-specific ops value after every await; they never mutate a copied struct. +[StructLayout(LayoutKind.Auto)] +internal struct CommandExecutionState +{ + internal int Phase; + internal CommandList Commands; + internal TimeSpan? PendingTimeout; + internal int CommandIndex; + internal PgClientFlow.Context Context; + internal bool ContextPublished; + internal CommandResult? Current; + internal bool CurrentPublished; + internal bool ReadFlowRfq; + internal bool ConsumerDetached; + internal bool ConsumerObservedCompletion; + internal Slon.Threading.Tasks.Sources.ManualResetValueTaskSourceCore ReadySource; + internal int ReadyCompletion; + internal Slon.Threading.Tasks.Sources.ManualResetValueTaskSourceCore PipelineTaskSource; + internal CancellationToken FlowToken; + internal CancellationTokenRegistration FlowRegistration; + // Scratch owned by the single active result consumer. Keeping it with the flow lets the + // resumable frames carry only their control state across suspension. + internal CancellationToken WindowToken; + // Every pre-consumed successor can suspend in FirstAsync simultaneously. Retain that frame with + // its reusable flow instead of competing for the builder's one-thread/one-core cache slots. + internal ValueTaskSourcePromise? FirstPromise; + internal CommandExecutionColdState? ColdState; + internal FlowHandoffEvent? HandoffEvent; + internal bool SyncHandoffClaimed; + internal int DrainStarted; + internal bool EnableActivationTimeout; + internal bool WaitForDrainOnDispose; +} + +/// Executes an ordered command list with consumer-owned synchronous or asynchronous result decoding. +[Experimental(ExperimentalDiagnostics.PostgreSqlLowerLayer)] +public sealed partial class CommandFlow : PgClientFlow, IValueTaskSource, IValueTaskSource +{ + static readonly TimeSpan ConsumerDrainCancellationGracePeriod = TimeSpan.FromSeconds(1); + CommandExecutionState _state; + CommandFlowObserver? _commandObserver; + object? _commandObserverState; - // Consumer disposal while waiting for the autonomous drain. - void MarkConsumerWaitForDrain() + CommandFlow(bool async, TimeSpan? pendingTimeout = null) + : base(supportsDeferredFlush: true) { - Volatile.Write(ref _consumerDisposed, true); - RequestCancel(default, CancellationScope.RemainingFlow, BackendCancellationTiming.AfterGrace, - BackendCancellationTiming.AtReadFrontier); + _state.CommandIndex = -1; + _state.EnableActivationTimeout = true; + _state.WaitForDrainOnDispose = true; + _state.PendingTimeout = pendingTimeout; + IsAsync = async; + if (!async) + _state.HandoffEvent = new(false); } - // A synchronous consumer cannot abandon its drive obligation while the body is parked. Give the - // cancellation path a delivery source so WakeBody transfers that obligation to the dedicated driver. - void MarkSyncConsumerGone() + public CommandFlow(bool async, params ReadOnlySpan commands) + : this(async) + => Initialize(async, commands); + + internal CommandFlow( + bool async, bool enableActivationTimeout, params ReadOnlySpan commands) + : this(async, commands) + => _state.EnableActivationTimeout = enableActivationTimeout; + + internal CommandFlow(bool async, CommandList commands, TimeSpan? pendingTimeout = null) + : this(async, pendingTimeout) + => Initialize(async, new CommandFlowOptions + { + Commands = commands, + PendingTimeout = pendingTimeout + }); + + public CommandFlow(bool async, in CommandFlowOptions options) + : this(async, options.PendingTimeout) + => Initialize(async, options); + + public CommandFlow Initialize(bool async, params ReadOnlySpan commands) + => Initialize(async, new CommandFlowOptions { Commands = new(commands) }); + + public CommandFlow Initialize(bool async, in CommandFlowOptions options) { - Volatile.Write(ref _consumerDisposed, true); - _ = GetOrCreateCancelDelivery(); - RequestCancel(default, CancellationScope.RemainingFlow, BackendCancellationTiming.AfterGrace, - BackendCancellationTiming.AtReadFrontier); - _callerInteractionCore.ResumeBody(runContinuationsAsynchronously: true); - _callerInteractionCore.WakeBody(useDedicatedDriver: true); - } - - // Enumeration already ended; only transfer the live body tail to autonomous ownership. - // There is no unfinished consumer work to justify a backend cancellation intent. - void MarkTerminalConsumerGone() - { - Volatile.Write(ref _consumerDisposed, true); - Volatile.Write(ref _draining, true); - } - - // Result publication orders these body-owned fields before the consumer can observe the - // CommandResult. IsComplete is consumer-owned after that handoff. A behavior-limited reader may - // consider its visible result final while later commands still exist, so use the physical - // command index rather than an ADO-visible result count. - bool IsFullyConsumedFinalResult - => _isResultReady && _commandIndex >= CommandCount - 1 - && _enumeratorCurrent is { IsComplete: true }; - - // A body-initiated drain keeps the consumer attached for terminal cancellation or close delivery. - void MarkBodyInitiatedDrain() => Volatile.Write(ref _draining, true); - - // Result-production state - RowDescription? _requestedRowDescription; - PgError? _pgError; - int _commandIndex = -1; - PgDecoder? _decoder; - bool _readFlowRfq; - - // Pipelined dispatch state. Lives here (not on PgClientFlow base) because the shared-promise - // optimization that needs these fields is CommandFlow-specific (see DispatchPipelinedRead). - Slon.Threading.Tasks.Sources.ManualResetValueTaskSourceCore _executePipelinedCore; - ValueTaskSourcePromise? _pipelinePromise; - Context _context; - bool _contextPublished; - ValueTask _task; - // Consumer terminality, body terminality, and framework release are distinct phases. Start and - // pre-start termination race during shutdown, so one atomic state owns that decision. - const int BodyNotStarted = 0; - const int BodyRunning = 1; - const int BodyTerminated = 2; - int _bodyState; - - sealed class CancellationState - { - internal CancellationToken CallerToken; - internal CancellationToken FlowToken; - internal CancellationTokenRegistration CallerRegistration; - internal CancellationTokenRegistration FlowRegistration; - internal bool Requested; - internal int Scope; - internal int Timing; - internal int SubsequentTiming; - internal CancellationToken DeliverToken; - internal TaskCompletionSource? Delivery; - internal object? EpisodeKey; - internal bool DeliverOce; - - internal void Reset() + IsAsync = async; + if (!async) + _state.HandoffEvent ??= new(false); + var commands = options.Commands; + if (commands.Count is 0) + return this; + _state.Commands = commands; + _state.PendingTimeout = options.PendingTimeout; + _commandObserver = options.Observer; + _commandObserverState = options.ObserverState; + if (options.Observer is { } observer) { - CallerToken = default; - FlowToken = default; - CallerRegistration.Dispose(); - CallerRegistration = default; - FlowRegistration.Dispose(); - FlowRegistration = default; - Requested = false; - Scope = (int)CancellationScope.None; - Timing = (int)BackendCancellationTiming.AfterGrace; - SubsequentTiming = (int)BackendCancellationTiming.AfterGrace; - DeliverToken = default; - Delivery = null; - EpisodeKey = null; - DeliverOce = false; + SetObserver(observer, options.ObserverState); + observer.OnStarted(this, options.ObserverState); } - } - CommandFlow() : base(supportsDeferredFlush: true) - { - _callerInteractionCore.Initialize(); + return this; } - // Interactive commands carry caller patience, so arm the activation timeout. - protected override bool EnableActivationTimeout => true; - protected override TimeSpan? PendingTimeout => _pendingTimeout; + internal override bool DefersSyncHandoff => true; + private protected override FlowHandoffEvent? HandoffEvent => _state.HandoffEvent; + protected override bool EnableActivationTimeout => _state.EnableActivationTimeout; + protected override TimeSpan? PendingTimeout => _state.PendingTimeout; internal override TimeSpan? BackendCancellationGracePeriod - => Volatile.Read(ref _consumerDisposed) ? ConsumerDrainCancellationGracePeriod : null; + => Volatile.Read(ref _state.ConsumerDetached) + ? ConsumerDrainCancellationGracePeriod + : null; - public CommandFlow(bool async, params ReadOnlySpan commands) : this() - => Initialize(async, commands); - public CommandFlow(bool async, in CommandFlowOptions options) : this() - => Initialize(async, options); + internal override void BindCallerToken(CancellationToken cancellationToken) + => _state.FlowToken = cancellationToken; + internal override CancellationToken MigrationCancellationToken + => _state.FlowToken; + + public Enumerator GetEnumerator() + => new(this, default); - private protected CommandFlow(bool async, TimeSpan? pendingTimeout) : this() + public Enumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) { - IsAsync = async; - _pendingTimeout = pendingTimeout; - _enumeratorMoveNextTaskSource.CanCompleteConcurrently = true; + // A missing enumeration token must not erase the token captured when the flow was queued. + if (cancellationToken.CanBeCanceled) + _state.FlowToken = cancellationToken; + return new(this, cancellationToken); } - public CommandFlow Initialize(bool async, params ReadOnlySpan commands) - => Initialize(async, options: new() { Commands = new(commands) }); + internal CommandResult? CurrentResult => _state.Current; + public int CommandCount => _state.Commands.Count; + public bool IsResultReady => Core.IsResultReady; + internal int VisibleCommandCount => _state.Commands.VisibleCount; + internal ValueTask MoveNextResultAsync(CancellationToken cancellationToken) + => Core.MoveNextAsync(cancellationToken); + internal void DisposeResults() => Core.Dispose(); + internal ValueTask DisposeResultsAsync() => Core.DisposeAsync(); + internal bool WaitForDrainOnDispose + { + get => _state.WaitForDrainOnDispose; + set => _state.WaitForDrainOnDispose = value; + } - public CommandFlow Initialize(bool async, in CommandFlowOptions options) + CommandFlowCore Core => new(new(this)); + + internal ValueTask ConsumeNonQueryAsync(CancellationToken cancellationToken = default) + => Core.ConsumeNonQueryAsync(cancellationToken); + + protected override ValueTask ExecuteAuto(Context context) + => Core.ExecuteAuto(context); + + internal Task CancelAsync() => Core.CancelAsync(); + + internal override bool ResetsSharedReadStateBeforeRelease => true; + protected override void OnStopping(Exception exception) => Core.OnStopping(exception); + protected override void OnAbort(Exception exception) => Core.OnAbort(exception); + internal override void Fail(Exception exception) => Core.Fail(exception); + protected override void OnReleasing(Exception? exception) => Core.OnReleasing(exception); + protected override void OnDiscarded() => Core.OnDiscarded(); + protected override void OnReset() => Core.OnReset(); + + readonly struct Ops(CommandFlow owner) : ICommandExecutionFlowOps { - IsAsync = async; - if (options.Observer is { } observer) - SetObserver(observer, options.ObserverState); - _commands = options.Commands; - _pendingTimeout = options.PendingTimeout; - options.Observer?.OnStarted(this, options.ObserverState); - // Arm before publication: teardown may complete the source concurrently even before enumeration. - _enumeratorMoveNextTaskSource.CanCompleteConcurrently = true; - return this; + readonly CommandFlow _owner = owner; + + public static Ops Create(PgClientFlow flow) => new((CommandFlow)flow); + public PgClientFlow Flow => _owner; + public ref CommandExecutionState GetField() => ref _owner._state; + public bool IsAsync + { + get => _owner.IsAsync; + set => _owner.IsAsync = value; + } + public bool IsAsyncAtDispatch => _owner.IsAsyncAtDispatch; + public bool HasSuccessfulActivation => _owner.HasSuccessfulActivation; + public void WaitForSyncHandoff() => _owner.WaitForSyncHandoff(); + public void OnCommandResult(CommandResult result) + => _owner._commandObserver?.OnCommandResult( + _owner, result, _owner._commandObserverState); + public void OnDrainStarted() + => _owner._commandObserver?.OnDrainStarted( + _owner, _owner._commandObserverState); + public void OnDiscarded() + => _owner.GetObserver(out var observerState)?.OnCompleting( + _owner, null, observerState); } +} + +internal interface ICommandExecutionFlowOps : IFieldRef + where TSelf : struct, ICommandExecutionFlowOps +{ + static abstract TSelf Create(PgClientFlow flow); + PgClientFlow Flow { get; } + bool IsAsync { get; set; } + bool IsAsyncAtDispatch { get; } + bool HasSuccessfulActivation { get; } + void WaitForSyncHandoff(); + void OnCommandResult(CommandResult result); + void OnDrainStarted(); + void OnDiscarded(); +} + +readonly struct CommandFlowCore(TOps ops) + where TOps : struct, ICommandExecutionFlowOps +{ + const int PhaseInitial = 0; + const int PhaseReading = 1; + const int PhaseResultReady = 2; + const int PhaseDraining = 3; + const int PhaseCompleted = 4; + + readonly TOps _ops = ops; + ref CommandExecutionState _state => ref _ops.GetField(); + [MethodImpl(MethodImplOptions.AggressiveInlining)] + void SetCurrent(CommandResult result) => _state.Current = result; + internal bool IsResultReady => Volatile.Read(ref _state.Phase) is PhaseResultReady; + bool IsSinglePublishedCommand + => _state.Commands.Count is 1 + && !_state.Commands.ItemRef(0).SuppressEnumeration; - // Declares internal non-query consumption and runs it to completion. The single entry point - // makes the ownership rule structural: no enumerator is exposed on this path, so mixing - // enumeration with internal consumption is unrepresentable. The declaration precedes any - // consumer-side gate release, so the body observes it at first wake and never publishes. + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] - internal async ValueTask ConsumeNonQueryAsync(CancellationToken cancellationToken = default) + internal async ValueTask ConsumeNonQueryAsync( + CancellationToken cancellationToken = default) { - _consumeNonQuery = true; - _nonQueryRecordsAffected = -1; - var enumerator = GetAsyncEnumerator(cancellationToken); + var recordsAffected = -1L; + if (cancellationToken.CanBeCanceled) + _state.FlowToken = cancellationToken; try { - // Release a body parked pre-publication and wake the flow. - _callerInteractionCore.ResumeBody(runContinuationsAsynchronously: false); - _callerInteractionCore.WakeBody(); - _ = await EnumeratorMoveNextTask.ConfigureAwait(false); - // Internal consumption owns error delivery. Errors collected during the internal - // drain have no other outlet: DisposeAsync deliberately skips a completed flow. - if (_drainErrors is { Count: > 0 } errors) - throw errors.Count == 1 ? errors[0] : new AggregateException(errors); - return _nonQueryRecordsAffected; + while (await MoveNextAsync(cancellationToken).ConfigureAwait(false)) + { + var result = _state.Current!; + await result.CompleteAsync().ConfigureAwait(false); + var affected = result.GetCommandComplete().BatchRecordsAffected; + if (affected >= 0) + recordsAffected = recordsAffected < 0 + ? affected + : checked(recordsAffected + affected); + } + return recordsAffected; } finally { - await enumerator.DisposeAsync().ConfigureAwait(false); + await DisposeAsync().ConfigureAwait(false); } } - // The token in force for the current read: the flow token once fired (whole-flow cancel), else the - // per-read token if cancelable, else the flow token. - CancellationToken EffectiveCancellationToken - => GetEffectiveCancellationToken(Volatile.Read(ref _cancellationState)); + CommandExecutionColdState GetOrCreateColdState() + => Volatile.Read(ref _state.ColdState) ?? + Interlocked.CompareExchange(ref _state.ColdState, new(), null) ?? _state.ColdState; + + bool IsClosed => Volatile.Read(ref _state.ColdState)?.CloseException is not null; + bool IsCancelRequested => Volatile.Read(ref _state.ColdState) is { CancelRequested: true }; + bool HasDecoder => _state.ContextPublished && _ops.HasSuccessfulActivation; - static CancellationToken GetEffectiveCancellationToken(CancellationState? cancellation) - => cancellation is null ? default - : cancellation.FlowToken.IsCancellationRequested ? cancellation.FlowToken - : cancellation.CallerToken.CanBeCanceled ? cancellation.CallerToken - : cancellation.FlowToken; + internal ValueTask ExecuteAuto(PgClientFlow.Context context) + { + _state.Context = context; + _state.ContextPublished = true; + ValueTask writeTask; + try + { + ref readonly var template = ref _state.Commands.ItemRef(_state.Commands.Count - 1); + var appendSync = !template.WithSync; + _state.ReadFlowRfq = appendSync; + // Caller cancellation never cancels wire I/O. The consumer observes the latched intent and + // drains its command to RFQ instead. + writeTask = _ops.IsAsync + ? _state.Commands.WriteCommandsAsync(context.GetEncoder(), appendSync, default) + : WriteCommandsResumable(context, appendSync); + // Observe synchronous faults here; pending writes remain the framework-owned trailing task. + if (writeTask.IsCompleted) + writeTask.GetAwaiter().GetResult(); + } + catch (Exception ex) + { + // The framework recovers the wire from this throw. Only the consumer needs the fault. + FaultReady(ex); + throw; + } - public int CommandCount => _commands.Count; - internal virtual int VisibleCommandCount => _commands.VisibleCount; - public bool IsResultReady => _isResultReady; + // Activation may precede or follow execution. Bridging it here guarantees a consumer resumes + // against a published context, and delivers an activation fault to the pipeline task when no + // consumer ever arrives. + var activation = context.GetDecoderAsync().ConfigureAwait(false); + if (activation.IsCompleted) + OnActivationSettled(onExecutorStrand: true); + else + activation.UnsafeOnCompleted(static state => + new CommandFlowCore(TOps.Create((PgClientFlow)state!)) + .OnActivationSettled(onExecutorStrand: + !((PgClientFlow)state!).ActivationWasDispatched), _ops.Flow); + return new(new FlowTasks(writeTask, new ValueTask((IValueTaskSource)_ops.Flow, _state.PipelineTaskSource.Version))); + } - public Enumerator GetEnumerator() + [MethodImpl(MethodImplOptions.NoInlining)] + ValueTask WriteCommandsResumable(PgClientFlow.Context context, bool appendSync) { - return new Enumerator(this); + var encoder = context.GetEncoder(); + ValueTask writeTask; + using (encoder.BeginResumableWriteScope()) + writeTask = _state.Commands.WriteCommandsResumable(encoder, appendSync); + return writeTask.IsCompleted ? writeTask : encoder.RunResumableTask(writeTask); } - // Bind at submission because eager writing precedes the first MoveNextAsync. - internal override void BindCallerToken(CancellationToken cancellationToken) - => GetOrCreateCancellationState().FlowToken = cancellationToken; - internal override CancellationToken MigrationCancellationToken - => Volatile.Read(ref _cancellationState)?.FlowToken ?? default; + // A detached activation callback may publish ready inline on that scheduler turn. ExecuteAuto + // can also observe an already-completed activation directly on the executor, while a pending + // zero-edge activation completes inline there; both executor cases need the scheduling firewall. + void OnActivationSettled(bool onExecutorStrand) + { + try + { + _ = _state.Context.GetDecoderAsync().ConfigureAwait(false).GetAwaiter().GetResult(); + } + catch (Exception fault) + { + // Preserve the activation's close, timeout, or cancellation identity. The pipeline task + // faults first so a consumer woken by the ready source never races its retirement. + CompletePipelineTask(fault, runContinuationsAsynchronously: true); + FaultReady(fault); + return; + } - public Enumerator GetAsyncEnumerator(CancellationToken cancellationToken = default) + var cancellationRequested = IsCancelRequested; + if (cancellationRequested) + RequestBackendCancellation(); + var context = _state.Context; + var stopping = context.StoppingToken.IsCancellationRequested; + var stoppingException = stopping ? context.FlowTerminationException : null; + var dispatchReady = (_ops.IsAsyncAtDispatch && onExecutorStrand) + || stopping || cancellationRequested; + if (!CompleteReady(null, runContinuationsAsynchronously: dispatchReady)) + { + // Teardown released the consumer while this flow waited for its turn. Nothing reads the + // response, the closing wire owns it. + CompletePipelineTask(null, runContinuationsAsynchronously: true); + return; + } + // Termination propagation enumerates pipeline positions best-effort. A flow can cross from + // the in-flight store into the activated slot while that pass is being taken and miss it. + // Recheck after publishing readiness: OnStopping arbitrates PhaseInitial against a consumer's + // PhaseReading claim, so exactly one side owns the decoder and eventual pipeline completion. + if (stopping) + { + OnStopping(stoppingException!); + return; + } + // A cancel latched before activation may have released its caller already. The response + // still has to reach RFQ, so drain it unless a consumer already owns the decoder. + if (cancellationRequested) + TryTakeOverDrain(); + } + + void FaultReady(Exception exception) { - // Body, teardown, and cancellation may complete the source concurrently. A missing per-call token - // must not replace the flow token captured at submission. - _enumeratorMoveNextTaskSource.CanCompleteConcurrently = true; - if (cancellationToken.CanBeCanceled) - GetOrCreateCancellationState().FlowToken = cancellationToken; - return new(this); + Interlocked.CompareExchange(ref GetOrCreateColdState().TerminalException, exception, null); + CompleteReady(exception, runContinuationsAsynchronously: true); } - protected override ValueTask ExecuteAuto(Context context) + bool CompleteReady(Exception? exception, bool runContinuationsAsynchronously) + { + if (Interlocked.CompareExchange(ref _state.ReadyCompletion, 1, 0) != 0) + return false; + if (exception is null) + _state.ReadySource.SetResult(true, runContinuationsAsynchronously); + else + _state.ReadySource.SetException(exception, runContinuationsAsynchronously); + return true; + } + + void CompletePipelineTask(Exception? exception, bool runContinuationsAsynchronously = false) { - if (!IsAsync && _callerInteractionCore.IsWaiting) - return ExecuteAfterHandoff(context); + if (Interlocked.Exchange(ref _state.Phase, PhaseCompleted) is PhaseCompleted) + return; + if (exception is null) + _state.PipelineTaskSource.SetResult(true, runContinuationsAsynchronously); + else + _state.PipelineTaskSource.SetException(exception, runContinuationsAsynchronously); + } - return new(ExecuteAutoCore(context)); + void EnsureSyncHandoff() + { + if (_ops.IsAsyncAtDispatch) + ThrowHelper.ThrowInvalidOperation( + "Synchronous result consumption requires a flow initialized for synchronous execution."); + if (_state.SyncHandoffClaimed) + return; + _ops.WaitForSyncHandoff(); + _state.SyncHandoffClaimed = true; } - [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] - async ValueTask ExecuteAfterHandoff(Context context) + internal bool MoveNext() + { + EnsureSyncHandoff(); + while (true) + { + var phase = Volatile.Read(ref _state.Phase); + switch (phase) + { + case PhaseInitial: + if (Interlocked.CompareExchange(ref _state.Phase, PhaseReading, PhaseInitial) != PhaseInitial) + continue; + _state.CommandIndex = 0; + return First(); + case PhaseResultReady: + if (Interlocked.CompareExchange(ref _state.Phase, PhaseReading, PhaseResultReady) != PhaseResultReady) + continue; + return NextBatch(); + case PhaseReading: + ThrowHelper.ThrowInvalidOperation("A read is already in progress on this flow."); + return false; + case PhaseDraining: + _ops.Flow.WaitForCompleteSynchronously(); + throw Volatile.Read(ref _state.ColdState)?.TerminalException + ?? ThrowHelper.ThrowInvalidOperation("The flow was disposed."); + default: + if (Volatile.Read(ref _state.ColdState)?.TerminalException is { } terminal) + ExceptionDispatchInfo.Throw(terminal); + return false; + } + } + } + + bool First() { try { - await YieldToCaller(); + WaitForReadySynchronously(); + Debug.Assert(!_state.ConsumerDetached); + RegisterCancellation(default); + var result = IsSinglePublishedCommand + ? ReadResult(0) + : ReadNextPublishedResult(); + return result is not null && PublishSynchronousResult(result); + } + catch (TimeoutException ex) + { + HandleReadTimeout(ex); + throw; } catch (Exception ex) { - TerminateBodyBeforeStart(); - CompleteEnumerationWithException(ex); + FaultFromOwner(ex); throw; } - - return ExecuteAutoCore(context); } - FlowTasks ExecuteAutoCore(Context context) + bool NextBatch() { - _context = context; - Volatile.Write(ref _contextPublished, true); - if (Volatile.Read(ref _cancellationState) is { } cancellation) - { - if (cancellation.FlowToken.IsCancellationRequested) - RequestCancel(cancellation.FlowToken, CancellationScope.RemainingFlow); - else if (Volatile.Read(ref cancellation.Requested)) - RequestBackendCancellation((BackendCancellationTiming)Volatile.Read(ref cancellation.Timing), - Volatile.Read(ref cancellation.Delivery)); - } - ValueTask writeTask; try { - // Writes are independent of consumer admission. Inter-result gates provide backpressure. - // Async writes use transport completion; sync writes use the resumable non-blocking path so - // the caller thread retains execution ownership across readiness waits. - var encoder = IsAsync ? default : context.GetEncoder(); - var appendSync = !_commands[CommandCount - 1].WithSync; - _readFlowRfq = appendSync; - if (IsAsync) + RegisterCancellation(default); + var result = _state.Current!; + var completeError = CompleteCurrentResult(); + _state.CurrentPublished = false; + if (Volatile.Read(ref _state.ColdState)?.TerminalException is { } consumerFault) { - // Caller cancellation never cancels wire I/O. A partially cancelled write requires - // protocol recovery and can strand already-pipelined successors; the body instead - // observes the latched intent and drains every written command to RFQ. - writeTask = _commands.WriteCommandsAsync(context.GetEncoder(), appendSync, default); + Interlocked.Exchange(ref _state.Phase, PhaseDraining); + NotifyDrainStarted(); + _state.ConsumerDetached = true; + Drain(); + ExceptionDispatchInfo.Throw(consumerFault); } - else + if (completeError is { TransactionStatus: TransactionStatus.Unknown }) + SkipDiscardedCommands(); + + _state.CommandIndex++; + if (IsSinglePublishedCommand) { - using (encoder.BeginResumableWriteScope()) - writeTask = _commands.WriteCommandsResumable(encoder, appendSync); + CompleteBatch(); + _state.ConsumerObservedCompletion = true; + return false; } + var next = ReadNextPublishedResult(); + if (next is not null) + return PublishSynchronousResult(next); - // Observe synchronous faults here; pending writes remain the framework-owned trailing task. - if (writeTask.IsCompleted) - writeTask.GetAwaiter().GetResult(); - else if (!IsAsync) - writeTask = encoder.RunResumableTask(writeTask); + return false; + } + catch (TimeoutException ex) + { + HandleReadTimeout(ex); + throw; } catch (Exception ex) { - TerminateBodyBeforeStart(); - CompleteEnumerationWithException(ex); + FaultFromOwner(ex); throw; } + } + + bool PublishSynchronousResult(CommandResult result) + { + _state.Current = result; + _state.CurrentPublished = true; + Interlocked.Exchange(ref _state.Phase, PhaseResultReady); + var context = _state.Context; + if (!IsClosed && context.StoppingToken.IsCancellationRequested) + Interlocked.CompareExchange(ref GetOrCreateColdState().CloseException, + context.FlowTerminationException, null); + if (!IsCancelRequested && !IsClosed) + return true; + + if (Interlocked.CompareExchange(ref _state.Phase, PhaseReading, PhaseResultReady) == PhaseResultReady) + Drain(); + else + _ops.Flow.WaitForCompleteSynchronously(); + throw Volatile.Read(ref _state.ColdState)?.TerminalException + ?? ThrowHelper.ThrowUnexpected("A latched flow completed without a terminal outcome."); + } - // Read and write run concurrently; the framework observes the trailing write before releasing - // the flow, preserving single-writer tenure without blocking reads behind socket backpressure. - return new FlowTasks( - trailingExecutionTask: writeTask, - pipelineTask: DispatchPipelinedRead(context, context.GetProtocolStatic().ReadPromise)); + void WaitForReadySynchronously() + { + var ready = new ValueTask((IValueTaskSource)_ops.Flow, _state.ReadySource.Version); + if (ready.IsCompleted) + _ = ready.GetAwaiter().GetResult(); + else + _ = ready.AsTask().GetAwaiter().GetResult(); } - // Defer state-machine creation until activation because all flows share one protocol-static promise. - ValueTask DispatchPipelinedRead(Context context, ValueTaskSourcePromise promise) + internal ValueTask MoveNextAsync(CancellationToken cancellationToken) { - // The shared promise may be tenured only after successful decoder activation. - var waiter = context.GetDecoderAsync().ConfigureAwait(false); - if (waiter.IsCompleted) + if (!_ops.IsAsyncAtDispatch) + return ValueTask.FromException(ThrowHelper.ThrowInvalidOperation( + "Asynchronous result consumption requires a flow initialized for asynchronous execution.")); + while (true) { - // Only successful activation owns the shared promise. A settled fault belongs to this flow's - // private completion source because it never claimed the wire. - if (!waiter.IsCompletedSuccessfully) + var phase = Volatile.Read(ref _state.Phase); + switch (phase) { - // Preserve the activation's close, timeout, or cancellation identity. - try { waiter.GetAwaiter().GetResult(); } - catch (Exception ex) { _executePipelinedCore.SetException(ex); } - return new ValueTask(this, _executePipelinedCore.Version); - } - // Handing the shared-promise-backed task to the framework is safe: the contract guarantees - // the waiter is consumed (releasing the promise tenure via GetResult's Reset) before the - // item's position is republished, so a successor's dispatch always finds the tenure released. - PromiseAsyncValueTaskMethodBuilder.Promise = promise; - try - { - return ExecutePipelined(context); - } - finally - { - PromiseAsyncValueTaskMethodBuilder.Promise = null; + case PhaseInitial: + if (cancellationToken.IsCancellationRequested) + return CancelBeforeRead(cancellationToken); + if (Interlocked.CompareExchange(ref _state.Phase, PhaseReading, PhaseInitial) != PhaseInitial) + continue; + _state.CommandIndex = 0; + _state.WindowToken = cancellationToken; + return FirstAsync(); + case PhaseResultReady: + if (cancellationToken.IsCancellationRequested) + return CancelBeforeRead(cancellationToken); + if (Interlocked.CompareExchange(ref _state.Phase, PhaseReading, PhaseResultReady) != PhaseResultReady) + continue; + RegisterCancellation(cancellationToken); + return NextBatchAsync(); + case PhaseReading: + return ValueTask.FromException( + ThrowHelper.ThrowInvalidOperation("A read is already in progress on this flow.")); + case PhaseDraining: + return AwaitTakeoverAsync(); + default: + return Volatile.Read(ref _state.ColdState)?.TerminalException is { } terminal + ? ValueTask.FromException(terminal) + : new(false); } } + } - _pipelinePromise = promise; - // Static continuation: a bridge into framework state, so no captured scheduling context is needed. - waiter.OnCompleted(static state => - { - var flow = (CommandFlow)state!; - var ctx = flow._context; - // A faulted activation never claimed the shared promise; complete only this flow's source. - var activation = ctx.GetDecoderAsync().GetAwaiter(); - if (!activation.IsCompletedSuccessfully) - { - // Preserve the activation's close, timeout, or cancellation identity. - try { activation.GetResult(); } - catch (Exception ex) { flow._executePipelinedCore.SetException(ex); } - return; - } - var promise = flow._pipelinePromise!; - PromiseAsyncValueTaskMethodBuilder.Promise = promise; - ValueTask task = flow.ExecutePipelined(ctx); - try - { - if (!task.IsCompleted) - { - flow._task = task; - ((IValueTaskSource)promise).OnCompleted(static state => - { - var flow = (CommandFlow)state!; - try - { - flow._task.GetAwaiter().GetResult(); - flow._executePipelinedCore.SetResult(true); - } - catch (Exception ex) - { - flow._executePipelinedCore.SetException(ex); - } - // This internal bridge runs no user code and requires no ExecutionContext flow. - }, flow, promise.Token, ValueTaskSourceOnCompletedFlags.None); - } - else - { - try - { - task.GetAwaiter().GetResult(); - flow._executePipelinedCore.SetResult(true); - } - catch (Exception ex) - { - flow._executePipelinedCore.SetException(ex); - } - } - } - finally - { - PromiseAsyncValueTaskMethodBuilder.Promise = null; - } - }, this); + // A pre-cancelled token releases the caller immediately. The wire still drains to RFQ. + ValueTask CancelBeforeRead(CancellationToken cancellationToken) + { + RequestCancel(cancellationToken, CommandExecutionCancellationScope.CurrentWindow); + return ValueTask.FromException(new OperationCanceledException(cancellationToken)); + } - return new ValueTask(this, _executePipelinedCore.Version); + // The consumer parks behind a takeover drain and receives the outcome that caused it. + async ValueTask AwaitTakeoverAsync() + { + await WaitForCompletionAsync().ConfigureAwait(false); + throw Volatile.Read(ref _state.ColdState)?.TerminalException ?? ThrowHelper.ThrowInvalidOperation("The flow was disposed."); + } + + ValueTask FirstAsync() + { + var ready = new ValueTask((IValueTaskSource)_ops.Flow, _state.ReadySource.Version); + if (ready.IsCompletedSuccessfully) + { + _ = ready.Result; + return FirstAfterReadyAsync(); + } + if (_state.FirstPromise is not { } promise) + return AwaitReadyPooledAsync(ready); + using (PromiseAsyncValueTaskMethodBuilder.BeginCallScope(promise)) + return AwaitReadyRetainedAsync(ready); } [RuntimeAsyncMethodGeneration(false)] - [AsyncMethodBuilder(typeof(PromiseAsyncValueTaskMethodBuilder))] - async ValueTask ExecutePipelined(Context context) + [AsyncMethodBuilder(typeof(PromiseAsyncValueTaskMethodBuilder<>))] + async ValueTask AwaitReadyRetainedAsync(ValueTask ready) { - // Pre-start teardown may have already claimed terminality. In that case the consumer has its - // fault and this late dispatch has no body tenure to establish. - if (Interlocked.CompareExchange(ref _bodyState, BodyRunning, BodyNotStarted) != BodyNotStarted) - return; try { - // If we have a continuation stored we must already be on the caller thread, - // otherwise we must make sure to unblock the executor (see comment in the write phase). - // This first handoff is body execution too: close may fault it, so it belongs inside the - // same terminal envelope that publishes the consumer fault and body termination. - if (!IsAsync && !_callerInteractionCore.HasHandoff) - await YieldToCaller(); - - // User cancellation must not cancel activation or wire I/O. The flow observes it after - // activation, drains itself to RFQ, then delivers OCE without invoking pipeline recovery. - _decoder = await context.GetDecoderAuto().ConfigureAwait(false); - var publishedResult = false; - while (++_commandIndex < CommandCount) - { - _isResultReady = false; - bool hasPreparedDescription; - bool suppressEnumeration; - bool describeForPreparation; - { - ref readonly var command = ref _commands.ItemRef(_commandIndex); - _decoder.UseReadTimeout(command.Timeout); - suppressEnumeration = command.SuppressEnumeration; - describeForPreparation = command.DescribeForPreparation; - hasPreparedDescription = command.Descriptor is { IsPrepared: true, PreparedRowDescription: not null } - && !command.DescribeOnly; - } - - // Registrations only latch and wake; terminal delivery remains body-owned. Dispose them - // before promise tenure ends so callbacks cannot reach the next flow. Do not rearm after - // consumer disposal, where a persistent cancellation could escape its intended wait. - if (Volatile.Read(ref _cancellationState) is { } cancellationAtReadStart && !IsDraining - && (cancellationAtReadStart.CallerToken.CanBeCanceled - || cancellationAtReadStart.FlowToken.CanBeCanceled)) - RegisterCancellationCallbacks(cancellationAtReadStart); - // After close, a fresh command must not consume bytes left by its predecessor. A draining - // flow may continue reading its own response to restore RFQ. - if (!IsDraining && context.IsProtocolClosed) - throw context.FlowTerminationException; - - ParameterTypeList describedParameterTypes = default; - if (describeForPreparation) - { - var rowDescription = context.GetProtocolStatic().RowDescription; - if (IsAsync) - { - (_pgError, describedParameterTypes, _requestedRowDescription) = - await _commands.ItemRef(_commandIndex) - .ReadPreparationDescriptionAsync(_decoder, rowDescription).ConfigureAwait(false); - } - else - { - (_pgError, describedParameterTypes, _requestedRowDescription) = - _commands.ItemRef(_commandIndex) - .ReadPreparationDescription(_decoder, rowDescription); - } - } - else if (IsAsync && hasPreparedDescription) - { - // Prepared commands with a known description have the compact BindComplete -> - // DataRow/CommandComplete prelude. Await the decoder directly so a read wake resumes - // this outer body rather than a nested parser coroutine; the second message normally - // comes from the same batch and is consumed synchronously. - if (!_decoder.TryMoveNext()) - { - if (!await _decoder.MoveNextAsync().ConfigureAwait(false)) - _decoder.ThrowUnexpectedEof(); - } - var message = _decoder.Current; - - if (message.EnsureExpectedOrError(PgTypes.BackendType.BindComplete) is { } bindError) - { - _pgError = bindError; - _requestedRowDescription = null; - } - else - { - if (!_decoder.TryMoveNext()) - { - if (!await _decoder.MoveNextAsync().ConfigureAwait(false)) - _decoder.ThrowUnexpectedEof(); - } - message = _decoder.Current; - message.DebugEnsureExpected(PgTypes.BackendType.DataRow, PgTypes.BackendType.CommandComplete); - _pgError = null; - _requestedRowDescription = null; - } - } - else if (IsAsync) - { - var rowDescription = context.GetProtocolStatic().RowDescription; - var read = _commands.ItemRef(_commandIndex).ReadUntilExecuteAsync(_decoder, rowDescription); - (_pgError, _requestedRowDescription) = await read.ConfigureAwait(false); - } - else - { - var rowDescription = context.GetProtocolStatic().RowDescription; - (_pgError, _requestedRowDescription) = _commands.ItemRef(_commandIndex) - .ReadUntilExecute(_decoder, rowDescription); - } - - // A draining consumer cannot observe CommandResult, so retain read errors for disposal. - var capturedThisCommand = false; - var readErrorIsOwnCancellation = _pgError is { } readError - && IsOwnCancellation(readError); - if (IsConsumingAutonomously && _pgError is { } readErrorToCapture - && !readErrorIsOwnCancellation) - { - (_drainErrors ??= new()).Add(PgErrorException.Create(readErrorToCapture)); - capturedThisCommand = true; - } - - // Await in-flight callbacks before releasing shared promise tenure. - var cancellation = Volatile.Read(ref _cancellationState); - if (cancellation is not null) - await DisposeCancellationRegistrations(cancellation).ConfigureAwait(false); - // Cancellation switches to autonomous drain; terminal delivery follows RFQ and tenure release. - var effectiveCancellationToken = GetEffectiveCancellationToken(cancellation); - if ((cancellation is { } && Volatile.Read(ref cancellation.Requested) - || effectiveCancellationToken.IsCancellationRequested) && !IsEnumerationCompleted) - { - cancellation ??= GetOrCreateCancellationState(); - if (effectiveCancellationToken.IsCancellationRequested) - cancellation.DeliverToken = effectiveCancellationToken; - cancellation.DeliverOce = true; - if (!IsDraining) - MarkBodyInitiatedDrain(); - } - - CommandResult result; - { - ref readonly var readState = ref context.GetProtocolStatic(); - readState.ResultMessageEnumerator.Initialize(this, _decoder); - result = _enumeratorCurrent ?? readState.CommandResult; - - ref readonly var resultCommand = ref _commands.ItemRef(_commandIndex); - var descriptor = resultCommand.Descriptor; - // We were preparing and we have no error from parse, make a prepared descriptor. - if (!descriptor.IsPrepared && !descriptor.CommandName.IsDefault - && (_pgError is not { } err || !err.Expected.Contains(PgTypes.BackendType.ParseComplete))) - { - descriptor = CommandDescriptor.CreatePrepared( - descriptor.CommandName, - describeForPreparation ? describedParameterTypes : descriptor.ParameterTypes, - _requestedRowDescription?.Preserve()); - } - result.Initialize(this, _commandIndex, descriptor, _requestedRowDescription, - !resultCommand.DescribeOnly, resultCommand.IsSimple(), _pgError); - } - ((CommandFlowObserver?)GetObserver(out var observerState)) - ?.OnCommandResult(this, result, observerState); - - // Disposal drains without another result handoff. Graceful close instead faults the - // attached consumer, then uses the same autonomous drain. Command errors remain results. - if (context.StoppingToken.IsCancellationRequested && !IsDraining - && !IsEnumerationCompleted) - { - // Latch the close (a consumer that Resets past this point self-delivers it), wake a - // parked consumer, then drain. - var close = context.FlowTerminationException; - _callerInteractionCore.SetCloseLatch(close); - CompleteEnumerationWithClose(close); - MarkBodyInitiatedDrain(); - } - var consumeInternally = IsConsumingNonQuery || suppressEnumeration; - if (!IsDraining && !consumeInternally) - { - // Eager async execution must wait for the consumer to arm generation zero before - // publishing its first result. Synchronous execution already runs on that caller. - if (!publishedResult && IsAsync) - { - await _callerInteractionCore.WaitForCaller(this).ConfigureAwait(false); - EnterStoppingDrainIfNeeded(context); - } - - if (!IsDraining && !IsConsumingNonQuery) - { - _isResultReady = true; - publishedResult = true; - // Result continuations run asynchronously so the body can reach the next gate - // before user code asks for the next result. Buffered batches then advance inline - // from MoveNextAsync instead of suspending one Task state machine per result. - SetResult(result); - - if (!IsDraining && !IsConsumingNonQuery) - { - if (IsAsync) - { - await _callerInteractionCore.WaitForCaller(this).ConfigureAwait(false); - EnterStoppingDrainIfNeeded(context); - } - else - await YieldToCaller(); - - /* The next MoveNext or MoveNextAsync call resumes here. */ - } - } - } - else if (!_drainModeEntered && IsAsyncAtDispatch && !IsAsync) - { - // An async I/O wake raced a synchronous disposer before the body reached its handoff. - if (WaitForDrainOnDispose) - { - // Hand the continuation to the disposer, which waits on the rendezvous rather than - // this task and can therefore drive the remaining drain without sync-over-async. - await YieldToCaller(); - } - else - { - // No disposer is waiting to drive; retain asynchronous background draining. - IsAsync = IsAsyncAtDispatch; - } - } - // Preserve the drive mode chosen on first drain entry. - _drainModeEntered = _drainModeEntered || IsDraining; - - // Disposing the message enumerator completes the command and, in drain mode, consumes its - // remaining rows. Re-read the current execution mode after every resumption. - (PgError Error, TransactionStatus TransactionStatus)? completeError; - // Consumption mode may change while the body is suspended; use the current value. - if (consumeInternally || IsConsumingNonQuery) - { - while (_decoder.Current.Header.Type is PgTypes.BackendType.DataRow) - { - if (!_decoder.TryMoveNext()) - await _decoder.GetNextAsync().ConfigureAwait(false); - } - result.CompleteNonQuery(_decoder.Current); - var completion = _commands.ItemRef(_commandIndex).CompleteAsync(_decoder); - completeError = await completion.ConfigureAwait(false); - if (_pgError is null && completeError is null) - { - var recordsAffected = result.GetCommandComplete().BatchRecordsAffected; - if (recordsAffected >= 0) - _nonQueryRecordsAffected = _nonQueryRecordsAffected < 0 - ? recordsAffected - : checked(_nonQueryRecordsAffected + recordsAffected); - } - } - else if (IsAsync) - { - var resultEnumerator = context.GetProtocolStatic().ResultMessageEnumerator; - await resultEnumerator.DisposeAsync().ConfigureAwait(false); - completeError = resultEnumerator.CompleteError; - } - else - { - var resultEnumerator = context.GetProtocolStatic().ResultMessageEnumerator; - resultEnumerator.Dispose(); - completeError = resultEnumerator.CompleteError; - } - - var resultErrorIsOwnCancellation = result.Error is { } resultError - && IsOwnCancellation(resultError); - if (suppressEnumeration && result.Error is { } suppressedError - && !resultErrorIsOwnCancellation) - { - (_drainErrors ??= new()).Add(PgErrorException.Create(suppressedError)); - capturedThisCommand = true; - if (!IsDraining) - MarkBodyInitiatedDrain(); - } - - { - // Accumulate each command's fresh error while draining, but do not duplicate an error - // already captured during its read phase or delivered to a live consumer. - var completeErrorIsOwnCancellation = completeError is { } completedWithError - && IsOwnCancellation(completedWithError.Error); - if ((consumeInternally || IsConsumingNonQuery || IsDraining && !_isResultReady) - && !capturedThisCommand && completeError is { } err - && !completeErrorIsOwnCancellation) - (_drainErrors ??= new()).Add(PgErrorException.Create(err.Error)); - } - - // Extended-query errors discard every following command through the next Sync. Skip - // those commands locally and consume the RFQ which is their only wire response. - if (completeError is { TransactionStatus: TransactionStatus.Unknown }) - { - while (++_commandIndex < CommandCount && !_commands[_commandIndex].WithSync) { } + await ready.ConfigureAwait(false); + } + catch (TimeoutException ex) + { + HandleReadTimeout(ex); + throw; + } + catch (Exception ex) + { + FaultFromOwner(ex); + throw; + } + return await FirstAfterReadyAsync().ConfigureAwait(false); + } - if (IsAsync) - await ReadRfqAsync(_decoder).ConfigureAwait(false); - else - ReadRfq(_decoder); + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] + async ValueTask AwaitReadyPooledAsync(ValueTask ready) + { + try + { + await ready.ConfigureAwait(false); + } + catch (TimeoutException ex) + { + HandleReadTimeout(ex); + throw; + } + catch (Exception ex) + { + FaultFromOwner(ex); + throw; + } + return await FirstAfterReadyAsync().ConfigureAwait(false); + } - // Reaching the end means the discarded segment terminated at our appended Sync. - if (_commandIndex == CommandCount) - _readFlowRfq = false; - } + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] + async ValueTask FirstAfterReadyAsync() + { + Exception? deliver; + try + { + Debug.Assert(!_state.ConsumerDetached); + RegisterCancellation(TakeWindowToken()); + var result = IsSinglePublishedCommand + ? await ReadResultAsync().ConfigureAwait(false) + : await ReadNextPublishedResultAsync().ConfigureAwait(false); + if (result is null) + return false; + _state.Current = result; + _state.CurrentPublished = true; + // Publish the idle state, then recheck the latches. A latch that landed between the read + // and this publication found no idle owner to take over, so this frame must act on it. + Interlocked.Exchange(ref _state.Phase, PhaseResultReady); + // Graceful stopping faults a result that arrives after the close began, as the ordinary + // flow does at each result boundary. Latch it so the drain delivers that close. + var context = _state.Context; + if (!IsClosed && context.StoppingToken.IsCancellationRequested) + Interlocked.CompareExchange(ref GetOrCreateColdState().CloseException, + context.FlowTerminationException, null); + if (!IsCancelRequested && !IsClosed) + return true; + if (Interlocked.CompareExchange(ref _state.Phase, PhaseReading, PhaseResultReady) == PhaseResultReady) + { + await DrainAsync().ConfigureAwait(false); } - - // The framework observes trailing write failure before releasing this flow. - if (_readFlowRfq) + else { - if (_decoder.TryMoveNext()) - { - var message = _decoder.Current; - if (message.EnsureExpectedOrError(PgTypes.BackendType.ReadyForQuery) is { } rfqError) - PgErrorException.Throw(rfqError); - } - else if (IsAsync) - { - await ReadRfqAsync(_decoder).ConfigureAwait(false); - } - else - { - ReadRfq(_decoder); - } + // The latching side took the decoder first. Park behind its drain. + await WaitForCompletionAsync().ConfigureAwait(false); } - - SetResult(null); + deliver = Volatile.Read(ref _state.ColdState)?.TerminalException; } - catch (PgClientClosedException) when (context.IsProtocolClosed) + catch (TimeoutException ex) { - // Scope to our own closure so a nested protocol's close doesn't get treated as ours. - // Latch the close so a consumer that Resets after this point self-delivers it. - _callerInteractionCore.SetCloseLatch(context.FlowTerminationException); - // A detached consumer treats close as drain completion; a live consumer observes the close. - if (IsDraining) - { - if (!IsEnumerationCompleted) - SetResult(null); - return; - } - CompleteEnumerationWithException(context.FlowTerminationException); + HandleReadTimeout(ex); throw; } - catch (OperationCanceledException ex) when (IsCancellationToken(ex.CancellationToken)) + catch (Exception ex) { - CompleteEnumerationWithException(ex); + FaultFromOwner(ex); throw; } - catch (TimeoutException ex) + throw deliver ?? ThrowHelper.ThrowUnexpected("A latched flow completed without a terminal outcome."); + } + + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] + async ValueTask NextBatchAsync() + { + try { - CompleteEnumerationWithException(ex); - RequestCancel(default, CancellationScope.RemainingFlow, BackendCancellationTiming.Immediate, - BackendCancellationTiming.AtReadFrontier, allowCompletedEnumeration: true); - if (context.IsProtocolClosed) - throw; - - // The timeout is terminal for the consumer, not for the body. Keep ownership of the - // command sequence and drain every remaining RFQ window; reaching each RFQ requests - // cancellation for the next window through OnCancellationWindowCompleted. Recovery is - // reserved for a failure of this semantic drain, where only wire obligations remain. - if (Volatile.Read(ref _cancellationState) is { } cancellation) - await DisposeCancellationRegistrations(cancellation).ConfigureAwait(false); - ((CommandFlowObserver?)GetObserver(out var observerState)) - ?.OnDrainStarted(this, observerState); - try + var result = _state.Current!; + await _state.Context.GetProtocolStatic() + .ResultMessageEnumerator.DisposeAsync().ConfigureAwait(false); + var skipDiscarded = _state.Context.GetProtocolStatic() + .ResultMessageEnumerator.CompleteError + is { TransactionStatus: TransactionStatus.Unknown }; + _state.CurrentPublished = false; + if (Volatile.Read(ref _state.ColdState)?.TerminalException is not null) { - while (context.OutstandingRfqCount != 0) - _ = await _decoder!.GetNextAuto().ConfigureAwait(false); + Interlocked.Exchange(ref _state.Phase, PhaseDraining); + NotifyDrainStarted(); + _state.ConsumerDetached = true; + await DrainAsync().ConfigureAwait(false); + ExceptionDispatchInfo.Throw( + Volatile.Read(ref _state.ColdState)!.TerminalException!); } - catch (TimeoutException) + if (skipDiscarded) + await SkipDiscardedCommandsAsync().ConfigureAwait(false); + + _state.CommandIndex++; + if (IsSinglePublishedCommand) { - // The semantic drain owns the same cancellation episode. A timeout here would - // otherwise bypass the outer catch and leave the episode unaware that its first - // read-timeout escalation produced no protocol progress. - RequestCancel(default, CancellationScope.RemainingFlow, - BackendCancellationTiming.Immediate, BackendCancellationTiming.AtReadFrontier, - allowCompletedEnumeration: true); - throw; + await CompleteBatchAsync().ConfigureAwait(false); + _state.ConsumerObservedCompletion = true; + return false; } - return; + var next = await ReadNextPublishedResultAsync().ConfigureAwait(false); + if (next is not null) + { + result = next; + _state.Current = result; + _state.CurrentPublished = true; + Interlocked.Exchange(ref _state.Phase, PhaseResultReady); + var context = _state.Context; + if (!IsClosed && context.StoppingToken.IsCancellationRequested) + Interlocked.CompareExchange(ref GetOrCreateColdState().CloseException, + context.FlowTerminationException, null); + if (!IsCancelRequested && !IsClosed) + return true; + + if (Interlocked.CompareExchange(ref _state.Phase, PhaseReading, PhaseResultReady) == PhaseResultReady) + await DrainAsync().ConfigureAwait(false); + else + await WaitForCompletionAsync().ConfigureAwait(false); + throw Volatile.Read(ref _state.ColdState)?.TerminalException + ?? ThrowHelper.ThrowUnexpected("A latched flow completed without a terminal outcome."); + } + + return false; } - catch (Exception ex) + catch (TimeoutException ex) { - CompleteEnumerationWithException(ex); + HandleReadTimeout(ex); throw; } - finally + catch (Exception ex) { - // The body is the sole owner of protocol-static row metadata. Recovery consumes only - // decoder/wire state, so a faulted body can release oversized storage while recovery - // retains the failed flow's framework tenure. - ref readonly var readState = ref context.GetProtocolStatic(); - readState.Reset(); - PublishBodyTerminated(); + FaultFromOwner(ex); + throw; } - void SetResult(CommandResult? next) + } + + CancellationToken TakeWindowToken() + { + var token = _state.WindowToken; + _state.WindowToken = default; + return token; + } + + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] + async ValueTask ReadNextPublishedResultAsync() + { + while (_state.CommandIndex < _state.Commands.Count) { - var completed = next is null; - if (completed) + SetCurrent(await ReadResultAsync().ConfigureAwait(false)); + _state.CurrentPublished = false; + if (!_state.Commands.ItemRef(_state.CommandIndex).SuppressEnumeration) + return _state.Current; + + var completeError = await CompleteCurrentResultAsync().ConfigureAwait(false); + var suppressedError = _state.Current!.Error; + if (suppressedError is null && completeError is { } completionError) + suppressedError = completionError.Error; + if (suppressedError is not null) { - _enumeratorCurrent = null; + var exception = PgErrorException.Create(suppressedError); + var cold = GetOrCreateColdState(); + Interlocked.CompareExchange(ref cold.TerminalException, exception, null); + (cold.DrainErrors ??= new()).Add(exception); + if (completeError is { TransactionStatus: TransactionStatus.Unknown }) + await SkipDiscardedCommandsAsync().ConfigureAwait(false); + _state.CommandIndex++; + _state.Current = null; + Interlocked.Exchange(ref _state.Phase, PhaseDraining); + NotifyDrainStarted(); + _state.ConsumerDetached = true; + await DrainAsync().ConfigureAwait(false); + throw exception; } - else - { - if (Volatile.Read(ref _cancellationState) is { } cancellation) - cancellation.CallerToken = default; - if (!ReferenceEquals(_enumeratorCurrent, next)) - _enumeratorCurrent = next; + _state.CommandIndex++; + } - } + if (_state.Current is null) + throw ThrowHelper.ThrowInvalidOperation("The flow contains no commands."); + await CompleteBatchAsync().ConfigureAwait(false); + _state.ConsumerObservedCompletion = true; + return null; + } - // Close is durable across generations; complete the current one without publishing a result. - if (_callerInteractionCore.CloseException is not null) - { - CompleteEnumerationWithClose(_callerInteractionCore.CloseException); - return; - } - if (completed) + CommandResult? ReadNextPublishedResult() + { + CommandResult? result = _state.Current; + while (_state.CommandIndex < _state.Commands.Count) + { + result = ReadResult(_state.CommandIndex); + _state.Current = result; + _state.CurrentPublished = false; + if (!_state.Commands.ItemRef(_state.CommandIndex).SuppressEnumeration) + return result; + + var completeError = CompleteCurrentResult(); + var suppressedError = result.Error; + if (suppressedError is null && completeError is { } completionError) + suppressedError = completionError.Error; + if (suppressedError is not null) { - // Publish durable terminal state and complete the current generation atomically with - // respect to consumer rearming. Completion dispatches asynchronously, so it cannot reenter - // this lock or the pipeline frame that still owns the shared promise. - using (_rearmLock.EnterScope()) - { - PublishEnumerationCompleted(); - CompleteEnumeration(); - } - return; + var exception = PgErrorException.Create(suppressedError); + var cold = GetOrCreateColdState(); + Interlocked.CompareExchange(ref cold.TerminalException, exception, null); + (cold.DrainErrors ??= new()).Add(exception); + if (completeError is { TransactionStatus: TransactionStatus.Unknown }) + SkipDiscardedCommands(); + _state.CommandIndex++; + _state.Current = null; + Interlocked.Exchange(ref _state.Phase, PhaseDraining); + NotifyDrainStarted(); + _state.ConsumerDetached = true; + Drain(); + throw exception; } - _enumeratorMoveNextTaskSource.SetResult(true, runContinuationsAsynchronously: true); + + _state.CommandIndex++; } - async ValueTask ReadRfqAsync(PgDecoder decoder) + if (result is null) + throw ThrowHelper.ThrowInvalidOperation("The flow contains no commands."); + CompleteBatch(); + _state.ConsumerObservedCompletion = true; + return null; + } + + // Dispatch before entering an async machine so each mutually-exclusive protocol shape carries + // only its own awaiter and scratch state. + ValueTask ReadResultAsync() + { + // After close, a fresh command must not consume bytes left by its predecessor. + if (_state.Context.IsProtocolClosed) + return ValueTask.FromException(_state.Context.FlowTerminationException); + ref readonly var command = ref _state.Commands.ItemRef(_state.CommandIndex); + _state.Context.Decoder.UseReadTimeout(command.Timeout); + if (command.DescribeForPreparation) + return ReadPreparationResultAsync(); + return command.Descriptor is { IsPrepared: true, PreparedRowDescription: not null } + && !command.DescribeOnly + ? ReadPreparedResultAsync() + : ReadUnpreparedResultAsync(); + } + + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] + async ValueTask ReadPreparationResultAsync() + { + ref readonly var command = ref _state.Commands.ItemRef(_state.CommandIndex); + var preparation = await command.ReadPreparationDescriptionAsync( + _state.Context.Decoder, + _state.Context.GetProtocolStatic().RowDescription) + .ConfigureAwait(false); + return InitializeResult( + _state.CommandIndex, preparation.Item1, preparation.Item3, preparation.Item2); + } + + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] + async ValueTask ReadPreparedResultAsync() + { + // Prepared commands with a known description have the compact BindComplete -> + // DataRow/CommandComplete prelude. Await the decoder directly so a read wake resumes this + // frame rather than a nested parser coroutine. + if (!_state.Context.Decoder.TryMoveNext()) { - var message = await decoder.GetNextAsync().ConfigureAwait(false); - if (message.EnsureExpectedOrError(PgTypes.BackendType.ReadyForQuery) is { } rfqError) - PgErrorException.Throw(rfqError); + if (!await _state.Context.Decoder.MoveNextAsync().ConfigureAwait(false)) + _state.Context.Decoder.ThrowUnexpectedEof(); } - - static void ReadRfq(PgDecoder decoder) + var message = _state.Context.Decoder.Current; + PgError? error; + if (message.EnsureExpectedOrError(PgTypes.BackendType.BindComplete) is { } bindError) { - var message = decoder.GetNext(); - if (message.EnsureExpectedOrError(PgTypes.BackendType.ReadyForQuery) is { } rfqError) - PgErrorException.Throw(rfqError); + error = bindError; } + else + { + if (!_state.Context.Decoder.TryMoveNext()) + { + if (!await _state.Context.Decoder.MoveNextAsync().ConfigureAwait(false)) + _state.Context.Decoder.ThrowUnexpectedEof(); + } + _state.Context.Decoder.Current.DebugEnsureExpected( + PgTypes.BackendType.DataRow, PgTypes.BackendType.CommandComplete); + error = null; + } + return InitializeResult(_state.CommandIndex, error, null); + } + + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] + async ValueTask ReadUnpreparedResultAsync() + { + ref readonly var command = ref _state.Commands.ItemRef(_state.CommandIndex); + var result = await command.ReadUntilExecuteAsync( + _state.Context.Decoder, + _state.Context.GetProtocolStatic().RowDescription) + .ConfigureAwait(false); + return InitializeResult(_state.CommandIndex, result.Item1, result.Item2); } - void RegisterCancellationCallbacks(CancellationState cancellation) + CommandResult ReadResult(int commandIndex) { - if (cancellation.CallerToken.CanBeCanceled) + var context = _state.Context; + var decoder = context.Decoder; + if (context.IsProtocolClosed) + throw context.FlowTerminationException; + ref readonly var command = ref _state.Commands.ItemRef(commandIndex); + decoder.UseReadTimeout(command.Timeout); + PgError? error; + RowDescription? requestedRowDescription; + ParameterTypeList? preparationParameterTypes = null; + if (command.DescribeForPreparation) + { + var preparation = command.ReadPreparationDescription( + decoder, context.GetProtocolStatic().RowDescription); + error = preparation.Item1; + preparationParameterTypes = preparation.Item2; + requestedRowDescription = preparation.Item3; + } + else { - Debug.Assert(IsAsync); - cancellation.CallerRegistration = cancellation.CallerToken.UnsafeRegister(static (state, token) - => ((CommandFlow)state!).RequestCancelAndWake(token, CancellationScope.CurrentWindow), this); + (error, requestedRowDescription) = command.ReadUntilExecute( + decoder, context.GetProtocolStatic().RowDescription); } - if (cancellation.FlowToken.CanBeCanceled) + return InitializeResult( + commandIndex, error, requestedRowDescription, preparationParameterTypes); + } + + CommandResult InitializeResult( + int commandIndex, PgError? error, RowDescription? requestedRowDescription, + ParameterTypeList? preparationParameterTypes = null) + { + var context = _state.Context; + ref readonly var readState = ref context.GetProtocolStatic(); + ref readonly var command = ref _state.Commands.ItemRef(commandIndex); + readState.ResultMessageEnumerator.Initialize(command, context.Decoder); + var result = readState.CommandResult; + var descriptor = command.Descriptor; + // A named unprepared statement that parsed becomes a prepared descriptor. + if (!descriptor.IsPrepared && !descriptor.CommandName.IsDefault + && (error is not { } err || !err.Expected.Contains(PgTypes.BackendType.ParseComplete))) { - cancellation.FlowRegistration = cancellation.FlowToken.UnsafeRegister(static (state, token) - => ((CommandFlow)state!).RequestCancelAndWake(token, CancellationScope.RemainingFlow), this); + descriptor = CommandDescriptor.CreatePrepared(descriptor.CommandName, + preparationParameterTypes ?? descriptor.ParameterTypes, + requestedRowDescription?.Preserve()); } + result.Initialize(_ops.Flow, commandIndex, descriptor, requestedRowDescription, + !command.DescribeOnly, command.IsSimple(), error); + _ops.OnCommandResult(result); + return result; + } + + async ValueTask<(PgError Error, TransactionStatus TransactionStatus)?> CompleteCurrentResultAsync() + { + await _state.Context.GetProtocolStatic() + .ResultMessageEnumerator.DisposeAsync().ConfigureAwait(false); + return _state.Context.GetProtocolStatic() + .ResultMessageEnumerator.CompleteError; + } + + (PgError Error, TransactionStatus TransactionStatus)? CompleteCurrentResult() + { + var enumerator = _state.Context.GetProtocolStatic().ResultMessageEnumerator; + enumerator.Dispose(); + return enumerator.CompleteError; + } + + async ValueTask SkipDiscardedCommandsAsync() + { + while (++_state.CommandIndex < _state.Commands.Count && !_state.Commands[_state.CommandIndex].WithSync) { } + await ReadRfqAsync().ConfigureAwait(false); + if (_state.CommandIndex == _state.Commands.Count) + _state.ReadFlowRfq = false; + } + + void SkipDiscardedCommands() + { + while (++_state.CommandIndex < _state.Commands.Count && !_state.Commands[_state.CommandIndex].WithSync) { } + ReadRfq(); + if (_state.CommandIndex == _state.Commands.Count) + _state.ReadFlowRfq = false; } + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] - async ValueTask DisposeCancellationRegistrations(CancellationState cancellation) + async ValueTask ReadRfqAsync() { - if (cancellation.CallerToken.CanBeCanceled) - await cancellation.CallerRegistration.DisposeAsync().ConfigureAwait(false); - if (cancellation.FlowToken.CanBeCanceled) - await cancellation.FlowRegistration.DisposeAsync().ConfigureAwait(false); + var message = await _state.Context.Decoder.GetNextAsync().ConfigureAwait(false); + if (message.EnsureExpectedOrError(PgTypes.BackendType.ReadyForQuery) is { } error) + PgErrorException.Throw(error); } - bool IsCancellationToken(CancellationToken token) + void ReadRfq() { - var cancellation = Volatile.Read(ref _cancellationState); - return cancellation is not null - && (token == cancellation.CallerToken || token == cancellation.FlowToken); + var message = _state.Context.Decoder.GetNext(); + if (message.EnsureExpectedOrError(PgTypes.BackendType.ReadyForQuery) is { } rfqError) + PgErrorException.Throw(rfqError); } - // Cancellation callbacks only latch intent and wake the body. The body delivers cancellation after - // it has restored the wire boundary and is ready to release execution tenure. - internal Task CancelAsync() + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] + async ValueTask CompleteBatchAsync() { - var delivery = GetOrCreateCancelDelivery(); - if (IsCompleted) - { - delivery.TrySetResult(); - return delivery.Task; - } - RequestCancelAndWake(default, CancellationScope.RemainingFlow); - if (IsCompleted) - delivery.TrySetResult(); - return delivery.Task; + if (_state.ReadFlowRfq) + await ReadRfqAsync().ConfigureAwait(false); + await DisposeRegistrationsAsync().ConfigureAwait(false); + Finish(); } - CancellationState GetOrCreateCancellationState() + void CompleteBatch() { - if (Volatile.Read(ref _cancellationState) is { } cancellation) - return cancellation; - var created = new CancellationState(); - return Interlocked.CompareExchange(ref _cancellationState, created, null) ?? created; + if (_state.ReadFlowRfq) + ReadRfq(); + DisposeRegistrations(); + Finish(); } - TaskCompletionSource GetOrCreateCancelDelivery() + // The wire is at this command's RFQ. Release the shared read objects, record the outcome the + // consumer must observe, then complete the pipeline task. + void Finish() { - var cancellation = GetOrCreateCancellationState(); - var delivery = Volatile.Read(ref cancellation.Delivery); - if (delivery is not null) - return delivery; - var created = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously); - return Interlocked.CompareExchange(ref cancellation.Delivery, created, null) ?? created; + _state.Context.GetProtocolStatic().Reset(); + _state.Current = null; + _state.CurrentPublished = false; + if (IsCancelRequested) + Interlocked.CompareExchange(ref GetOrCreateColdState().TerminalException, + new OperationCanceledException(_state.ColdState!.DeliverToken), null); + else if (Volatile.Read(ref _state.ColdState)?.CloseException is { } close) + Interlocked.CompareExchange(ref GetOrCreateColdState().TerminalException, close, null); + CompletePipelineTask(null); } - bool RequestCancel(CancellationToken token, CancellationScope scope, - BackendCancellationTiming timing = BackendCancellationTiming.AfterGrace, - BackendCancellationTiming subsequentTiming = BackendCancellationTiming.AfterGrace, - bool allowCompletedEnumeration = false) + bool IsOwnCancellation(PgError error) + => IsCancelRequested && error.SqlState == PgErrorCodes.QueryCanceled; + + // The owning frame failed. The read error is the pipeline task's failure, which the framework + // recovers or drains. Later consumer calls replay it. + void FaultFromOwner(Exception exception) { - if (IsEnumerationCompleted && !allowCompletedEnumeration) - return false; - var cancellation = GetOrCreateCancellationState(); - cancellation.DeliverToken = token; - var observedScope = Volatile.Read(ref cancellation.Scope); - while ((int)scope > observedScope) + Interlocked.CompareExchange(ref GetOrCreateColdState().TerminalException, exception, null); + if (Volatile.Read(ref _state.Phase) == PhaseCompleted) + return; + DisposeRegistrations(); + _state.Current = null; + _state.CurrentPublished = false; + if (HasDecoder) + _state.Context.GetProtocolStatic().Reset(); + CompletePipelineTask(exception); + } + + // Takes an idle decoder for a drain. Returns false when a frame already owns it, which will observe + // the latch itself, or when the flow reached its terminal. + bool TryTakeOverDrain() + { + while (true) { - var priorScope = Interlocked.CompareExchange(ref cancellation.Scope, (int)scope, observedScope); - if (priorScope == observedScope) - break; - observedScope = priorScope; + var phase = Volatile.Read(ref _state.Phase); + if (phase is PhaseInitial && !HasDecoder) + return false; + if (phase is not (PhaseInitial or PhaseResultReady)) + return false; + if (Interlocked.CompareExchange(ref _state.Phase, PhaseDraining, phase) != phase) + continue; + NotifyDrainStarted(); + // Decoder takeover does not imply consumer abandonment. Explicit cancellation and + // graceful close also drain autonomously while retaining their consumer semantics. + Slon.Threading.SchedulingContext.SubmitDetached( + static state => + _ = new CommandFlowCore(TOps.Create((PgClientFlow)state!)).DrainAsync(), + _ops.Flow, + preferLocal: false); + return true; } - Volatile.Write(ref cancellation.Requested, true); - Volatile.Write(ref _draining, true); - var observedTiming = Volatile.Read(ref cancellation.Timing); - while ((int)timing > observedTiming) + } + + void NotifyDrainStarted() + { + if (Interlocked.Exchange(ref _state.DrainStarted, 1) is 0) + _ops.OnDrainStarted(); + } + + // Autonomous drain. Owns the decoder until the pipeline task completes. Never throws. + async ValueTask DrainAsync() + { + try { - var priorTiming = Interlocked.CompareExchange(ref cancellation.Timing, (int)timing, observedTiming); - if (priorTiming == observedTiming) - break; - observedTiming = priorTiming; + if (_state.Current is null) + { + await new ValueTask((IValueTaskSource)_ops.Flow, _state.ReadySource.Version).ConfigureAwait(false); + if (_state.CommandIndex < 0) + _state.CommandIndex = 0; + if (_state.CommandIndex >= _state.Commands.Count) + { + await CompleteBatchAsync().ConfigureAwait(false); + return; + } + SetCurrent(await ReadResultAsync().ConfigureAwait(false)); + } + + while (true) + { + var completeError = await CompleteCurrentResultAsync().ConfigureAwait(false); + CaptureDrainError(_state.Current!, completeError); + _state.CurrentPublished = false; + if (completeError is { TransactionStatus: TransactionStatus.Unknown }) + await SkipDiscardedCommandsAsync().ConfigureAwait(false); + if (++_state.CommandIndex >= _state.Commands.Count) + break; + SetCurrent(await ReadResultAsync().ConfigureAwait(false)); + } + await CompleteBatchAsync().ConfigureAwait(false); } - var observedSubsequentTiming = Volatile.Read(ref cancellation.SubsequentTiming); - while ((int)subsequentTiming > observedSubsequentTiming) + catch (TimeoutException ex) { - var priorTiming = Interlocked.CompareExchange(ref cancellation.SubsequentTiming, - (int)subsequentTiming, observedSubsequentTiming); - if (priorTiming == observedSubsequentTiming) - break; - observedSubsequentTiming = priorTiming; + // A timeout during semantic drain escalates the same cancellation episode immediately. + // The pipeline failure then hands any remaining wire obligation to recovery. + RequestCancel(default, CommandExecutionCancellationScope.RemainingFlow, + BackendCancellationTiming.Immediate, BackendCancellationTiming.AtReadFrontier); + FaultFromOwner(ex); + } + catch (Exception ex) + { + FaultFromOwner(ex); } - var delivery = Volatile.Read(ref cancellation.Delivery); - RequestBackendCancellation(timing, delivery); - return true; } - void RequestCancelAndWake(CancellationToken token, CancellationScope scope) + // A read timeout is terminal for the consumer but not for the wire obligation. This frame has + // released the decoder read tenure, so transfer ownership to an autonomous semantic drain while + // the original MoveNext returns its timeout immediately. + void HandleReadTimeout(TimeoutException exception) { - if (!RequestCancel(token, scope)) - return; - var delivery = Volatile.Read(ref _cancellationState) is { } cancellation - ? Volatile.Read(ref cancellation.Delivery) - : null; - _callerInteractionCore.ResumeBody(runContinuationsAsynchronously: true); - _callerInteractionCore.WakeBody(useDedicatedDriver: !IsAsync && delivery is not null); + Interlocked.CompareExchange( + ref GetOrCreateColdState().TerminalException, exception, null); + _state.ConsumerDetached = true; + NotifyDrainStarted(); + Interlocked.Exchange(ref _state.Phase, PhaseDraining); + RequestCancel(default, CommandExecutionCancellationScope.RemainingFlow, + BackendCancellationTiming.Immediate, BackendCancellationTiming.AtReadFrontier); + Slon.Threading.SchedulingContext.SubmitDetached( + static state => + _ = new CommandFlowCore(TOps.Create((PgClientFlow)state!)).DrainAsync(), + _ops.Flow, + preferLocal: false); } - void RequestBackendCancellation(BackendCancellationTiming timing = BackendCancellationTiming.AfterGrace, - TaskCompletionSource? delivery = null) + void Drain() { - // Cancellation which wins before body entry remains latched in the cancellation state and is - // replayed immediately after the execution context is published. - if (Volatile.Read(ref _contextPublished) && Volatile.Read(ref _cancellationState) is { } cancellation) + try { - var key = Volatile.Read(ref cancellation.EpisodeKey); - if (key is null) + var result = _state.Current; + if (result is null) + { + // Synchronous disposal before any read. The activation bridge normally completed long + // ago, so this bridge is rarely more than a status check. + new ValueTask((IValueTaskSource)_ops.Flow, _state.ReadySource.Version).AsTask().GetAwaiter().GetResult(); + if (_state.CommandIndex < 0) + _state.CommandIndex = 0; + if (_state.CommandIndex >= _state.Commands.Count) + { + CompleteBatch(); + return; + } + result = ReadResult(_state.CommandIndex); + } + + while (true) { - var created = new object(); - key = Interlocked.CompareExchange(ref cancellation.EpisodeKey, created, null) ?? created; + var completeError = CompleteCurrentResult(); + CaptureDrainError(result, completeError); + _state.CurrentPublished = false; + if (completeError is { TransactionStatus: TransactionStatus.Unknown }) + SkipDiscardedCommands(); + if (++_state.CommandIndex >= _state.Commands.Count) + break; + result = ReadResult(_state.CommandIndex); } - _context.RequestBackendCancellation(this, CancellationWindow, timing, delivery, - key, Volatile.Read(ref cancellation.Scope), - (BackendCancellationTiming)Volatile.Read(ref cancellation.SubsequentTiming)); + CompleteBatch(); + } + catch (Exception ex) + { + FaultFromOwner(ex); } } - protected override void OnCancellationWindowCompleted(int completedWindow, int remainingWindowCount) - { } + internal ValueTask DisposeAsync() + { + while (true) + { + var phase = Volatile.Read(ref _state.Phase); + switch (phase) + { + case PhaseInitial: + case PhaseResultReady: + if (Interlocked.CompareExchange(ref _state.Phase, PhaseDraining, phase) != phase) + continue; + NotifyDrainStarted(); + _state.ConsumerDetached = true; + if (_state.Current is { IsComplete: false } + || _state.CommandIndex + 1 < _state.Commands.Count) + RequestConsumerDrainCancellation(); + return _state.WaitForDrainOnDispose ? DisposeDrainAsync() : FireAndForgetDrain(); + case PhaseReading: + _state.ConsumerDetached = true; + NotifyDrainStarted(); + RequestConsumerDrainCancellation(); + return _state.WaitForDrainOnDispose ? DisposeCompletedAsync() : default; + default: + return !_state.WaitForDrainOnDispose || _state.ConsumerObservedCompletion + ? default + : DisposeCompletedAsync(); + } + } + } - bool IsOwnCancellation(PgError error) + // Drains on the disposer's frame, then waits for framework release so a drain error can surface. + async ValueTask DisposeDrainAsync() { - if (Volatile.Read(ref _cancellationState) is not { } cancellation - || !Volatile.Read(ref cancellation.Requested) || error.SqlState != PgErrorCodes.QueryCanceled) - return false; - // PgDecoder records each ErrorResponse arrival exactly once. Classification may inspect the - // preserved error through several command-result paths, so it must not report the same strike. - return true; + await DrainAsync().ConfigureAwait(false); + await DisposeCompletedAsync().ConfigureAwait(false); } - void EnterStoppingDrainIfNeeded(Context context) + ValueTask FireAndForgetDrain() { - if (_callerInteractionCore.CloseException is { } close && context.StoppingToken.IsCancellationRequested - && !IsDraining && !IsEnumerationCompleted) + _ = DrainAsync(); + return default; + } + + async ValueTask DisposeCompletedAsync() + { + await WaitForCompletionAsync().ConfigureAwait(false); + ThrowDrainErrors(); + } + + // Flow completion is independent of errors accumulated while draining. A close is a clean + // terminal for a disposing consumer. + async ValueTask WaitForCompletionAsync() + { + try + { + await _ops.Flow.WaitForComplete().ConfigureAwait(false); + } + catch (PgClientClosedException) + { + } + } + + internal void Dispose() + { + if (!_ops.IsAsyncAtDispatch) + EnsureSyncHandoff(); + while (true) { - CompleteEnumerationWithException(close); - MarkBodyInitiatedDrain(); + var phase = Volatile.Read(ref _state.Phase); + switch (phase) + { + case PhaseInitial: + case PhaseResultReady: + if (Interlocked.CompareExchange(ref _state.Phase, PhaseDraining, phase) != phase) + continue; + NotifyDrainStarted(); + _state.ConsumerDetached = true; + if (_state.Current is { IsComplete: false } + || _state.CommandIndex + 1 < _state.Commands.Count) + RequestConsumerDrainCancellation(); + Drain(); + if (_state.WaitForDrainOnDispose) + DisposeCompleted(); + return; + case PhaseReading: + _state.ConsumerDetached = true; + NotifyDrainStarted(); + RequestConsumerDrainCancellation(); + if (_state.WaitForDrainOnDispose) + DisposeCompleted(); + return; + default: + if (_state.WaitForDrainOnDispose && !_state.ConsumerObservedCompletion) + DisposeCompleted(); + return; + } } } - // Publish progress unconditionally. A terminal delivery can lose its task-source CAS to an already- - // completed generation; a synchronous disposer must still observe this sticky level rather than park. - void SignalPumpProgress() - => _callerInteractionCore.SignalProgress(); + void DisposeCompleted() + { + try + { + _ops.Flow.WaitForCompleteSynchronously(); + } + catch (PgClientClosedException) + { + } + ThrowDrainErrors(); + } - void CompleteEnumerationWithException(Exception ex) + void CaptureDrainError(CommandResult result, + (PgError Error, TransactionStatus TransactionStatus)? completeError) { - // Close state must survive task-source rearming, including flows whose body never started. - if (ex is PgClientClosedException or PgCollateralException) - _callerInteractionCore.SetCloseLatch(ex); - if (IsEnumerationCompleted) + if (!_state.ConsumerDetached || _state.CurrentPublished) + return; + var error = result.Error ?? completeError?.Error; + if (error is null || IsOwnCancellation(error)) return; - // Teardown may race the consumer. The task source is the completion authority; - // _enumeratorCompleted follows only when this call wins the current generation. - if (_enumeratorMoveNextTaskSource.TrySetException(ex, runContinuationsAsynchronously: true)) - PublishEnumerationCompleted(); - // A faulted body will not publish another continuation. - SignalPumpProgress(); - // Wire recovery remains with the body or the framework recovery flow; this method only completes - // the consumer-facing generation. + var cold = GetOrCreateColdState(); + (cold.DrainErrors ??= new()).Add(PgErrorException.Create(error)); } - void PublishBodyTerminated() + void ThrowDrainErrors() { - Volatile.Write(ref _bodyState, BodyTerminated); - SignalPumpProgress(); + if (Volatile.Read(ref _state.ColdState)?.DrainErrors is not { Count: > 0 } errors) + return; + if (errors.Count is 1) + ExceptionDispatchInfo.Throw(errors[0]); + throw new AggregateException(errors); } - bool TerminateBodyBeforeStart() - => Interlocked.CompareExchange(ref _bodyState, BodyTerminated, BodyNotStarted) == BodyNotStarted; + // Give the current window a chance to finish naturally. Once disposal advances into an unread + // successor, dispatch at its read frontier instead of paying the grace period for every window. + void RequestConsumerDrainCancellation() + => RequestCancel(default, CommandExecutionCancellationScope.RemainingFlow, + BackendCancellationTiming.AfterGrace, + BackendCancellationTiming.AtReadFrontier); - bool IsBodyRunning => Volatile.Read(ref _bodyState) == BodyRunning; - bool IsBodyTerminated => Volatile.Read(ref _bodyState) == BodyTerminated; + void RegisterCancellation(CancellationToken callerToken) + { + // Keep the second token/registration pair off ordinary flow objects. Default-token traffic + // does not need cancellation state at all and remains the allocation/footprint hot path. + var cancellation = Volatile.Read(ref _state.ColdState); + if (callerToken.CanBeCanceled || cancellation is not null) + { + cancellation ??= GetOrCreateColdState(); + if (callerToken != cancellation.CallerToken) + { + var callerRegistration = cancellation.CallerRegistration; + cancellation.CallerRegistration = default; + callerRegistration.Dispose(); + cancellation.CallerToken = callerToken; + } + if (callerToken.CanBeCanceled && cancellation.CallerRegistration == default) + cancellation.CallerRegistration = callerToken.UnsafeRegister(static (state, token) + => new CommandFlowCore(TOps.Create((PgClientFlow)state!)).RequestCancel( + token, CommandExecutionCancellationScope.CurrentWindow), _ops.Flow); + } - // Source handoff finishes before body/consumer rendezvous begins, so both reuse the same wait event. - private protected override FlowHandoffEvent? HandoffEvent => _callerInteractionCore.GetWaitEvent(); + if (_state.FlowToken.CanBeCanceled && _state.FlowRegistration == default) + _state.FlowRegistration = _state.FlowToken.UnsafeRegister(static (state, token) + => new CommandFlowCore(TOps.Create((PgClientFlow)state!)).RequestCancel( + token, CommandExecutionCancellationScope.RemainingFlow), _ops.Flow); + } - // Return the rendezvous directly. An async wrapper could signal the disposer before registering the - // body's continuation, causing a late ThreadPool dispatch instead of caller-thread drain execution. - FlowCallerInteractionCore.CallerHandoffAwaitable YieldToCaller() + ValueTask DisposeRegistrationsAsync() { - FieldRef> fieldRef; - unsafe + var cancellation = Volatile.Read(ref _state.ColdState); + var callerRegistration = cancellation?.CallerRegistration ?? default; + if (callerRegistration == default && _state.FlowRegistration == default) + return default; + if (cancellation is not null) + cancellation.CallerRegistration = default; + var flowRegistration = _state.FlowRegistration; + _state.FlowRegistration = default; + return DisposeRegistrationsAsync(callerRegistration, flowRegistration); + + static async ValueTask DisposeRegistrationsAsync( + CancellationTokenRegistration callerRegistration, + CancellationTokenRegistration flowRegistration) { - fieldRef = FieldRef>.Create(&GetCallerInteractionCore, this); + await callerRegistration.DisposeAsync().ConfigureAwait(false); + await flowRegistration.DisposeAsync().ConfigureAwait(false); } - return _callerInteractionCore.YieldToCaller(fieldRef); } - static ref FlowCallerInteractionCore GetCallerInteractionCore(CommandFlow instance) - => ref instance._callerInteractionCore; + void DisposeRegistrations() + { + var cancellation = Volatile.Read(ref _state.ColdState); + var callerRegistration = cancellation?.CallerRegistration ?? default; + if (cancellation is not null) + cancellation.CallerRegistration = default; + var flowRegistration = _state.FlowRegistration; + _state.FlowRegistration = default; + callerRegistration.Dispose(); + flowRegistration.Dispose(); + } - protected override void OnAbort(Exception exception) => FaultCaller(exception); + // Cancellation only latches intent and requests a backend cancel. The frame owning the decoder + // delivers it after the wire is back at RFQ. An idle flow drains autonomously first. + void RequestCancel(CancellationToken token, CommandExecutionCancellationScope scope, + BackendCancellationTiming timing = BackendCancellationTiming.AfterGrace, + BackendCancellationTiming subsequentTiming = BackendCancellationTiming.AfterGrace) + { + if (Volatile.Read(ref _state.Phase) == PhaseCompleted) + return; + var cancellation = GetOrCreateColdState(); + cancellation.DeliverToken = token; + RaiseCancellationScope(cancellation, scope); + RaiseCancellationTiming(ref cancellation.Timing, timing); + RaiseCancellationTiming(ref cancellation.SubsequentTiming, subsequentTiming); + Interlocked.Exchange(ref cancellation.CancelRequested, true); + if (HasDecoder) + RequestBackendCancellation(); + TryTakeOverDrain(); + } - // Graceful stopping is the early wire-close wake and is idempotent across heartbeat ticks. - protected override void OnStopping(Exception exception) + static void RaiseCancellationScope(CommandExecutionColdState cancellation, CommandExecutionCancellationScope scope) { - if (!IsBodyRunning || !IsAsync) + var requested = (int)scope; + var current = Volatile.Read(ref cancellation.Scope); + while (current < requested) { - FaultCaller(exception); - return; + var observed = Interlocked.CompareExchange( + ref cancellation.Scope, requested, current); + if (observed == current) + return; + current = observed; } - - // Resume normally so the body observes the close latch and drains; abort faults the gate. - _callerInteractionCore.SetCloseLatch(exception); - _callerInteractionCore.ResumeBody(runContinuationsAsynchronously: true); } - // Wake a running body so it owns fault delivery; directly fault a flow whose body never started. - void FaultCaller(Exception exception) + static void RaiseCancellationTiming( + ref int location, BackendCancellationTiming timing) { - if (TerminateBodyBeforeStart()) + var requested = (int)timing; + var current = Volatile.Read(ref location); + while (current < requested) { - CompleteEnumerationWithException(exception); - // A synchronous flow may already have entered ExecuteAfterHandoff while its inner read - // body is still NotStarted. Terminating that body does not complete the outer execution: - // fault and wake its initial handoff so the framework task can settle and release tenure. - _callerInteractionCore.FaultBodyWait(exception); - _callerInteractionCore.WakeBody(); - return; + var observed = Interlocked.CompareExchange(ref location, requested, current); + if (observed == current) + return; + current = observed; } - - // A concurrent body start may have beaten the pre-start terminal claim. - if (IsBodyRunning) - _callerInteractionCore.FaultBodyWait(exception); - else - CompleteEnumerationWithException(exception); } - internal void Fail(Exception exception) => FaultCaller(exception); + internal Task CancelAsync() + { + var cancellation = GetOrCreateColdState(); + var delivery = Volatile.Read(ref cancellation.Delivery) + ?? Interlocked.CompareExchange(ref cancellation.Delivery, + new(TaskCreationOptions.RunContinuationsAsynchronously), null) + ?? cancellation.Delivery; + if (Volatile.Read(ref _state.Phase) is PhaseCompleted) + { + delivery.TrySetResult(); + return delivery.Task; + } + RequestCancel(default, CommandExecutionCancellationScope.RemainingFlow); + if (Volatile.Read(ref _state.Phase) is PhaseCompleted) + delivery.TrySetResult(); + return delivery.Task; + } - protected override void OnReleasing(Exception? exception) + void RequestBackendCancellation() { - if (Volatile.Read(ref _cancellationState) is { } cancellation) - Volatile.Read(ref cancellation.Delivery)?.TrySetResult(); - _commands.Return(); + var cancellation = GetOrCreateColdState(); + var episodeKey = Volatile.Read(ref cancellation.EpisodeKey); + if (episodeKey is null) + { + var created = new object(); + episodeKey = Interlocked.CompareExchange( + ref cancellation.EpisodeKey, created, null) ?? created; + } + _state.Context.RequestBackendCancellation( + _ops.Flow, _ops.Flow.CancellationWindow, + (BackendCancellationTiming)Volatile.Read(ref cancellation.Timing), + Volatile.Read(ref cancellation.Delivery), + episodeKey, + Math.Max(Volatile.Read(ref cancellation.Scope), (int)CommandExecutionCancellationScope.CurrentWindow), + (BackendCancellationTiming)Volatile.Read(ref cancellation.SubsequentTiming)); } - protected override void OnDiscarded() + // Graceful stop. An unactivated flow releases its consumer, the closing wire owns its response. + // An idle activated flow drains itself to RFQ so the pipeline can complete. + internal void OnStopping(Exception exception) { - // Discarded flows never enter the base release path. - GetObserver(out var observerState)?.OnCompleting(this, null, observerState); - _commands.Return(); + Interlocked.CompareExchange(ref GetOrCreateColdState().CloseException, exception, null); + if (CompleteReady(exception, runContinuationsAsynchronously: true)) + { + Interlocked.CompareExchange(ref GetOrCreateColdState().TerminalException, exception, null); + return; + } + TryTakeOverDrain(); } - protected override void OnReset() + // Forceful abort. No frame can read a dead wire, so an idle owner faults the pipeline task + // directly. A frame in flight fails on its own read. + internal void OnAbort(Exception exception) { - Debug.Assert(IsPending || IsCompleted); - _commandIndex = -1; - _executePipelinedCore.Reset(); - _enumeratorMoveNextTaskSource.Reset(); - // Disarm while idle in the pool (no teardown can target a non-live flow). Initialize re-arms it - // before the flow is queued, so a live flow is always in concurrent-completion mode. - _enumeratorMoveNextTaskSource.CanCompleteConcurrently = false; - _enumeratorCurrent = default; - _enumeratorCompleted = false; - _isResultReady = false; - _callerInteractionCore.Reset(); - if (_cancellationState is { } cancellation) + Interlocked.CompareExchange(ref GetOrCreateColdState().CloseException, exception, null); + if (CompleteReady(exception, runContinuationsAsynchronously: true)) { - cancellation.Reset(); - _cancellationState = null; + Interlocked.CompareExchange(ref GetOrCreateColdState().TerminalException, exception, null); + return; } - _drainErrors = null; - _consumeNonQuery = false; - _nonQueryRecordsAffected = 0; - _consumerDisposed = false; - _draining = false; - _drainModeEntered = false; - WaitForDrainOnDispose = true; - // Dispatch state is per-tenure. - _pipelinePromise = null; - _contextPublished = false; - _context = default; - _task = default; - _bodyState = BodyNotStarted; - _consumerAdvanced = false; - } - - FlowCallerInteractionCoreResult IValueTaskSource.GetResult(short token) - => _callerInteractionCore.ConsumeGateResult(token); - - ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) - => _callerInteractionCore.GateStatus(token); - - void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) - { - _callerInteractionCore.OnGateCompleted(continuation, state, token, flags); - // Drain is a sticky level. Recheck it after registering so an earlier gate edge cannot be lost - // across gate reset. This is the body's suspending stack, so the wake must dispatch asynchronously. - if (IsDraining) + while (true) { - // A synchronous takeover would already have resumed the body inline. Reaching this callback - // means autonomous execution still owns the body. - IsAsync = true; - _callerInteractionCore.ResumeBody(runContinuationsAsynchronously: true); + var phase = Volatile.Read(ref _state.Phase); + if (phase is not (PhaseInitial or PhaseResultReady)) + return; + if (Interlocked.CompareExchange(ref _state.Phase, PhaseCompleted, phase) != phase) + continue; + Interlocked.CompareExchange(ref GetOrCreateColdState().TerminalException, exception, null); + _state.PipelineTaskSource.SetException(exception, runContinuationsAsynchronously: true); + return; } } - // Backing for the pipelined-dispatch ValueTask. Returned to the framework when activation - // hasn't fired yet. Nested callback completes it when ExecutePipelined finishes. - void IValueTaskSource.GetResult(short token) => _executePipelinedCore.GetResult(token); - ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _executePipelinedCore.GetStatus(token); + internal void Fail(Exception exception) + { + // A result callback failed on the frame that owns the decoder. Its throw propagates there. + Interlocked.CompareExchange(ref GetOrCreateColdState().TerminalException, exception, null); + } + + internal void OnReleasing(Exception? exception) + { + Volatile.Read(ref _state.ColdState)?.Delivery?.TrySetResult(); + DisposeRegistrations(); + _state.Commands.Return(); + } + + internal void OnDiscarded() + { + _ops.OnDiscarded(); + _state.Commands.Return(); + } + + internal void OnReset() + { + _state.Phase = PhaseInitial; + _state.CommandIndex = -1; + _state.Context = default; + _state.ContextPublished = false; + _state.Current = null; + _state.CurrentPublished = false; + _state.ReadFlowRfq = false; + _state.ConsumerDetached = false; + _state.ConsumerObservedCompletion = false; + _state.ReadySource.Reset(); + _state.PipelineTaskSource.Reset(); + _state.ReadyCompletion = 0; + _state.DrainStarted = 0; + _state.FlowToken = default; + _state.FlowRegistration = default; + _state.WindowToken = default; + _state.FirstPromise ??= new(); + _state.ColdState = null; + _state.SyncHandoffClaimed = false; + _state.HandoffEvent?.ResetInteraction(); + _state.WaitForDrainOnDispose = true; + } + +} + +public sealed partial class CommandFlow +{ + bool IValueTaskSource.GetResult(short token) => _state.ReadySource.GetResult(token); + ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _state.ReadySource.GetStatus(token); + void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) + => _state.ReadySource.OnCompleted(continuation, state, token, flags); + + void IValueTaskSource.GetResult(short token) => _state.PipelineTaskSource.GetResult(token); + ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _state.PipelineTaskSource.GetStatus(token); void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) - => _executePipelinedCore.OnCompleted(continuation, state, token, flags); + => _state.PipelineTaskSource.OnCompleted(continuation, state, token, flags); + public readonly struct Enumerator : IEnumerator, IAsyncEnumerator + { + readonly CommandFlow? _flow; + readonly CancellationToken _cancellationToken; + + public Enumerator(CommandFlow flow) + : this(flow, default) + { } + + internal Enumerator(CommandFlow flow, CancellationToken cancellationToken) + { + _flow = flow; + _cancellationToken = cancellationToken; + } + + public Enumerator GetAsyncEnumerator() => this; + + public Enumerator GetEnumerator() => this; + + public bool MoveNext() => _flow?.Core.MoveNext() ?? false; + + public ValueTask MoveNextAsync() => MoveNextAsync(_cancellationToken); + + public ValueTask MoveNextAsync(CancellationToken cancellationToken) + => _flow is null ? new(false) : _flow.Core.MoveNextAsync(cancellationToken); + + public CommandResult Current => _flow?._state.Current ?? default!; + + object? IEnumerator.Current => Current; + + void IEnumerator.Reset() => throw new NotSupportedException(); + + public ValueTask DisposeAsync() => _flow is null ? default : _flow.Core.DisposeAsync(); + + public void Dispose() => _flow?.Core.Dispose(); + } } diff --git a/Slon/Pg/Protocol/Flows/FlowCallerInteractionCore.cs b/Slon/Pg/Protocol/Flows/FlowCallerInteractionCore.cs index 8bf0042..1ee22a5 100644 --- a/Slon/Pg/Protocol/Flows/FlowCallerInteractionCore.cs +++ b/Slon/Pg/Protocol/Flows/FlowCallerInteractionCore.cs @@ -26,16 +26,17 @@ public Exception SetCloseLatch(Exception exception) // Inline completion transfers the body to a synchronous caller; asynchronous completion preserves // autonomous execution without running the body on the signaller's stack. Slon.Threading.Tasks.Sources.ManualResetValueTaskSourceCore _gate; - public void Initialize() - { - _gate.CanCompleteConcurrently = true; - } + int _gateCompletionClaim; + public void Initialize() => _gateCompletionClaim = 0; public ValueTask WaitForCaller(IValueTaskSource source) => new(source, _gate.Version); // The IVTS facade forwards through this gate surface. public void ResumeBody(bool runContinuationsAsynchronously) - => _gate.TrySetResult(default!, runContinuationsAsynchronously); + { + if (Interlocked.CompareExchange(ref _gateCompletionClaim, 1, 0) == 0) + _gate.SetResult(default!, runContinuationsAsynchronously); + } public System.Threading.Tasks.Sources.ValueTaskSourceStatus GateStatus(short token) => _gate.GetStatus(token); public void OnGateCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) => _gate.OnCompleted(continuation, state, token, flags); @@ -44,6 +45,7 @@ public TResult ConsumeGateResult(short token) { var result = _gate.GetResult(token); _gate.Reset(); + Volatile.Write(ref _gateCompletionClaim, 0); return result; } @@ -52,7 +54,8 @@ public void FaultBodyWait(Exception exception) // Latch first (monotone), then fault the gate, so a consumer observing the gate fault always // reads the latch set. var latched = SetCloseLatch(exception); - _gate.TrySetException(latched, runContinuationsAsynchronously: true); + if (Interlocked.CompareExchange(ref _gateCompletionClaim, 1, 0) == 0) + _gate.SetException(latched, runContinuationsAsynchronously: true); // The synchronous disposer may not have created its event yet. Publish a sticky progress // level so WaitForContinuation observes the close even when this Set would have been a no-op. SignalProgress(); @@ -144,7 +147,10 @@ static void Dispatch(Action continuation, bool useDedicatedDriver) { if (!useDedicatedDriver) { - ThreadPool.UnsafeQueueUserWorkItem(static s => ((Action)s!)(), (object)continuation); + Slon.Threading.SchedulingContext.SubmitDetached( + static state => ((Action)state!)(), + continuation, + preferLocal: false); return; } @@ -177,7 +183,9 @@ public void SignalProgress() bool ConsumeProgress() => Interlocked.Exchange(ref _progressSignaled, 0) != 0; - public CallerHandoffAwaitable YieldToCaller(FieldRef> fieldRef) + public CallerHandoffAwaitable YieldToCaller(TFieldRef fieldRef) + where TFieldRef : struct, + IFieldRef> => new(fieldRef); public void Reset() @@ -187,13 +195,16 @@ public void Reset() _wakeRequested = false; _closeException = null; _gate.Reset(); + Volatile.Write(ref _gateCompletionClaim, 0); } - public readonly struct CallerHandoffAwaitable(FieldRef> fieldRef) + public readonly struct CallerHandoffAwaitable(TFieldRef fieldRef) + where TFieldRef : struct, + IFieldRef> { public Awaiter GetAwaiter() => new(fieldRef); - public readonly struct Awaiter(FieldRef> fieldRef) : ICriticalNotifyCompletion + public readonly struct Awaiter(TFieldRef fieldRef) : ICriticalNotifyCompletion { public bool IsCompleted => false; @@ -202,14 +213,14 @@ public void GetResult() // Surface a pending cancellation set by FaultBodyWait. The gate source is // only completed when cancellation fires. In the normal sync-flow handoff path // it stays Pending and we just return. - var gate = fieldRef.Invoke()._gate; + var gate = fieldRef.GetField()._gate; if (gate.GetStatus(gate.Version) != System.Threading.Tasks.Sources.ValueTaskSourceStatus.Pending) gate.GetResult(gate.Version); } public void OnCompleted(Action continuation) { - ref var field = ref fieldRef.Invoke(); + ref var field = ref fieldRef.GetField(); var waitEvent = field.GetWaitEvent(); if (!ReferenceEquals(waitEvent.HandoffContinuation, continuation)) Volatile.Write(ref waitEvent.HandoffContinuation, continuation); diff --git a/Slon/Pg/Protocol/PgClientFlow.cs b/Slon/Pg/Protocol/PgClientFlow.cs index 5bcd53e..5018095 100644 --- a/Slon/Pg/Protocol/PgClientFlow.cs +++ b/Slon/Pg/Protocol/PgClientFlow.cs @@ -16,6 +16,9 @@ protected internal virtual void OnCompleted(PgClientFlow flow, Exception? except // them from the migratable source backlog, while standalone low-level flows remain initialized. abstract class PgClientFlowBindingContext; +readonly struct FlowCompletion; +readonly struct FlowActivation; + sealed class FlowHandoffEvent : ManualResetEventSlim { PgClientFlowSource.State? _placementSource; @@ -74,8 +77,14 @@ internal void ResetInteraction() } [Experimental(ExperimentalDiagnostics.PostgreSqlLowerLayer)] -public abstract class PgClientFlow : IValueTaskSource, IValueTaskSource, IThreadPoolWorkItem +public abstract class PgClientFlow : IValueTaskSource, IValueTaskSource, IThreadPoolWorkItem { + sealed class OptionalState + { + internal ManualResetEventSlim? CompletionEvent; + internal CancellationTokenRegistration ActivationRegistration; + } + PgClientProtocol.Control? _pendingActivationControl; FlowEnqueueOptions _enqueueOptions; bool _ownsWireCapacity; @@ -123,12 +132,19 @@ internal virtual void Bind(PgClientFlowBindingContext? context) { } /// Initialize overwrote the first's item before its Execute ran (one pending activation /// per flow tenure makes the field safe). internal void PrepareActivationDispatch(PgClientProtocol.Control control) - => _pendingActivationControl = control; + { + _activationWasDispatched = true; + _pendingActivationControl = control; + } void IThreadPoolWorkItem.Execute() { var control = _pendingActivationControl; - Debug.Assert(control is not null); + if (control is null) + { + ExecuteDetachedWorkItem(); + return; + } _pendingActivationControl = null; // The decoder bind already ran synchronously at activation; this dispatch is only the body // wake. Skip it for a flow the abort retired before the dispatch ran: its activation source is @@ -138,13 +154,17 @@ void IThreadPoolWorkItem.Execute() control!.Activate(this); } + private protected virtual void ExecuteDetachedWorkItem() + => ThrowHelper.ThrowInvalidOperation("The flow has no detached work item pending."); + readonly bool _supportsDeferredFlush; internal bool SupportsDeferredFlush => _supportsDeferredFlush; - Action? _decoderOnHeartbeatAction; // TODO should we have this here? int _rfqCount; int _cancellationWindow; internal int CancellationWindow => Volatile.Read(ref _cancellationWindow); bool _lastMessageInducesRfq; + bool _activationWasDispatched; + internal bool ActivationWasDispatched => _activationWasDispatched; // We store the IsAsync value at bind time so the protocol can keep track of pipeline stalls correctly. bool _isAsyncAtDispatch; // Tri-state int (0 = unset, 1 = true, 2 = false) instead of bool? so reads / writes can be @@ -166,15 +186,15 @@ void IThreadPoolWorkItem.Execute() // the flow types; the flow-as-result is the free disambiguator (the old Slon.Protocols // pattern). At most one pending waiter per tenure; post-completion awaits resolve // synchronously. - Slon.Threading.Tasks.Sources.ManualResetValueTaskSourceCore _completionCore; - ManualResetEventSlim? _completionEvent; + Slon.Threading.Tasks.Sources.ManualResetValueTaskSourceCore _completionCore; + OptionalState? _optionalState; // 1 while a WaitForComplete token is live (set at capture, cleared after GetResult consumed the // core). Guards reuse: Reset bumps the core's version, so it must not run while this is set. int _completionWaiterPending; // Activation state. - Slon.Threading.Tasks.Sources.ManualResetValueTaskSourceCore _activationTaskSource; - CancellationTokenRegistration _activationCancellationTokenRegistration; + Slon.Threading.Tasks.Sources.ManualResetValueTaskSourceCore _activationTaskSource; + int _activationClaim; TimeSpan _remainingActivationTimeout; bool _pendingTimeoutStarted; @@ -268,8 +288,55 @@ internal void WaitForSyncHandoff() protected PgClientFlow(bool supportsDeferredFlush = false) { _supportsDeferredFlush = supportsDeferredFlush; - _activationTaskSource.CanCompleteConcurrently = true; - _completionCore.CanCompleteConcurrently = true; + } + + OptionalState GetOrCreateOptionalState() + => Volatile.Read(ref _optionalState) + ?? Interlocked.CompareExchange(ref _optionalState, new(), null) + ?? _optionalState; + + CancellationTokenRegistration TakeActivationRegistration() + { + var state = Volatile.Read(ref _optionalState); + if (state is null) + return default; + var registration = state.ActivationRegistration; + state.ActivationRegistration = default; + return registration; + } + + bool TrySetActivationResult(bool runContinuationsAsynchronously) + { + if (Interlocked.CompareExchange(ref _activationClaim, 1, 0) != 0) + return false; + _activationTaskSource.SetResult(default, runContinuationsAsynchronously); + return true; + } + + bool TrySetActivationException(Exception exception, bool runContinuationsAsynchronously) + { + if (Interlocked.CompareExchange(ref _activationClaim, 1, 0) != 0) + return false; + _activationTaskSource.SetException(exception, runContinuationsAsynchronously); + return true; + } + + void ResetActivationSource() + { + _activationTaskSource.Reset(); + Volatile.Write(ref _activationClaim, 0); + } + + void CompleteFlow(Exception? exception) + { + // Completion has one structural owner: either pipeline retirement, or the source drain for + // an item which was never dispatched. Substitution transfers failed-item completion to the + // policy; it does not add a competing completer. + Debug.Assert(_completionCore.GetStatus(_completionCore.Version) is ValueTaskSourceStatus.Pending); + if (exception is null) + _completionCore.SetResult(default, runContinuationsAsynchronously: true); + else + _completionCore.SetException(exception, runContinuationsAsynchronously: true); } protected void SetObserver(PgClientFlowObserver observer, object? state) @@ -308,7 +375,7 @@ internal void DiscardUnqueued() // synchronously). The token is checked on entry only: the park itself is not cancelable, the // signal fires on every exit path (terminal, fault delivery, teardown), including the // cancel-delivered terminal. - internal ValueTask WaitForComplete(CancellationToken cancellationToken = default) + internal ValueTask WaitForComplete(CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); // Publish waiter-pending before the token capture: a reuse path that observes the flag @@ -326,23 +393,31 @@ internal ValueTask WaitForComplete(CancellationToken cancellationT // Synchronous consumers must not block on an async continuation whose dispatch requires another // scheduler turn. The event is allocated only for that uncommon path and is signaled after the // completion core has published its result. - internal PgClientFlow WaitForCompleteSynchronously(CancellationToken cancellationToken = default) + internal void WaitForCompleteSynchronously(CancellationToken cancellationToken = default) { cancellationToken.ThrowIfCancellationRequested(); Volatile.Write(ref _completionWaiterPending, 1); var token = _completionCore.Version; - var completionEvent = _completionEvent; + if (_completionCore.GetStatus(token) is not ValueTaskSourceStatus.Pending) + { + _ = ((IValueTaskSource)this).GetResult(token); + return; + } + + var state = GetOrCreateOptionalState(); + var completionEvent = state.CompletionEvent; if (completionEvent is null) { var created = new ManualResetEventSlim(); - completionEvent = Interlocked.CompareExchange(ref _completionEvent, created, null) ?? created; + completionEvent = Interlocked.CompareExchange( + ref state.CompletionEvent, created, null) ?? created; if (!ReferenceEquals(completionEvent, created)) created.Dispose(); } if (_completionCore.GetStatus(token) is ValueTaskSourceStatus.Pending) completionEvent.Wait(); - return ((IValueTaskSource)this).GetResult(token); + _ = ((IValueTaskSource)this).GetResult(token); } /// True while a completion waiter holds an unconsumed token on this tenure's signal. Reuse @@ -365,20 +440,21 @@ public void Reset() // rollover, a fail-loud TimeoutException at worst). If cancellation-aware flows become reusable, // that same reference-plus-stamp identity should be the cancellation coordinator's owner: a raw // reference cannot distinguish retained attribution from a later tenure whose window restarts at - // zero. Until the stamp lands, refuse to recycle a timeout-armed flow rather than let the race - // silently reappear. - if (EnableActivationTimeout) - ThrowHelper.ThrowInvalidOperation("Cannot pool a flow with EnableActivationTimeout: a recycled instance can be wrong-tenure-completed by a stale activation timeout. Implement generation-checked completion first."); + // zero. Production must refuse to recycle a timeout-armed flow until that stamp lands. + // EXPERIMENT ONLY: deliberately permit timeout-armed flows to be recycled so the unified + // flow's allocation-free throughput ceiling can be measured. This is not safe against a + // stale activation-timeout heartbeat completing a later tenure. _started = false; _completed = false; // Version bump per tenure. Cross-tenure completer staleness rests on the done -> torn-down // -> retired layering (Complete precedes recycle), the same basis as the rest of this reset. _completionCore.Reset(); - _completionEvent?.Reset(); - _activationTaskSource.Reset(); + Volatile.Read(ref _optionalState)?.CompletionEvent?.Reset(); + ResetActivationSource(); _rfqCount = 0; _cancellationWindow = 0; _lastMessageInducesRfq = false; + _activationWasDispatched = false; HandoffEvent?.ResetPlacement(); _pendingTimeoutStarted = false; _enqueueOptions = FlowEnqueueOptions.None; @@ -398,6 +474,14 @@ public void Reset() protected virtual void OnHeartbeat(TimeSpan interval) {} protected virtual void OnAbort(Exception exception) {} + /// True when the flow resets the protocol-static read objects before its pipeline task completes + /// and hands out no result past its terminal. The idle edge then keeps those objects for the next + /// flow instead of replacing them to protect a handle retained past completion. + internal virtual bool ResetsSharedReadStateBeforeRelease => false; + /// Terminates the flow from a result callback. Result-producing flows route the fault to their + /// consumer and framework completion. Other flows never hand out results. + internal virtual void Fail(Exception exception) + => throw new InvalidOperationException("This flow does not produce command results.", exception); /// Graceful-shutdown observation point. Fires while StoppingToken is set but before the /// AbortToken escalation. Flow types whose body can park on a non-IO rendezvous (CommandFlow's /// GateTask) override this to wake it so the body short-circuits instead of waiting for @@ -407,10 +491,13 @@ protected virtual void OnStopping(Exception exception) {} /// No overridable hook may run after that signal: a completion waiter can immediately Reset and /// enqueue the same object for its next tenure. protected virtual void OnReleasing(Exception? exception) {} - protected virtual void OnCancellationWindowCompleted(int completedWindow, int remainingWindowCount) {} protected virtual void OnDiscarded() {} protected virtual void OnReset() {} + private protected bool HasSuccessfulActivation + => _activationTaskSource.GetStatus(_activationTaskSource.Version) + is ValueTaskSourceStatus.Succeeded; + // The per-flow handoff rendezvous primitive for the (wait-list-free) sync source handoff: non-null only // for a flow that needs a caller takeover (a sync CommandFlow with a parked caller). The source signals // it when it dequeues-and-holds the flow for that caller (OnExecutorSuspended), and the caller parks on @@ -422,7 +509,7 @@ protected virtual void OnReset() {} // bare flow ref. Keeps the handoff primitive off PgClientFlow's internal API, like _rfqCount. private protected virtual FlowHandoffEvent? HandoffEvent => null; - PgClientFlow IValueTaskSource.GetResult(short token) + FlowCompletion IValueTaskSource.GetResult(short token) { // Consume-then-clear: the release store orders the core consumption before the flag clear, // so a reuse path's acquire read of "not pending" proves the token's lifetime ended. Fault @@ -436,21 +523,22 @@ PgClientFlow IValueTaskSource.GetResult(short token) Volatile.Write(ref _completionWaiterPending, 0); } } - ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _completionCore.GetStatus(token); - void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) + ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _completionCore.GetStatus(token); + void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) => _completionCore.OnCompleted(continuation, state, token, flags); - PgDecoder IValueTaskSource.GetResult(short token) => _activationTaskSource.GetResult(token); - ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _activationTaskSource.GetStatus(token); - void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) + FlowActivation IValueTaskSource.GetResult(short token) => _activationTaskSource.GetResult(token); + ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _activationTaskSource.GetStatus(token); + void IValueTaskSource.OnCompleted(Action continuation, object? state, short token, ValueTaskSourceOnCompletedFlags flags) => _activationTaskSource.OnCompleted(continuation, state, token, flags); - protected readonly struct Context + protected internal readonly struct Context { readonly ExecutionControl _executionControl; internal Context(ExecutionControl executionControl) => _executionControl = executionControl; + internal PgDecoder Decoder => _executionControl.Decoder; /// Graceful drain signal. Poll at handoff/coordination boundaries (per-CommandResult for /// CommandFlow) to switch to drain mode. I/O keeps running so the wire reaches a clean @@ -484,6 +572,11 @@ public PgEncoder GetEncoder() internal ValueTask WaitForCancellationAttempt() => _executionControl.WaitForCancellationAttempt(); + internal void SubmitDetached(Action action, object? state, bool preferLocal = true) + => _executionControl.SubmitDetached(action, state, preferLocal); + internal void SubmitDetached(IThreadPoolWorkItem workItem, bool preferLocal = true) + => _executionControl.SubmitDetached(workItem, preferLocal); + internal void RequestBackendCancellation(PgClientFlow instigator, int window, BackendCancellationTiming timing, TaskCompletionSource? delivery, object episodeKey, int scope, BackendCancellationTiming subsequentTiming) @@ -492,7 +585,7 @@ internal void RequestBackendCancellation(PgClientFlow instigator, int window, internal void RequestBackendCancellation(PgClientFlow instigator, int window, BackendCancellationTiming timing, TaskCompletionSource? delivery = null) => RequestBackendCancellation(instigator, window, timing, delivery, new object(), - (int)Flows.CommandFlow.CancellationScope.CurrentWindow, timing); + (int)Flows.CommandExecutionCancellationScope.CurrentWindow, timing); /// Returns an awaitable for the decoder. Activation is a cross-flow rendezvous completed by /// another flow's thread, so GetResult throws if not yet completed - async bodies await, @@ -513,7 +606,7 @@ public DecoderAwaitable GetDecoderAuto(CancellationToken cancellationToken = def // compiler checks IsCompleted and only schedules via (Unsafe)OnCompleted(Action) when not ready. // Direct dispatchers (CommandFlow's shared-promise pattern) instead use IsCompleted + // (Unsafe)OnCompleted(Action, object?) to register without a closure allocation. - protected readonly struct DecoderAwaitable : ICriticalNotifyCompletion + protected internal readonly struct DecoderAwaitable : ICriticalNotifyCompletion { readonly ExecutionControl control; readonly CancellationToken cancellationToken; @@ -588,7 +681,7 @@ public void UnsafeOnCompleted(Action continuation, object? state) // The ConfigureAwait(false) variant: skips scheduling-context capture. Action overloads are // for the C# `await` syntax (compiler calls UnsafeOnCompleted on ICriticalNotifyCompletion). - protected readonly struct ConfiguredDecoderAwaitable : ICriticalNotifyCompletion + protected internal readonly struct ConfiguredDecoderAwaitable : ICriticalNotifyCompletion { readonly ExecutionControl control; readonly CancellationToken cancellationToken; @@ -659,6 +752,7 @@ public void UnsafeOnCompleted(Action continuation, object? state) internal readonly struct ExecutionControl(PgClientFlow flow, PgClientProtocol.Control control) { internal PgClientFlow Flow => flow; + internal PgDecoder Decoder => control.Decoder; public bool SupportsDeferredFlush => flow is { _supportsDeferredFlush: true, _isAsyncAtDispatch: true }; public bool StallsPipeline => !SupportsDeferredFlush; @@ -666,6 +760,11 @@ internal readonly struct ExecutionControl(PgClientFlow flow, PgClientProtocol.Co public bool HasQueuedFlow => control.HasQueuedFlow; public bool IsInlineDrive => control.IsInlineDrive; + internal void SubmitDetached(Action action, object? state, bool preferLocal = true) + => control.SubmitDetached(action, state, preferLocal); + internal void SubmitDetached(IThreadPoolWorkItem workItem, bool preferLocal = true) + => control.SubmitDetached(workItem, preferLocal); + // Small optimization to allow us to skip the final sync message if we can piggyback on the flow's final rfq. public bool LastMessageInducesRfq => flow._lastMessageInducesRfq; @@ -708,6 +807,13 @@ public void OnMessageWrite(PgTypes.FrontendType type) } } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + internal static bool ShouldHandle(PgTypes.BackendType type) + => type is PgTypes.BackendType.ReadyForQuery + or PgTypes.BackendType.NoticeResponse + or PgTypes.BackendType.NotificationResponse + or PgTypes.BackendType.ParameterStatus; + /// Try-shape sync attempt: returns true if the message was processed without I/O. handled is /// true if the protocol layer consumed it (caller skips and pulls the next), false if it /// should be surfaced to the flow. Returns false only when a handler genuinely needs async @@ -716,13 +822,9 @@ public void OnMessageWrite(PgTypes.FrontendType type) [MethodImpl(MethodImplOptions.AggressiveInlining)] public bool TryHandleMessage(in BackendMessage backendMessage, out bool handled) { - if (backendMessage.Header.Type - is PgTypes.BackendType.ReadyForQuery - or PgTypes.BackendType.NoticeResponse - or PgTypes.BackendType.NotificationResponse - or PgTypes.BackendType.ParameterStatus) + if (ShouldHandle(backendMessage.Header.Type)) { - return TryHandleMessageCore(backendMessage, out handled); + return TryHandleKnownMessage(backendMessage, out handled); } handled = false; return true; @@ -734,16 +836,12 @@ or PgTypes.BackendType.NotificationResponse [MethodImpl(MethodImplOptions.AggressiveInlining)] public ValueTask HandleMessageAuto(in BackendMessage backendMessage) { - return backendMessage.Header.Type - is PgTypes.BackendType.ReadyForQuery - or PgTypes.BackendType.NoticeResponse - or PgTypes.BackendType.NotificationResponse - or PgTypes.BackendType.ParameterStatus + return ShouldHandle(backendMessage.Header.Type) ? HandleMessageAutoCore(backendMessage) : new(false); } [MethodImpl(MethodImplOptions.NoInlining)] - bool TryHandleMessageCore(BackendMessage backendMessage, out bool handled) + internal bool TryHandleKnownMessage(BackendMessage backendMessage, out bool handled) { switch (backendMessage.Header.Type) { @@ -751,7 +849,6 @@ bool TryHandleMessageCore(BackendMessage backendMessage, out bool handled) flow._rfqCount -= 1; var completedWindow = flow._cancellationWindow++; control.OnFlowRfq(flow, backendMessage, completedWindow, flow._rfqCount); - flow.OnCancellationWindowCompleted(completedWindow, flow._rfqCount); handled = false; return true; case PgTypes.BackendType.NoticeResponse: @@ -779,7 +876,6 @@ ValueTask HandleMessageAutoCore(BackendMessage backendMessage) flow._rfqCount -= 1; var completedWindow = flow._cancellationWindow++; control.OnFlowRfq(flow, backendMessage, completedWindow, flow._rfqCount); - flow.OnCancellationWindowCompleted(completedWindow, flow._rfqCount); goto default; case PgTypes.BackendType.NoticeResponse: // We sink all notices (this includes RAISE notices) and expect those to end up on the flow for user retrieval/logging. @@ -844,22 +940,18 @@ public ValueTask ExecuteAuto() [MethodImpl(MethodImplOptions.NoInlining)] ValueTask ExecuteSynchronously() => flow.ExecuteAuto(new(this)); - public void Activate(PgDecoder decoder) + public void Activate() { - flow._activationCancellationTokenRegistration.Dispose(); + var activationRegistration = flow.TakeActivationRegistration(); + activationRegistration.Dispose(); // If none of the cancellations triggered, we have a problem, throw. - if (!flow._activationTaskSource.TrySetResult(decoder, runContinuationsAsynchronously: false) + if (!flow.TrySetActivationResult(runContinuationsAsynchronously: false) && !(flow._remainingActivationTimeout <= TimeSpan.Zero) && !control.AbortToken.IsCancellationRequested - && !flow._activationCancellationTokenRegistration.Token.IsCancellationRequested) + && !activationRegistration.Token.IsCancellationRequested) ThrowHelper.ThrowInvalidOperation("Flow was already activated unexpectedly."); } - public void RegisterDecoderOnHeartbeat(Action action) - { - flow._decoderOnHeartbeatAction = action; - } - public void OnHeartbeat(TimeSpan interval) { if (PropagateTermination()) @@ -867,7 +959,9 @@ public void OnHeartbeat(TimeSpan interval) OnActivationHeartbeat(interval); - flow._decoderOnHeartbeatAction?.Invoke(interval); + if (flow._activationTaskSource.GetStatus(flow._activationTaskSource.Version) + is ValueTaskSourceStatus.Succeeded) + control.Decoder.OnHeartbeat(interval); flow.OnHeartbeat(interval); } @@ -885,7 +979,7 @@ public bool PropagateTermination() if (control.AbortToken.IsCancellationRequested && !flow._completed) { var ex = control.FlowTerminationException; - flow._activationTaskSource.TrySetException(ex, runContinuationsAsynchronously: true); + flow.TrySetActivationException(ex, runContinuationsAsynchronously: true); flow.OnAbort(ex); return true; } @@ -906,7 +1000,7 @@ public void OnActivationHeartbeat(TimeSpan interval) if (flow._remainingActivationTimeout != Timeout.InfiniteTimeSpan && flow._remainingActivationTimeout != TimeSpan.Zero && flow._activationTaskSource.GetStatus(flow._activationTaskSource.Version) is ValueTaskSourceStatus.Pending && (flow._remainingActivationTimeout -= interval) <= TimeSpan.Zero) - flow._activationTaskSource.TrySetException(new TimeoutException("Operation timed out waiting for activation."), runContinuationsAsynchronously: true); + flow.TrySetActivationException(new TimeoutException("Operation timed out waiting for activation."), runContinuationsAsynchronously: true); } /// Fail a never-started flow drained from the backlog at shutdown with the wire-death reason. The @@ -941,14 +1035,14 @@ public void DetachForMigration(PgClientFlowSource source) Debug.Assert(!flow._started && !flow._completed); // No body has awaited activation yet, so there is no activation-cycle registration to // dismantle and reconstruct on the replacement source. - Debug.Assert(flow._activationCancellationTokenRegistration == default); + Debug.Assert(Volatile.Read(ref flow._optionalState)?.ActivationRegistration == default); // Forceful shutdown may have propagated the retired wire's abort into this inert flow // before the source drain transferred it. No body can have observed the pre-start gate; // reset that wire-local verdict so replacement dispatch can activate the same operation. if (IsDecoderSettled) { Debug.Assert(control.AbortToken.IsCancellationRequested); - flow._activationTaskSource.Reset(); + flow.ResetActivationSource(); } if (flow.HandoffEvent?.PlacementSource is not null) flow.DetachPlacementSource(source.SourceState); @@ -972,7 +1066,7 @@ public void Release(Exception? exception = null) flow._completed = true; if (flow.OwnsAdmissionBarrier) control.ReleaseAdmissionBarrier(); - flow._activationCancellationTokenRegistration.Dispose(); + flow.TakeActivationRegistration().Dispose(); var observer = flow._observer; var observerState = flow._observerState; try { flow.OnReleasing(exception); } @@ -996,11 +1090,8 @@ public void Release(Exception? exception = null) // Async continuation dispatch: completers run in retirement/teardown contexts where // inline caller continuations are a re-entrancy hazard, the contract the old TCS's // RunContinuationsAsynchronously carried, minus its unconditional thread-pool destination. - if (exception is not null) - flow._completionCore.TrySetException(exception, runContinuationsAsynchronously: true); - else - flow._completionCore.TrySetResult(flow, runContinuationsAsynchronously: true); - flow._completionEvent?.Set(); + flow.CompleteFlow(exception); + Volatile.Read(ref flow._optionalState)?.CompletionEvent?.Set(); // The completed observer runs from CompleteItem in the advancer/retirement work-item // context: a raw throw would crash that thread unobserved. Don't swallow either - a // throwing completed observer means the consumer-side integration is broken, so the @@ -1051,17 +1142,21 @@ public bool IsDecoderReady public bool IsDecoderSettled => flow._activationTaskSource.GetStatus(flow._activationTaskSource.Version) is not ValueTaskSourceStatus.Pending; public PgDecoder GetDecoderResult() - => flow._activationTaskSource.GetResult(flow._activationTaskSource.Version); + { + _ = flow._activationTaskSource.GetResult(flow._activationTaskSource.Version); + return control.Decoder; + } public void OnDecoder(Action continuation, object? state, ValueTaskSourceOnCompletedFlags flags) => flow._activationTaskSource.OnCompleted(continuation, state, flow._activationTaskSource.Version, flags); // Bridge to Task for callers that need to block (sync flow body using // .GetAwaiter().GetResult()) or to compose with Task-based combinators. MVTSC has no // blocking GetResult of its own, so this is the only safe sync-wait path. - public Task GetDecoderTask(CancellationToken cancellationToken) + public async Task GetDecoderTask(CancellationToken cancellationToken) { RegisterActivationCancellation(cancellationToken); - return new ValueTask(flow, flow._activationTaskSource.Version).AsTask(); + await new ValueTask(flow, flow._activationTaskSource.Version).ConfigureAwait(false); + return control.Decoder; } // Registers caller cancellation against the activation source so a flow can unwind @@ -1071,11 +1166,12 @@ public void RegisterActivationCancellation(CancellationToken cancellationToken) { if (!cancellationToken.CanBeCanceled) return; - if (flow._activationCancellationTokenRegistration != default) + var state = flow.GetOrCreateOptionalState(); + if (state.ActivationRegistration != default) ThrowHelper.ThrowInvalidOperation("Concurrent activation result awaits are not supported."); - flow._activationCancellationTokenRegistration = cancellationToken.UnsafeRegister( + state.ActivationRegistration = cancellationToken.UnsafeRegister( static (state, token) => - ((PgClientFlow)state!)._activationTaskSource.TrySetException(new OperationCanceledException(token), runContinuationsAsynchronously: true), + ((PgClientFlow)state!).TrySetActivationException(new OperationCanceledException(token), runContinuationsAsynchronously: true), flow); } } diff --git a/Slon/Pg/Protocol/PgClientFlowSource.cs b/Slon/Pg/Protocol/PgClientFlowSource.cs index f65d521..0bbeb6e 100644 --- a/Slon/Pg/Protocol/PgClientFlowSource.cs +++ b/Slon/Pg/Protocol/PgClientFlowSource.cs @@ -543,6 +543,7 @@ WaitForNextAwaitable WaitCore() // Only reached on real write backpressure (the flush didn't complete inline), so a pooled // box is plenty - the promise-reuse builder would be overkill for how rarely this fires. + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] async ValueTask FlushThenWaitAsync(ValueTask flushTask) { diff --git a/Slon/Pg/Protocol/PgClientProtocol.Cancellation.cs b/Slon/Pg/Protocol/PgClientProtocol.Cancellation.cs index 19395c4..0902d5a 100644 --- a/Slon/Pg/Protocol/PgClientProtocol.Cancellation.cs +++ b/Slon/Pg/Protocol/PgClientProtocol.Cancellation.cs @@ -55,7 +55,7 @@ internal void RequestServerCancellation(PgClientFlow instigator, coordinator = GetOrCreateCancellationCoordinatorLocked(); } coordinator.RequestCancellation(instigator, window, timing, delivery, - episodeKey, scope == (int)Flows.CommandFlow.CancellationScope.RemainingFlow, + episodeKey, scope == (int)Flows.CommandExecutionCancellationScope.RemainingFlow, subsequentTiming); // Flow release publishes its coordinator down-edge before IsCompleted. Recheck after // publishing the episode so either that down-edge or this observation retires it. diff --git a/Slon/Pg/Protocol/PgClientProtocol.cs b/Slon/Pg/Protocol/PgClientProtocol.cs index 55f326e..3048674 100644 --- a/Slon/Pg/Protocol/PgClientProtocol.cs +++ b/Slon/Pg/Protocol/PgClientProtocol.cs @@ -36,6 +36,13 @@ public enum FlowEnqueueOptions : byte AllowMigration = 4 } +[Experimental(ExperimentalDiagnostics.PostgreSqlLowerLayer)] +public enum PgClientProtocolHeartbeatMode : byte +{ + Automatic, + External +} + interface IProtocolStatic { ref readonly T Value { get; } @@ -79,6 +86,7 @@ public PgClientProtocolOptions(PgClientOptions options) public Encoding DefaultClientEncoding { get; set; } public TimeSpan FlowActivationTimeout { get; set; } public TimeSpan HeartbeatInterval { get; set; } = Heartbeat.DefaultInterval; + public PgClientProtocolHeartbeatMode HeartbeatMode { get; set; } public TimeSpan ReadTimeout { get; set; } = PgClientOptions.DefaultReadTimeout; public TimeSpan WriteTimeout { get; set; } = TimeSpan.FromSeconds(10); // Allocation-free grace before starting a backend CancelRequest. Heartbeat supplies the clock. @@ -86,7 +94,7 @@ public PgClientProtocolOptions(PgClientOptions options) public TimeSpan CancellationTimeout { get; set; } = TimeSpan.FromSeconds(10); public TimeSpan CancellationRetryInterval { get; set; } = TimeSpan.FromSeconds(1); internal PgSessionResetOptions SessionReset { get; set; } = new(); - public int DataRowStreamingThreshold { get; set; } = BackendMessageBatch.Segmenter.DefaultDataRowStreamingThreshold; + public int DataRowStreamingThreshold { get; set; } = BackendMessageCursor.DefaultDataRowStreamingThreshold; public int MaxInFlightFlowsPerWire { get; set; } public ILoggerFactory LoggerFactory { get; set; } = NullLoggerFactory.Instance; // Datasource bootstrap supplies this. A standalone raw protocol can omit backend identity and @@ -177,7 +185,6 @@ await _protocol.StartAsync(_transport, _flow, hosting, cancellationToken) TransportConnection _connection = null!; IOutputWriter _pipeWriter = null!; ProtocolDataWriter _protocolDataWriter = null!; - PipeSegmentEnumerator _pipeSegmentEnumerator = null!; PgDecoder _pgDecoder = null!; LoadObserver? _loadObserver; @@ -298,22 +305,30 @@ public static PgClientProtocol Create(PgClientProtocolOptions protocolOptions) void Initialize(TransportConnection connection, Hosting hosting) { + PipelineScheduler? connectionScheduler = null; + if ((_options.ExecutionScheduler is null || _options.ActivationScheduler is null) + && connection.Scheduler is { } scheduler) + connectionScheduler = new DelegatedPipelineScheduler(scheduler); + _executionScheduler = _options.ExecutionScheduler + ?? connectionScheduler ?? PipelineScheduler.ThreadPool; _activationScheduler = _options.ActivationScheduler + ?? connectionScheduler ?? PipelineScheduler.ThreadPool; _connection = connection; _pipeWriter = connection.Writer as IOutputWriter ?? new PipeOutputWriter(connection.Writer); _protocolDataWriter = new(_pipeWriter, PgClientOptions.PreStartupEncoding, connection.WaitUntilWritable, AbortToken, FlowControl, _options.WriteTimeout); - _pipeSegmentEnumerator = new(connection.Reader, - new(_options.DataRowStreamingThreshold), ownsReader: true); - _pgDecoder = new(_pipeSegmentEnumerator, AbortToken, _options.ReadTimeout, _options.ReadTimeoutArmed); + _pgDecoder = new(connection.Reader, _options.DataRowStreamingThreshold, + AbortToken, _options.ReadTimeout, _options.ReadTimeoutArmed, + ownsReader: true); _admissionAvailable = hosting.AdmissionAvailable; _loadObserver = hosting.LoadObserver; - if (!hosting.DrivesHeartbeat) + if (!hosting.DrivesHeartbeat && + _options.HeartbeatMode is PgClientProtocolHeartbeatMode.Automatic) { _heartbeat = new(_options.HeartbeatInterval, _options.TimeProvider, _logger); _heartbeat.Register(period => Heartbeat(period)); @@ -474,7 +489,7 @@ static void ReleaseTransportOnStartFailure(TransportConnection connection, Excep connection.Reader.Complete(reason); } - async ValueTask StartAsync(StartupFlow flow, ValueTask flowCompletion, CancellationToken cancellationToken = default) + async ValueTask StartAsync(StartupFlow flow, ValueTask flowCompletion, CancellationToken cancellationToken = default) { _source = PgClientFlowSource.Create( this, FlowControl, _executionScheduler, _options.MaxInFlightFlowsPerWire); @@ -616,10 +631,10 @@ public bool TryQueue(PgClientFlow flow, FlowEnqueueOptions options = FlowEnqueue if ((options & FlowEnqueueOptions.RequireExistingPipeline) != 0) { - if (!TryQueueFlow(flow, options, static protocol => protocol.PipelineDepth > 0, this)) + if (!TryQueueFlow(flow, options, requireExistingPipeline: true)) return false; } - else if (!TryQueueFlow(flow, options, null, (object?)null)) + else if (!TryQueueFlow(flow, options)) return false; try @@ -744,14 +759,14 @@ internal async ValueTask BeginExclusiveScopeAsync( return scope; } - bool TryQueueFlow(PgClientFlow flow, FlowEnqueueOptions options, - Func? predicate = null, TState state = default!) - => TryQueueFlow(flow, ProtocolStatus.Ready, options, predicate, state); + bool TryQueueFlow(PgClientFlow flow, FlowEnqueueOptions options, + bool requireExistingPipeline = false) + => TryQueueFlow(flow, ProtocolStatus.Ready, options, requireExistingPipeline); bool TryQueueFlow(PgClientFlow flow, ProtocolStatus requiredStatus) - => TryQueueFlow(flow, requiredStatus, FlowEnqueueOptions.None); - bool TryQueueFlow(PgClientFlow flow, ProtocolStatus requiredStatus, - FlowEnqueueOptions options, Func? predicate = null, TState state = default!) + => TryQueueFlow(flow, requiredStatus, FlowEnqueueOptions.None, requireExistingPipeline: false); + bool TryQueueFlow(PgClientFlow flow, ProtocolStatus requiredStatus, + FlowEnqueueOptions options, bool requireExistingPipeline) { // A handoff-capable sync flow is held at its FIFO turn. Consumer-driven flows defer the // handoff; self-driven flows take it as part of admission. @@ -764,7 +779,7 @@ bool TryQueueFlow(PgClientFlow flow, ProtocolStatus requiredStatus, (requiredStatus is ProtocolStatus.Ready && _admissionBlocked != 0)) return false; - if (predicate?.Invoke(state) == false) + if (requireExistingPipeline && _pipeline.Depth == 0) return false; if (requiredStatus is ProtocolStatus.Ready && @@ -1125,31 +1140,53 @@ internal ValueTask Heartbeat(TimeSpan period) return new(); } + public ValueTask HeartbeatAsync(TimeSpan elapsed) + => Heartbeat(elapsed); + void PropagateFlowHeartbeat(TimeSpan period) { var control = FlowControl; + control.BeginHeartbeatObservation(); try { - _source.OnActivationHeartbeat(period); - } - catch (Exception ex) - { - SlonLogMessages.UnobservedCallbackException( - _logger, ex, "the source heartbeat callback"); - } - - foreach (var flow in GetFlows()) - { - try + // Backlog observation will get its own bounded source frontier. Holding the gate for + // this phase is sufficient for the experiment and free when the backlog is empty. + lock (control.HeartbeatObservationLock) { - flow.GetExecutionControl(control).OnHeartbeat(period); + try + { + _source.OnActivationHeartbeat(period); + } + catch (Exception ex) + { + SlonLogMessages.UnobservedCallbackException( + _logger, ex, "the source heartbeat callback"); + } } - catch (Exception ex) + + var flows = _pipeline.GetEnumerator(out var frontier); + while (flows.MoveNext()) { - SlonLogMessages.UnobservedCallbackException( - _logger, ex, "a flow heartbeat callback"); + lock (control.HeartbeatObservationLock) + { + if (frontier.IsRetired(flows.Position)) + continue; + try + { + flows.Current.GetExecutionControl(control).OnHeartbeat(period); + } + catch (Exception ex) + { + SlonLogMessages.UnobservedCallbackException( + _logger, ex, "a flow heartbeat callback"); + } + } } } + finally + { + control.EndHeartbeatObservation(); + } } void PropagateFlowTermination() @@ -1202,6 +1239,18 @@ public Policy(PgClientProtocol protocol, Control control, ExclusiveScopeState? l [MethodImpl(MethodImplOptions.AggressiveInlining)] public void CompleteItem(PgClientFlow item, Exception? exception) + { + if (_control.IsHeartbeatObserving) + { + lock (_control.HeartbeatObservationLock) + CompleteItemCore(item, exception); + return; + } + CompleteItemCore(item, exception); + } + + [MethodImpl(MethodImplOptions.AggressiveInlining)] + void CompleteItemCore(PgClientFlow item, Exception? exception) { if (exception is PgClientClosedException && _control.ClosedException is not null) exception = _control.FlowTerminationException; @@ -1323,22 +1372,53 @@ public ValueTask ExecuteItemAsync(PgClientFlow item, bool pi PromiseAsyncValueTaskMethodBuilder.Promise = null; } - [RuntimeAsyncMethodGeneration(false)] - [AsyncMethodBuilder(typeof(PromiseAsyncValueTaskMethodBuilder<>))] - static async ValueTask ExecuteCore( + static ValueTask ExecuteCore( Control control, PgClientFlow item, CancellationToken cancellationToken) { - await control.WaitForCancellationAttempt().ConfigureAwait(false); + var pending = control.WaitForCancellationAttempt(); + if (!pending.IsCompletedSuccessfully) + return AwaitPriorWork(control, item, cancellationToken, pending, checkFlush: true); + pending.GetAwaiter().GetResult(); // A flow may defer this flush only when its first phase cannot wait for decoder input. if (!item.SupportsDeferredFlush && control.UnflushedBytes != 0) - await control.FlushAsync(cancellationToken).ConfigureAwait(false); + { + pending = control.FlushAsync(cancellationToken); + if (!pending.IsCompletedSuccessfully) + return AwaitPriorWork(control, item, cancellationToken, pending, checkFlush: false); + pending.GetAwaiter().GetResult(); + } + + var execution = control.Execute(item); + if (!execution.IsCompletedSuccessfully) + return AwaitExecution(execution); + var tasks = execution.GetAwaiter().GetResult(); + return new ValueTask( + new PipelineItemResult(tasks.TrailingExecutionTask, tasks.PipelineTask)); + } + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PromiseAsyncValueTaskMethodBuilder<>))] + static async ValueTask AwaitPriorWork( + Control control, PgClientFlow item, CancellationToken cancellationToken, + ValueTask pending, bool checkFlush) + { + await pending.ConfigureAwait(false); + if (checkFlush && !item.SupportsDeferredFlush && control.UnflushedBytes != 0) + await control.FlushAsync(cancellationToken).ConfigureAwait(false); var tasks = await control.Execute(item).ConfigureAwait(false); - return new PipelineItemResult(tasks.TrailingExecutionTask, tasks.PipelineTask); + return new(tasks.TrailingExecutionTask, tasks.PipelineTask); + } + + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PromiseAsyncValueTaskMethodBuilder<>))] + static async ValueTask AwaitExecution(ValueTask execution) + { + var tasks = await execution.ConfigureAwait(false); + return new(tasks.TrailingExecutionTask, tasks.PipelineTask); } - // Stock builder (no shared promise) for pipeline-task recovery. Body identical to ExecuteCore. + // Stock builder (no shared promise) for pipeline-task recovery. static async ValueTask ExecutePipelineTaskRecovery( Control control, PgClientFlow item, CancellationToken cancellationToken) { @@ -1356,9 +1436,9 @@ public void ActivateHeadItem(PgClientFlow item, bool preferAsync = true) // Only the body wake is deferred below. _control.BindDecoder(item); - // Inline-activate when the framework allows it (preferAsync=false) or the flow is sync: - // sync flows park on a kernel wait-handle signal, bounded cost, safe under the advancer - // latch. Async flows can attach arbitrary await continuations, so they go through TP. + // Inline-activate when the framework allows it (preferAsync=false) or the flow is sync. + // CommandFlow retains whether an async activation was dispatched so its consumer-facing + // ready publication supplies the scheduling firewall only when activation ran inline. if (preferAsync && item.IsAsyncAtDispatch) { // The flow itself is the work item: an immutable (flow, control) pairing per queued @@ -1501,8 +1581,31 @@ sealed class PipelineSlots( } } - internal sealed class Control(PgClientProtocol protocol, bool poolFacing) : IProtocolStatic + internal sealed class Control(PgClientProtocol protocol, bool poolFacing) : + IProtocolStatic, + IProtocolStatic { + readonly Lock _heartbeatObservationLock = new(); + int _heartbeatObserving; + + internal Lock HeartbeatObservationLock => _heartbeatObservationLock; + internal bool IsHeartbeatObserving => Volatile.Read(ref _heartbeatObserving) != 0; + + internal void BeginHeartbeatObservation() + { + lock (_heartbeatObservationLock) + { + Debug.Assert(_heartbeatObserving == 0); + Volatile.Write(ref _heartbeatObserving, 1); + } + } + + internal void EndHeartbeatObservation() + { + lock (_heartbeatObservationLock) + Volatile.Write(ref _heartbeatObserving, 0); + } + // The pipeline whose slots this Control reads, bound right after that pipeline is created. The // outer (pool-facing) Control reads the protocol's own pipeline; an exclusive flow's inner // Control reads its inner pipeline - both through the same IPipelineSlots handle, so any @@ -1525,6 +1628,10 @@ public void BindPipeline( public bool IsInlineDrive => _source.IsInlineDrive; public long UnflushedBytes => protocol.UnflushedBytes; public ValueTask FlushAsync(CancellationToken cancellationToken) => protocol.FlushAsync(cancellationToken); + internal void SubmitDetached(Action action, object? state, bool preferLocal = true) + => protocol._activationScheduler.SubmitDetached(action, state, preferLocal); + internal void SubmitDetached(IThreadPoolWorkItem workItem, bool preferLocal = true) + => protocol._activationScheduler.SubmitDetached(workItem, preferLocal); PgClientFlow? _cancellationActivatedFlow; internal (PgClientFlow? Owner, int Window) CancellationActivation { @@ -1599,7 +1706,7 @@ public void BindShells(PgDecoder decoder, ProtocolDataWriter writer) _writer = writer; } - PgDecoder Decoder => _decoder ?? protocol._pgDecoder; + internal PgDecoder Decoder => _decoder ?? protocol._pgDecoder; public PgClientFlow? ExecutingFlow => _slots.Executing; public PgClientFlow? ActivatedFlow => _slots.Activated; @@ -1812,7 +1919,7 @@ internal void BindDecoder(PgClientFlow flow) // off the executor via the TP dispatch. Safe to lag the flow's retirement: TrySetResult no-ops // on a flow the abort already faulted. internal void Activate(PgClientFlow flow) - => flow.GetExecutionControl(this).Activate(Decoder); + => flow.GetExecutionControl(this).Activate(); // Self-evict route for the flow layer's release-callback seam (see ExecutionControl.Release). internal void FailProtocol(Exception? reason) => protocol.FailProtocol(reason); @@ -1833,9 +1940,15 @@ internal void AssignCancellationBoundary(PgClientFlow flow, int window) internal void OnReleasing(PgClientFlow flow) { + // Draghi clears ActivatedFlow at the exact zero edge but retains the old activation turn + // through this callback. Null is therefore both the idle witness and an exclusive + // pre-release window in which no successor reader can activate. + var idle = ActivatedFlow is null; + Decoder.EndResultBuffering(flow); + if (idle) + Decoder.ReleaseReadBufferAtIdle(); protocol._serverParameterState.CommitFlow(); ClearCancellationActivation(flow); - var idle = ActivatedFlow is null; protocol.OnFlowReleased(flow, poolFacing && idle); // Inner exclusive-scope subflows are not pool load units; only the outer pipeline reports // admission-to-retirement lifetimes to its host. @@ -1844,8 +1957,9 @@ internal void OnReleasing(PgClientFlow flow) flow.GetExecutionControl(this).StallsPipeline); // Draghi clears the activated slot only at the exact idle edge and before CompleteItem. - // Release the shared read objects before the flow's terminal observer can reuse them. - if (idle) + // Release the shared read objects before the flow's terminal observer can reuse them, + // unless the flow already reset them and hands out nothing past its terminal. + if (idle && !flow.ResetsSharedReadStateBeforeRelease) _commandFlowReadState = new(); } @@ -1863,5 +1977,8 @@ internal void OnIdle() CommandFlow.ReadState _commandFlowReadState = new(); ref readonly CommandFlow.ReadState IProtocolStatic.Value => ref _commandFlowReadState; + readonly CommandFlow.ReadPromiseState _commandFlowReadPromiseState = new(); + ref readonly CommandFlow.ReadPromiseState IProtocolStatic.Value + => ref _commandFlowReadPromiseState; } } diff --git a/Slon/Pg/Protocol/PgDecoder.cs b/Slon/Pg/Protocol/PgDecoder.cs index 2621378..4d4cfb4 100644 --- a/Slon/Pg/Protocol/PgDecoder.cs +++ b/Slon/Pg/Protocol/PgDecoder.cs @@ -19,10 +19,10 @@ namespace Slon.Pg.Protocol; public sealed class PgDecoder: IEnumerator, IAsyncEnumerator { readonly ProtocolReadPipe _pipe; + readonly StreamPipeReader? _directReader; readonly CancellationToken _abortToken; readonly TimeSpan _defaultReadTimeout; readonly Action? _readTimeoutArmed; - readonly Action _onHeartbeatAction; CancellationTokenSource _cancellationTokenSource; TimeSpan _readTimeout; @@ -32,6 +32,7 @@ public sealed class PgDecoder: IEnumerator, IAsyncEnumerator messageBatchEnumerator, - CancellationToken abortToken, TimeSpan defaultReadTimeout, Action? readTimeoutArmed = null) - : this(new ProtocolReadPipe(messageBatchEnumerator), abortToken, defaultReadTimeout, readTimeoutArmed) + internal PgDecoder(PipeReader reader, int dataRowStreamingThreshold, + CancellationToken abortToken, TimeSpan defaultReadTimeout, + Action? readTimeoutArmed = null, bool ownsReader = false) + : this(new ProtocolReadPipe( + reader, dataRowStreamingThreshold, ownsReader), + abortToken, defaultReadTimeout, readTimeoutArmed) { } internal ProtocolReadPipe Pipe => _pipe; internal Encoding ClientEncoding => _control.ClientEncoding; + internal void SetCurrentMessageLength(long messageLength) + => _pipe.SetCurrentMessageLength(messageLength); + + internal void CompleteCurrentMessage() + => _pipe.CompleteCurrentMessage(); + + internal bool ResultBuffering + { + get => _resultBufferingOwner is not null; + set + { + if (value) + { + var owner = CurrentExecutionControl.Flow; + if (!ReferenceEquals(_resultBufferingOwner, owner)) + _resultBufferingOwner = owner; + _pipe.EnableResultRetention(); + } + else + { + if (_resultBufferingOwner is null) + return; + _resultBufferingOwner = null; + _pipe.EndResultRetention(); + } + } + } + + internal void EndResultBuffering(PgClientFlow owner) + { + if (!ReferenceEquals(_resultBufferingOwner, owner)) + return; + _resultBufferingOwner = null; + _pipe.EndResultRetention(); + } + + internal void ReleaseReadBufferAtIdle() + => _pipe.ReleaseReadBufferAtIdle(); + + void ValidateResultBufferingOwner() + { + var owner = _resultBufferingOwner; + if (owner is not null + && !ReferenceEquals(CurrentExecutionControl.Flow, owner)) + EndResultBuffering(owner); + } + + void PrepareRead() + { + ValidateResultBufferingOwner(); + _pipe.PrepareRead(); + } + + bool CompleteRead( + in ReadResult result, CancellationToken cancellationToken, out bool completed) + { + ValidateResultBufferingOwner(); + return _pipe.CompleteRead( + result, cancellationToken, out completed); + } + + bool ReadNext(TimeSpan timeout) + { + ValidateResultBufferingOwner(); + return _pipe.MoveNext(timeout); + } + + bool TryBeginDirectRead(CancellationToken cancellationToken, out ValueTask task) + { + if (_directReader is { SupportsDirectRead: true } directReader) + { + task = directReader.BeginDirectRead(cancellationToken); + return true; + } + task = default; + return false; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + bool CompleteDirectRead(int length, CancellationToken cancellationToken, + out ValueTask next, out bool readFinished, out bool completed) + { + if (!_directReader!.CompleteDirectRead(length, cancellationToken, out next, out var result)) + { + readFinished = false; + completed = false; + return false; + } + readFinished = true; + return CompleteRead(result, cancellationToken, out completed); + } + + void AbortDirectRead() => _directReader!.AbortDirectRead(); // Builds a scope-bound shell over the shared pipe with the scope's abort token. internal static PgDecoder CreateScopeShell(PgDecoder baseShell, CancellationToken abortToken, TimeSpan defaultReadTimeout) => new(baseShell._pipe, abortToken, defaultReadTimeout, baseShell._readTimeoutArmed); @@ -128,8 +225,6 @@ internal void Initialize(PgClientProtocol.Control control) if (!ReferenceEquals(_control, control)) _control = control; _pipe.BindDecoder(this); - // TODO we want a heartbeat setup directly through the protocol on construction. - CurrentExecutionControl.RegisterDecoderOnHeartbeat(_onHeartbeatAction); } /// @@ -143,19 +238,19 @@ public void UseReadTimeout(TimeSpan timeout) void RestoreDefaultReadTimeout() => _readTimeout = _defaultReadTimeout; - internal bool TryContinueCurrentMessage(SequencePosition consumed, long consumedLength, out CurrentSegmentBuffer result) - => _pipe.TryContinueCurrentMessage(consumed, consumedLength, out result); + internal bool TrySlideCurrentMessage(SequencePosition consumed, long consumedLength, out CurrentMessageBuffer result) + => _pipe.TrySlideCurrentMessage(consumed, consumedLength, out result); - internal ValueTask ContinueCurrentMessageAsync( + internal ValueTask SlideCurrentMessageAsync( SequencePosition consumed, long consumedLength, CancellationToken cancellationToken) { EnsureUsableCts(); - if (_pipe.TryContinueCurrentMessage(consumed, consumedLength, out var result)) + if (_pipe.TrySlideCurrentMessage(consumed, consumedLength, out var result)) return new(result); return Core(cancellationToken); - async ValueTask Core(CancellationToken cancellationToken) + async ValueTask Core(CancellationToken cancellationToken) { var timeoutSet = false; var frontierFlow = EnterCancellationReadFrontier(); @@ -165,8 +260,9 @@ async ValueTask Core(CancellationToken cancellationToken) { ArmReadTimeout(); timeoutSet = true; - return await _pipe.ContinueCurrentMessageAsync( + var read = await _pipe.BeginSlideCurrentMessageAsync( consumed, consumedLength, _cancellationTokenSource.Token).ConfigureAwait(false); + return _pipe.CompleteCurrentMessageRead(read, _cancellationTokenSource.Token); } catch (Exception ex) when (_cancellationTokenSource.IsCancellationRequested) { @@ -186,7 +282,7 @@ async ValueTask Core(CancellationToken cancellationToken) } } - internal CurrentSegmentBuffer ContinueCurrentMessage( + internal CurrentMessageBuffer SlideCurrentMessage( SequencePosition consumed, long consumedLength) { var timeoutSet = false; @@ -194,7 +290,8 @@ internal CurrentSegmentBuffer ContinueCurrentMessage( { ArmReadTimeout(); timeoutSet = true; - return _pipe.ContinueCurrentMessage(consumed, consumedLength, GetRemainingTimeout()); + return _pipe.SlideCurrentMessage( + consumed, consumedLength, GetRemainingTimeout()); } catch (Exception) when (_abortToken.IsCancellationRequested && _control.ClosedException is not null) { @@ -211,49 +308,62 @@ internal CurrentSegmentBuffer ContinueCurrentMessage( } } - internal bool TryExtendCurrentMessage(out CurrentSegmentBuffer result) + internal bool TryExtendCurrentMessage(out CurrentMessageBuffer result) => _pipe.TryExtendCurrentMessage(out result); - internal ValueTask ExtendCurrentMessageAsync(CancellationToken cancellationToken) + internal ValueTask ExtendCurrentMessageAsync(CancellationToken cancellationToken) { EnsureUsableCts(); if (_pipe.TryExtendCurrentMessage(out var result)) return new(result); - return Core(cancellationToken); + return ReadCurrentMessageAsync(cancellationToken, bufferAll: false); + } - async ValueTask Core(CancellationToken cancellationToken) + internal ValueTask BufferCurrentMessageAsync( + CancellationToken cancellationToken) + { + EnsureUsableCts(); + return ReadCurrentMessageAsync(cancellationToken, bufferAll: true); + } + + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] + async ValueTask ReadCurrentMessageAsync( + CancellationToken cancellationToken, bool bufferAll) + { + var timeoutSet = false; + var frontierFlow = EnterCancellationReadFrontier(); + var registration = cancellationToken.UnsafeRegister( + static (state, _) => ((CancellationTokenSource)state!).Cancel(), _cancellationTokenSource); + try { - var timeoutSet = false; - var frontierFlow = EnterCancellationReadFrontier(); - var registration = cancellationToken.UnsafeRegister( - static (state, _) => ((CancellationTokenSource)state!).Cancel(), _cancellationTokenSource); - try - { - ArmReadTimeout(); - timeoutSet = true; - return await _pipe.ExtendCurrentMessageAsync( - _cancellationTokenSource.Token).ConfigureAwait(false); - } - catch (Exception ex) when (_cancellationTokenSource.IsCancellationRequested) - { - throw TranslateReadCancellation(ex, cancellationToken); - } - catch (EndOfStreamException ex) - { - throw TranslateEof(ex); - } - finally - { - LeaveCancellationReadFrontier(frontierFlow); - registration.Dispose(); - if (timeoutSet) - SetRemainingTimeout(Timeout.InfiniteTimeSpan); - } + ArmReadTimeout(); + timeoutSet = true; + var read = await (bufferAll + ? _pipe.BeginBufferCurrentMessageAsync(_cancellationTokenSource.Token) + : _pipe.BeginExtendCurrentMessageAsync(_cancellationTokenSource.Token)) + .ConfigureAwait(false); + return _pipe.CompleteCurrentMessageRead(read, _cancellationTokenSource.Token); + } + catch (Exception ex) when (_cancellationTokenSource.IsCancellationRequested) + { + throw TranslateReadCancellation(ex, cancellationToken); + } + catch (EndOfStreamException ex) + { + throw TranslateEof(ex); + } + finally + { + LeaveCancellationReadFrontier(frontierFlow); + registration.Dispose(); + if (timeoutSet) + SetRemainingTimeout(Timeout.InfiniteTimeSpan); } } - internal CurrentSegmentBuffer ExtendCurrentMessage() + internal CurrentMessageBuffer ExtendCurrentMessage() { var timeoutSet = false; try @@ -277,7 +387,7 @@ internal CurrentSegmentBuffer ExtendCurrentMessage() } } - void OnHeartbeat(TimeSpan elapsed) + internal void OnHeartbeat(TimeSpan elapsed) { var ticks = Interlocked.Exchange(ref _remainingTimeoutTicks, ClaimedTimeoutTicks); if (ticks == ClaimedTimeoutTicks) @@ -394,25 +504,24 @@ public ValueTask MoveNextAsync(CancellationToken cancellationToken = defau return new(true); } - if (pipe.TryMoveNextBatch(out var completed)) - continue; - if (completed) - return new(ReadCompleted()); + PrepareRead(); var readToken = _cancellationTokenSource.Token; var frontierFlow = EnterCancellationReadFrontier(); try { retryRead: - if (pipe.TryBeginDirectRead(readToken, out var directReadTask)) + if (TryBeginDirectRead(readToken, out var directReadTask)) { try { while (true) { if (!directReadTask.IsCompletedSuccessfully) - return MoveNextAsyncCore(null, directReadTask, null, cancellationToken, frontierFlow); - if (pipe.CompleteDirectRead(directReadTask.Result, readToken, out directReadTask, out var readFinished, out var directReadCompleted)) + return AwaitDirectRead(directReadTask, cancellationToken, frontierFlow); + if (CompleteDirectRead(directReadTask.Result, readToken, + out directReadTask, out var readFinished, + out var directReadCompleted)) break; if (!readFinished) continue; @@ -428,19 +537,12 @@ public ValueTask MoveNextAsync(CancellationToken cancellationToken = defau } catch { - pipe.AbortDirectRead(); + AbortDirectRead(); throw; } } - var readTask = pipe.ReadAsync(readToken); - if (!readTask.IsCompletedSuccessfully) - return MoveNextAsyncCore(readTask, null, null, cancellationToken, frontierFlow); - LeaveCancellationReadFrontier(frontierFlow); - if (pipe.TryMoveNextBatch(readTask.Result, _cancellationTokenSource.Token, out var readCompleted)) - continue; - if (readCompleted) - return new(ReadCompleted()); + return BeginPipeRead(pipe, readToken, cancellationToken, frontierFlow); } catch (Exception ex) when (_cancellationTokenSource.IsCancellationRequested) { @@ -461,8 +563,140 @@ public ValueTask MoveNextAsync(CancellationToken cancellationToken = defau [MethodImpl(MethodImplOptions.NoInlining)] - async ValueTask MoveNextAsyncCore(ValueTask? readTask, ValueTask? directReadTask, ValueTask? messageHandledTask, CancellationToken cancellationToken, PgClientFlow? frontierFlow = null) + ValueTask BeginPipeRead( + ProtocolReadPipe pipe, + CancellationToken readToken, + CancellationToken cancellationToken, + PgClientFlow frontierFlow) + => MoveNextAsyncCore( + pipe.ReadAsync(readToken), null, null, + cancellationToken, frontierFlow); + + [MethodImpl(MethodImplOptions.NoInlining)] + ValueTask AwaitDirectRead( + ValueTask directReadTask, + CancellationToken cancellationToken, + PgClientFlow frontierFlow) + => cancellationToken.CanBeCanceled + ? MoveNextDirectWithCancellationAsync(directReadTask, cancellationToken, frontierFlow) + : MoveNextDirectAsync(directReadTask, frontierFlow); + + [MethodImpl(MethodImplOptions.NoInlining)] + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] + async ValueTask MoveNextDirectAsync( + ValueTask directReadTask, + PgClientFlow frontierFlow) { + var timeoutSet = false; + try + { + while (true) + { + try + { + if (!timeoutSet) + { + ArmReadTimeout(); + timeoutSet = true; + } + var length = await directReadTask.ConfigureAwait(false); + if (CompleteDirectRead(length, + _cancellationTokenSource.Token, + out var nextDirectRead, out var readFinished, + out var readCompleted)) + { + LeaveCancellationReadFrontier(frontierFlow); + frontierFlow = null!; + } + else if (!readFinished) + { + directReadTask = nextDirectRead; + continue; + } + else + { + LeaveCancellationReadFrontier(frontierFlow); + frontierFlow = null!; + if (readCompleted) + return ReadCompleted(); + } + } + catch (Exception ex) + { + AbortDirectRead(); + if (frontierFlow is not null) + { + LeaveCancellationReadFrontier(frontierFlow); + frontierFlow = null!; + } + if (_cancellationTokenSource.IsCancellationRequested) + throw TranslateReadCancellation(ex, default); + if (ex is EndOfStreamException eof) + throw TranslateEof(eof); + throw; + } + + while (TryMoveNext(_pipe)) + { + var handleTask = CurrentExecutionControl.HandleMessageAuto(_pipe.Current); + if (!handleTask.IsCompletedSuccessfully) + { + if (!await handleTask.ConfigureAwait(false)) + return true; + continue; + } + if (!handleTask.Result) + return true; + } + + PrepareRead(); + var token = _cancellationTokenSource.Token; + frontierFlow = EnterCancellationReadFrontier(); + if (!TryBeginDirectRead(token, out directReadTask)) + return await MoveNextAsyncCore( + _pipe.ReadAsync(token), null, null, + default, frontierFlow).ConfigureAwait(false); + } + } + finally + { + if (frontierFlow is not null) + LeaveCancellationReadFrontier(frontierFlow); + if (timeoutSet) + SetRemainingTimeout(Timeout.InfiniteTimeSpan); + } + } + + [MethodImpl(MethodImplOptions.NoInlining)] + async ValueTask MoveNextDirectWithCancellationAsync( + ValueTask directReadTask, + CancellationToken cancellationToken, + PgClientFlow frontierFlow) + { + var registration = cancellationToken.UnsafeRegister( + static (state, _) => ((CancellationTokenSource)state!).Cancel(), + _cancellationTokenSource); + try + { + return await MoveNextDirectAsync(directReadTask, frontierFlow).ConfigureAwait(false); + } + catch (Exception) when (cancellationToken.IsCancellationRequested) + { + throw new OperationCanceledException(cancellationToken); + } + finally + { + registration.Dispose(); + } + } + + [MethodImpl(MethodImplOptions.NoInlining)] + async ValueTask MoveNextAsyncCore(ValueTask? readTask, + ValueTask? directReadTask, ValueTask? messageHandledTask, + CancellationToken cancellationToken, PgClientFlow? frontierFlow = null) + { + var pipe = _pipe; var timeoutSet = false; var registration = cancellationToken.UnsafeRegister(static (state, _) => ((CancellationTokenSource)state!).Cancel(), _cancellationTokenSource); try @@ -489,7 +723,9 @@ async ValueTask MoveNextAsyncCore(ValueTask? readTask, ValueTa LeaveCancellationReadFrontier(frontierFlow!); frontierFlow = null; readTask = null; - if (_pipe.TryMoveNextBatch(result, _cancellationTokenSource.Token, out var readCompleted)) + if (CompleteRead( + result, _cancellationTokenSource.Token, + out var readCompleted)) continue; if (readCompleted) return ReadCompleted(); @@ -514,7 +750,10 @@ async ValueTask MoveNextAsyncCore(ValueTask? readTask, ValueTa timeoutSet = true; } var length = await pendingDirectRead.ConfigureAwait(false); - if (_pipe.CompleteDirectRead(length, _cancellationTokenSource.Token, out var nextDirectRead, out var readFinished, out var readCompleted)) + if (CompleteDirectRead(length, + _cancellationTokenSource.Token, + out var nextDirectRead, out var readFinished, + out var readCompleted)) { LeaveCancellationReadFrontier(frontierFlow!); frontierFlow = null; @@ -534,7 +773,7 @@ async ValueTask MoveNextAsyncCore(ValueTask? readTask, ValueTa } catch (Exception ex) { - _pipe.AbortDirectRead(); + AbortDirectRead(); if (frontierFlow is not null) { LeaveCancellationReadFrontier(frontierFlow); @@ -548,9 +787,9 @@ async ValueTask MoveNextAsyncCore(ValueTask? readTask, ValueTa } } - while (TryMoveNext(_pipe)) + while (TryMoveNext(pipe)) { - var handleTask = CurrentExecutionControl.HandleMessageAuto(_pipe.Current); + var handleTask = CurrentExecutionControl.HandleMessageAuto(pipe.Current); if (!handleTask.IsCompletedSuccessfully) { messageHandledTask = handleTask; @@ -562,19 +801,16 @@ async ValueTask MoveNextAsyncCore(ValueTask? readTask, ValueTa if (messageHandledTask.HasValue) continue; - if (_pipe.TryMoveNextBatch(out var completed)) - continue; - if (completed) - return ReadCompleted(); + PrepareRead(); try { var token = _cancellationTokenSource.Token; frontierFlow = EnterCancellationReadFrontier(); - if (_pipe.TryBeginDirectRead(token, out var nextDirectRead)) + if (TryBeginDirectRead(token, out var nextDirectRead)) directReadTask = nextDirectRead; else - readTask = _pipe.ReadAsync(token); + readTask = pipe.ReadAsync(token); } catch (Exception ex) when (_cancellationTokenSource.IsCancellationRequested) { throw TranslateReadCancellation(ex, cancellationToken); } @@ -626,6 +862,21 @@ public BackendMessage Current get => _pipe.Current; } + internal BackendMessage.Accessor CurrentAccessor + { + [MethodImpl(MethodImplOptions.AggressiveInlining)] + get => _pipe.CurrentAccessor; + } + + internal ReadOnlyMemory CurrentBufferedBody + { + [MethodImpl(MethodImplOptions.AggressiveInlining)] + get => _pipe.CurrentBufferedBody; + } + + internal PgTypes.BackendType CurrentType => _pipe.CurrentType; + internal bool CurrentBuffered => _pipe.CurrentBuffered; + [MethodImpl(MethodImplOptions.AggressiveInlining)] public bool TryGetCurrent(out BackendMessage message) { @@ -686,47 +937,60 @@ bool TryMoveNext(ProtocolReadPipe pipe) internal bool TryMoveNext() => TryMoveNextCore(); bool TryMoveNextCore() + { + if (!_pipe.TryPeekNext(out var header)) + { + PrepareRead(); + return false; + } + + var type = header.Type; + // Only auto-handled messages need the transactional peek slot: their handler may need + // I/O and decline the synchronous path. Ordinary messages can publish directly. + if (PgClientFlow.ExecutionControl.ShouldHandle(type)) + return TryMoveNextAutoHandled(type); + + _pipe.PublishPeeked(); + if (type is PgTypes.BackendType.ErrorResponse) + ObserveMessage(_pipe.Current); + return true; + } + + [MethodImpl(MethodImplOptions.NoInlining)] + bool TryMoveNextAutoHandled(PgTypes.BackendType type) { while (true) { - while (_pipe.TryPeekNextType(out var type)) + Debug.Assert(PgClientFlow.ExecutionControl.ShouldHandle(type)); + var handled = false; + if (type is PgTypes.BackendType.ReadyForQuery) + RestoreDefaultReadTimeout(); + if (!CurrentExecutionControl.TryHandleKnownMessage(_pipe.Peeked, out handled)) + return false; + + _pipe.PublishPeeked(); + if (!handled) { - // Only auto-handled messages need the transactional peek slot: their handler may need - // I/O and decline the synchronous path. Ordinary messages can publish directly. - if (type is not (PgTypes.BackendType.ReadyForQuery - or PgTypes.BackendType.NoticeResponse - or PgTypes.BackendType.NotificationResponse - or PgTypes.BackendType.ParameterStatus)) - { - var moved = TryMoveNext(_pipe); - Debug.Assert(moved); - return true; - } - - if (!_pipe.TryPeekNext(out _)) - break; - var handled = false; - if (type is PgTypes.BackendType.ReadyForQuery) - RestoreDefaultReadTimeout(); - if (!CurrentExecutionControl.TryHandleMessage(_pipe.Peeked, out handled)) - { - goto unavailable; - } - TryMoveNext(_pipe); - if (handled) - continue; + if (type is PgTypes.BackendType.ErrorResponse) + ObserveMessage(_pipe.Current); return true; } - // The current batch is exhausted. Descend through any bytes the PipeReader already owns - // before reporting unavailable; only a genuinely pending physical read should make the - // async caller install its continuation tree. - if (!_pipe.TryMoveNextBatch(out _)) - break; - } + if (!_pipe.TryPeekNext(out var header)) + { + PrepareRead(); + return false; + } - unavailable: - return false; + type = header.Type; + if (!PgClientFlow.ExecutionControl.ShouldHandle(type)) + { + _pipe.PublishPeeked(); + if (type is PgTypes.BackendType.ErrorResponse) + ObserveMessage(_pipe.Current); + return true; + } + } } // Auto-switch read, mirroring the encoder's FlushAuto: a sync flow takes the BLOCKING read path @@ -749,6 +1013,7 @@ public ValueTask GetNextAsync() return default; } + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] async ValueTask GetNextAsyncCore(ValueTask task) { @@ -805,12 +1070,11 @@ bool MoveNextSlow() { LeaveCancellationReadFrontier(frontierFlow); } - pipe.CommitBatch(); if (!success) return ReadCompleted(); if (!TryMoveNext(pipe)) - ThrowHelper.ThrowInvalidOperation("No message in a new batch"); + continue; } // HandleMessageAuto is always sync-completing (every branch returns a diff --git a/Slon/Pg/Protocol/PgEncoder.cs b/Slon/Pg/Protocol/PgEncoder.cs index 31c4d30..4ff7e55 100644 --- a/Slon/Pg/Protocol/PgEncoder.cs +++ b/Slon/Pg/Protocol/PgEncoder.cs @@ -1,4 +1,6 @@ +using System.Buffers.Binary; using System.Collections.Immutable; +using System.Diagnostics; using System.Runtime.CompilerServices; using System.Text; using Slon.Runtime; @@ -233,6 +235,82 @@ public void WriteBind(EncodedCString commandName) _writer.WriteUShort(1); // all binary } + // Executes a prepared statement without parameters and with default result formats: Bind on the + // unnamed portal, an optional portal Describe, an optional Execute, then syncCount Sync messages, + // all written through one reserved span. A flow appends at most one Sync of its own beside the + // command's, so syncCount is bounded to two. + internal void WritePreparedExecution(EncodedCString commandName, bool describe, bool execute, int syncCount) + { + WritePreparedExecutionCore(_writer, ClientEncoding, commandName, describe, execute, syncCount); + _executionControl.OnMessageWrite(FrontendType.Bind); + if (describe) + _executionControl.OnMessageWrite(FrontendType.Describe); + if (execute) + _executionControl.OnMessageWrite(FrontendType.Execute); + for (var i = 0; i < syncCount; i++) + _executionControl.OnMessageWrite(FrontendType.Sync); + } + + // The complete sequence is sized before its span is acquired and published only after every + // message has been filled. Incremental message tracking remains for streaming writers. + internal static void WritePreparedExecutionCore(ProtocolDataWriter writer, Encoding encoding, + EncodedCString commandName, bool describe, bool execute, int syncCount) + { + ArgumentOutOfRangeException.ThrowIfNegative(syncCount); + ArgumentOutOfRangeException.ThrowIfGreaterThan(syncCount, 2); + + const int header = sizeof(byte) + sizeof(uint); + const int describeBody = sizeof(byte) + 1; // 'P' and the unnamed portal + const int executeBody = sizeof(byte) + sizeof(int); // unnamed portal, all rows + var commandNameBytes = commandName.AsNullTerminatedSpan(encoding); + var bindBody = checked(1 + commandNameBytes.Length + 4 * sizeof(ushort)); + var total = checked(header + bindBody + + (describe ? header + describeBody : 0) + + (execute ? header + executeBody : 0) + + syncCount * header); + var span = writer.GetCompleteMessagesSpan(total).Slice(0, total); + + WriteHeader(span, FrontendType.Bind, bindBody); + span[header] = 0; // unnamed portal + commandNameBytes.CopyTo(span.Slice(header + 1)); + var formats = span.Slice(header + 1 + commandNameBytes.Length); + BinaryPrimitives.WriteUInt16BigEndian(formats, 0); // parameter format codes + BinaryPrimitives.WriteUInt16BigEndian(formats.Slice(2), 0); // parameters + BinaryPrimitives.WriteUInt16BigEndian(formats.Slice(4), 1); // result format codes + BinaryPrimitives.WriteUInt16BigEndian(formats.Slice(6), 1); // all binary + span = span.Slice(header + bindBody); + + if (describe) + { + WriteHeader(span, FrontendType.Describe, describeBody); + span[header] = (byte)'P'; + span[header + 1] = 0; + span = span.Slice(header + describeBody); + } + + if (execute) + { + WriteHeader(span, FrontendType.Execute, executeBody); + span[header] = 0; // unnamed portal + BinaryPrimitives.WriteUInt32BigEndian(span.Slice(header + 1), 0); // all rows + span = span.Slice(header + executeBody); + } + + for (var i = 0; i < syncCount; i++) + { + WriteHeader(span, FrontendType.Sync, 0); + span = span.Slice(header); + } + Debug.Assert(span.IsEmpty); + writer.AdvanceCompleteMessages(total); + + static void WriteHeader(Span span, FrontendType type, int bodyLength) + { + span[0] = type.ToByte(); + BinaryPrimitives.WriteUInt32BigEndian(span.Slice(1), checked((uint)(sizeof(uint) + bodyLength))); + } + } + public void WriteBind(EncodedCString commandName = default, EncodedCString portalName = default, ParameterSource parameters = default, ImmutableArray resultFormats = default) { diff --git a/Slon/Pg/Protocol/PgFlowSourceDriver.cs b/Slon/Pg/Protocol/PgFlowSourceDriver.cs index 1a72561..55f7799 100644 --- a/Slon/Pg/Protocol/PgFlowSourceDriver.cs +++ b/Slon/Pg/Protocol/PgFlowSourceDriver.cs @@ -8,6 +8,8 @@ namespace Slon.Pg.Protocol; /// sealed class PgFlowSourceDriver { + static readonly Action RunAction = static state => ((PgFlowSourceDriver)state!).Run(); + readonly PgClientFlowSource.State _source; readonly SourceWakeEvent _wakeEvent; readonly Action _signalHeldSyncFlow; @@ -53,7 +55,7 @@ public void Drive(bool runContinuationsAsynchronously) return; if (runContinuationsAsynchronously) - _wakeEvent.Scheduler.SubmitDetached(static driver => driver.Run(), this); + _wakeEvent.Scheduler.SubmitDetached(RunAction, (object?)this); else Run(); } @@ -114,7 +116,7 @@ void OnWaitReady(SourceWakeEvent.WaitReadyContext context) } void ScheduleRun() - => _wakeEvent.Scheduler.SubmitDetached(static driver => driver.Run(), this); + => _wakeEvent.Scheduler.SubmitDetached(RunAction, (object?)this); void Run() { @@ -149,7 +151,7 @@ void Run() } if (transfer) - _wakeEvent.Scheduler.SubmitDetached(static driver => driver.Run(), this); + _wakeEvent.Scheduler.SubmitDetached(RunAction, (object?)this); return; } } diff --git a/Slon/Pg/Protocol/ProtocolDataWriter.cs b/Slon/Pg/Protocol/ProtocolDataWriter.cs index 78bb332..d5d251b 100644 --- a/Slon/Pg/Protocol/ProtocolDataWriter.cs +++ b/Slon/Pg/Protocol/ProtocolDataWriter.cs @@ -90,6 +90,10 @@ internal object GetParameterWriterState(ParameterWriter writer) public Memory GetMemory(int sizeHint = 0) => _pipe.GetMemory(sizeHint); public Span GetSpan(int sizeHint = 0) => _pipe.GetSpan(sizeHint); public void Advance(int count) => _pipe.Advance(count); + internal Span GetCompleteMessagesSpan(int totalLength) + => _pipe.GetCompleteMessagesSpan(totalLength); + internal void AdvanceCompleteMessages(int totalLength) + => _pipe.AdvanceCompleteMessages(totalLength); internal const long UnflushedBytesFlushThreshold = ProtocolWritePipe.UnflushedBytesFlushThreshold; diff --git a/Slon/Pg/Protocol/ProtocolReadPipe.cs b/Slon/Pg/Protocol/ProtocolReadPipe.cs index 29e7bc0..9362006 100644 --- a/Slon/Pg/Protocol/ProtocolReadPipe.cs +++ b/Slon/Pg/Protocol/ProtocolReadPipe.cs @@ -1,107 +1,489 @@ +using System.Buffers; +using System.Diagnostics; using System.IO.Pipelines; using Slon.Pipelines; namespace Slon.Pg.Protocol; -// Shared per-protocol read-side wire state. One instance per protocol, behind any number of -// PgDecoder shells (the base protocol shell plus a per-exclusive-scope shell). The single-pump -// invariant means only one shell ever drives this pipe at a time, so the batch enumerator and -// message context are safe to share. Token-bearing concerns (CTS, abort translation, read-timeout -// countdown, CurrentExecutionControl, and framing/handler loops) live in the shell; the pipe -// exposes only the raw read primitives and message iteration. -sealed class ProtocolReadPipe(PipeSegmentEnumerator messageBatchEnumerator) +// Shared per-protocol read-side wire state. One cursor parses each backend message once, and an +// incomplete message directly drives the next PipeReader grant. +sealed class ProtocolReadPipe( + PipeReader reader, int dataRowStreamingThreshold, bool ownsReader = false) { + enum PendingRead : byte { None, Messages, Slide, Extend } + readonly BackendMessageContext _messageContext = new(); + ReadOnlySequence _activeBuffer; + SequencePosition _examined; + SequencePosition _retainedStart; + // Offset of _retainedStart within _activeBuffer. Reset to zero whenever a new PipeReader grant + // begins at that position; non-zero only while advancing the retention origin within one grant. + long _retainedOffset; + long _currentMessageOffset; + long _currentMessageLength = -1; + long _pendingCursorOffset; + long _pendingSkipLength; + int _minimumReadSize; + PendingRead _pendingRead; + bool _hasActiveRead; + bool _retainsResult; + public PipeReader PipeReader => reader; public BackendMessage Current => _messageContext.Current; + public BackendMessage.Accessor CurrentAccessor => _messageContext.CurrentAccessor; public bool CurrentIsError => _messageContext.CurrentIsError; - public bool TryGetCurrent(out BackendMessage message) => _messageContext.TryGetCurrent(out message); + public PgTypes.BackendType CurrentType => _messageContext.CurrentType; + public bool CurrentBuffered => _messageContext.CurrentBuffered; + public ReadOnlyMemory CurrentBufferedBody => _messageContext.CurrentBufferedBody; + public bool TryGetCurrent(out BackendMessage message) + => _messageContext.TryGetCurrent(out message); public bool TryMoveNext() => _messageContext.TryMoveNext(); - public bool TryPeekNextType(out PgTypes.BackendType type) => _messageContext.TryPeekNextType(out type); - public bool TryPeekNext(out BackendHeader header) => _messageContext.TryPeekNext(out header); + public bool TryPeekNext(out BackendHeader header) + => _messageContext.TryPeekNext(out header); + public void PublishPeeked() => _messageContext.PublishPeeked(); public BackendMessage Peeked => _messageContext.Peeked; public void BindDecoder(PgDecoder decoder) => _messageContext.BindDecoder(decoder); - public bool TryMoveNextBatch(out bool completed) + public void PrepareRead() { - _messageContext.RetireCurrentBatch(); - if (!messageBatchEnumerator.TryMoveNext(out completed)) - return false; - CommitBatch(); - return true; + if (_pendingRead is PendingRead.Messages) + return; + if (_pendingRead is not PendingRead.None) + ThrowHelper.ThrowInvalidOperation( + "The current message still has a pending read."); + + var retainsResult = _retainsResult; + + if (!_hasActiveRead) + { + _pendingCursorOffset = 0; + _minimumReadSize = BackendHeader.ByteCount; + _pendingRead = PendingRead.Messages; + return; + } + + if (_currentMessageLength > 0) + { + PrepareAfterPartialMessage(retainsResult); + return; + } + + if (!_messageContext.TryGetReadRequirement( + out var cursorConsumedLength, out var requiredLength)) + ThrowHelper.ThrowInvalidOperation( + "The current backend-message cursor has not been exhausted."); + + var unreadOffset = checked(_pendingCursorOffset + cursorConsumedLength); + var retainedOffset = retainsResult ? _retainedOffset : 0; + var unread = _activeBuffer.GetPosition(unreadOffset); + _pendingCursorOffset = retainsResult + ? checked(unreadOffset - retainedOffset) + : 0; + _messageContext.RetireCursor( + retainProjections: retainsResult); + reader.AdvanceTo( + retainsResult ? _retainedStart : unread, _examined); + _hasActiveRead = false; + _activeBuffer = default; + _retainedOffset = 0; + _currentMessageLength = -1; + _currentMessageOffset = 0; + _minimumReadSize = int.CreateSaturating(requiredLength); + _pendingRead = PendingRead.Messages; + } + + void PrepareAfterPartialMessage(bool retainsResult) + { + var current = _activeBuffer.Slice(_currentMessageOffset); + _messageContext.RetireCursor( + retainProjections: retainsResult); + if (retainsResult) + { + _pendingCursorOffset = checked( + _currentMessageOffset + _currentMessageLength - _retainedOffset); + var examined = current.Length >= _currentMessageLength + ? current.GetPosition(_currentMessageLength) + : _examined; + reader.AdvanceTo(_retainedStart, examined); + _minimumReadSize = int.CreateSaturating( + _pendingCursorOffset + BackendHeader.ByteCount); + } + else if (current.Length >= _currentMessageLength) + { + var unread = current.GetPosition(_currentMessageLength); + reader.AdvanceTo(unread, unread); + _pendingCursorOffset = 0; + _minimumReadSize = BackendHeader.ByteCount; + } + else + { + _pendingSkipLength = _currentMessageLength - current.Length; + reader.AdvanceTo(_activeBuffer.End, _examined); + _pendingCursorOffset = _pendingSkipLength; + _minimumReadSize = int.CreateSaturating( + _pendingSkipLength + BackendHeader.ByteCount); + } + + _hasActiveRead = false; + _activeBuffer = default; + _retainedOffset = 0; + _currentMessageLength = -1; + _currentMessageOffset = 0; + _pendingRead = PendingRead.Messages; } public ValueTask ReadAsync(CancellationToken cancellationToken) - => messageBatchEnumerator.ReadAsync(cancellationToken); + => _minimumReadSize > 0 + ? reader.ReadAtLeastAsync(_minimumReadSize, cancellationToken) + : reader.ReadAsync(cancellationToken); - public bool TryBeginDirectRead(CancellationToken cancellationToken, out ValueTask task) - => messageBatchEnumerator.TryBeginDirectRead(cancellationToken, out task); + public bool CompleteRead( + in ReadResult result, CancellationToken cancellationToken, + out bool completed) + { + if (_pendingRead is not PendingRead.Messages) + ThrowHelper.ThrowInvalidOperation("No protocol read is pending."); + _pendingRead = PendingRead.None; + _minimumReadSize = 0; + if (result.IsCanceled) + ThrowHelper.ThrowOperationCanceled(cancellationToken); + if (result.Buffer.IsEmpty && result.IsCompleted) + { + completed = true; + _messageContext.RetireCursor(); + return false; + } + if (result.Buffer.IsEmpty) + { + completed = false; + return false; + } + + _activeBuffer = result.Buffer; + _examined = result.Buffer.End; + _retainedStart = result.Buffer.Start; + _retainedOffset = 0; + _hasActiveRead = true; + if (_pendingSkipLength > 0) + { + if (result.Buffer.Length < _pendingSkipLength) + { + if (result.IsCompleted) + throw new EndOfStreamException( + "The pipe completed within a backend message."); + _pendingSkipLength -= result.Buffer.Length; + reader.AdvanceTo(result.Buffer.End, result.Buffer.End); + _hasActiveRead = false; + _activeBuffer = default; + _pendingCursorOffset = _pendingSkipLength; + _minimumReadSize = int.CreateSaturating( + _pendingSkipLength + BackendHeader.ByteCount); + _pendingRead = PendingRead.Messages; + completed = false; + return false; + } + _pendingCursorOffset = _pendingSkipLength; + _pendingSkipLength = 0; + } + _currentMessageOffset = _pendingCursorOffset; + var cursorBuffer = _pendingCursorOffset is 0 + ? result.Buffer + : result.Buffer.Slice(_pendingCursorOffset); + if (cursorBuffer.IsEmpty) + { + completed = result.IsCompleted; + if (completed && !_retainsResult) + _messageContext.RetireCursor(); + if (!completed) + { + reader.AdvanceTo( + _retainsResult ? _retainedStart : result.Buffer.End, + result.Buffer.End); + _hasActiveRead = false; + _activeBuffer = default; + if (!_retainsResult) + _pendingCursorOffset = 0; + _minimumReadSize = _retainsResult + ? int.CreateSaturating( + _pendingCursorOffset + BackendHeader.ByteCount) + : BackendHeader.ByteCount; + _pendingRead = PendingRead.Messages; + } + return false; + } + _currentMessageLength = -1; + var cursor = new BackendMessageCursor( + cursorBuffer, dataRowStreamingThreshold); + _messageContext.SetCursor(cursor); + completed = false; + return true; + } - public bool CompleteDirectRead(int length, CancellationToken cancellationToken, out ValueTask next, out bool readFinished, out bool completed) + public bool TrySlideCurrentMessage( + SequencePosition consumed, long consumedLength, + out CurrentMessageBuffer result) { - if (!messageBatchEnumerator.CompleteDirectRead(length, cancellationToken, out next, out readFinished, out completed)) + PrepareCurrentMessageRead(consumed, consumedLength, PendingRead.Slide); + if (!reader.TryRead(out var read)) + { + result = default; return false; - CommitBatch(); + } + result = CompleteCurrentMessageRead(read); return true; } - public void AbortDirectRead() => messageBatchEnumerator.AbortDirectRead(); + public ValueTask BeginSlideCurrentMessageAsync( + SequencePosition consumed, long consumedLength, + CancellationToken cancellationToken) + { + PrepareCurrentMessageRead(consumed, consumedLength, PendingRead.Slide); + return reader.ReadAsync(cancellationToken); + } - public bool TryMoveNextBatch(ReadResult result, CancellationToken cancellationToken, out bool completed) + public CurrentMessageBuffer SlideCurrentMessage( + SequencePosition consumed, long consumedLength, TimeSpan timeout) { - _messageContext.RetireCurrentBatch(); - if (!messageBatchEnumerator.TryMoveNext(result, cancellationToken, out completed)) + if (reader is not StreamPipeReader syncReader) + throw new NotSupportedException( + "Underlying pipe reader does not support synchronous reads."); + PrepareCurrentMessageRead(consumed, consumedLength, PendingRead.Slide); + return CompleteCurrentMessageRead(syncReader.Read(timeout)); + } + public bool TryExtendCurrentMessage(out CurrentMessageBuffer result) + { + PrepareCurrentMessageRead( + _retainedStart, consumedLength: 0, PendingRead.Extend); + if (!reader.TryRead(out var read)) + { + result = default; return false; - CommitBatch(); + } + result = CompleteCurrentMessageRead(read); return true; } - public bool TryContinueCurrentMessage(SequencePosition consumed, long consumedLength, out CurrentSegmentBuffer result) - => messageBatchEnumerator.TryContinueCurrentSegment(consumed, consumedLength, out result); + public ValueTask BeginExtendCurrentMessageAsync( + CancellationToken cancellationToken) + { + PrepareCurrentMessageRead( + _retainedStart, consumedLength: 0, PendingRead.Extend); + return reader.ReadAsync(cancellationToken); + } - public ValueTask ContinueCurrentMessageAsync( - SequencePosition consumed, long consumedLength, CancellationToken cancellationToken) - => messageBatchEnumerator.ContinueCurrentSegmentAsync(consumed, consumedLength, cancellationToken); + public ValueTask BeginBufferCurrentMessageAsync( + CancellationToken cancellationToken) + { + PrepareCurrentMessageRead( + _retainedStart, consumedLength: 0, PendingRead.Extend); + var requiredLength = checked(_currentMessageOffset + _currentMessageLength); + return requiredLength <= int.MaxValue + ? reader.ReadAtLeastAsync((int)requiredLength, cancellationToken) + : reader.ReadAsync(cancellationToken); + } - public CurrentSegmentBuffer ContinueCurrentMessage( - SequencePosition consumed, long consumedLength, TimeSpan timeout) - => messageBatchEnumerator.ContinueCurrentSegment(consumed, consumedLength, timeout); + public CurrentMessageBuffer ExtendCurrentMessage(TimeSpan timeout) + { + if (reader is not StreamPipeReader syncReader) + throw new NotSupportedException( + "Underlying pipe reader does not support synchronous reads."); + PrepareCurrentMessageRead( + _retainedStart, consumedLength: 0, PendingRead.Extend); + return CompleteCurrentMessageRead(syncReader.Read(timeout)); + } - public bool TryExtendCurrentMessage(out CurrentSegmentBuffer result) - => messageBatchEnumerator.TryExtendCurrentSegment(out result); + void PrepareCurrentMessageRead( + SequencePosition consumed, long consumedLength, PendingRead mode) + { + if (_pendingRead == mode) + return; + if (_pendingRead is not PendingRead.None || !_hasActiveRead + || _currentMessageLength <= 0) + ThrowHelper.ThrowInvalidOperation( + "The current message is not awaiting more data."); + if (mode is PendingRead.Slide + && (consumedLength <= 0 || consumedLength >= _currentMessageLength)) + throw new ArgumentOutOfRangeException(nameof(consumedLength)); + + var retainedOffset = _retainsResult ? _retainedOffset : 0; + reader.AdvanceTo( + mode is PendingRead.Slide && !_retainsResult + ? consumed + : _retainedStart, + _examined); + if (mode is PendingRead.Slide) + { + if (_retainsResult) + { + _currentMessageOffset = checked( + _activeBuffer.Slice(0, consumed).Length - retainedOffset); + } + else + { + _retainedStart = consumed; + _currentMessageOffset = 0; + } + _currentMessageLength -= consumedLength; + } + else if (_retainsResult) + { + _currentMessageOffset = checked( + _currentMessageOffset - retainedOffset); + } + _retainedOffset = 0; + _hasActiveRead = false; + _activeBuffer = default; + _pendingRead = mode; + } - public ValueTask ExtendCurrentMessageAsync(CancellationToken cancellationToken) - => messageBatchEnumerator.ExtendCurrentSegmentAsync(cancellationToken); + public CurrentMessageBuffer CompleteCurrentMessageRead( + in ReadResult result, CancellationToken cancellationToken = default) + { + if (_pendingRead is not (PendingRead.Slide or PendingRead.Extend)) + ThrowHelper.ThrowInvalidOperation( + "No current-message read is pending."); + _pendingRead = PendingRead.None; + if (result.IsCanceled) + ThrowHelper.ThrowOperationCanceled(cancellationToken); - public CurrentSegmentBuffer ExtendCurrentMessage(TimeSpan timeout) - => messageBatchEnumerator.ExtendCurrentSegment(timeout); + _activeBuffer = result.Buffer; + _examined = result.Buffer.End; + _retainedStart = result.Buffer.Start; + _hasActiveRead = true; + var current = result.Buffer.Slice(_currentMessageOffset); + if (result.IsCompleted && current.Length < _currentMessageLength) + throw new EndOfStreamException( + "The pipe completed within a backend message."); + var complete = current.Length >= _currentMessageLength; + var buffer = complete + ? current.Slice(0, _currentMessageLength) + : current; + _examined = complete ? buffer.End : result.Buffer.End; + return new(buffer, complete); + } - public ValueTask MoveNextAsync(CancellationToken cancellationToken) + public async ValueTask MoveNextAsync( + CancellationToken cancellationToken) { - _messageContext.RetireCurrentBatch(); - return messageBatchEnumerator.MoveNextAsync(cancellationToken); + while (true) + { + PrepareRead(); + var read = await ReadAsync(cancellationToken).ConfigureAwait(false); + if (CompleteRead( + read, cancellationToken, out var completed)) + return true; + if (completed) + return false; + } } public bool MoveNext(TimeSpan timeout) { - _messageContext.RetireCurrentBatch(); - return messageBatchEnumerator.MoveNext(timeout); + if (reader is not StreamPipeReader syncReader) + throw new NotSupportedException( + "Underlying pipe reader does not support synchronous reads."); + while (true) + { + PrepareRead(); + var read = _minimumReadSize > 0 + ? syncReader.ReadAtLeast(_minimumReadSize, timeout) + : syncReader.Read(timeout); + if (CompleteRead( + read, CancellationToken.None, out var completed)) + return true; + if (completed) + return false; + } + } + + public void SetCurrentMessageLength(long messageLength) + { + ArgumentOutOfRangeException.ThrowIfNegativeOrZero(messageLength); + _currentMessageLength = messageLength; + } + + public void CompleteCurrentMessage() + => _currentMessageLength = -1; + + public void EnableResultRetention() + { + if (!_hasActiveRead || _pendingRead is not PendingRead.None) + ThrowHelper.ThrowInvalidOperation( + "Result retention requires an active backend message."); + if (_retainsResult) + { + _ = _messageContext.Current; + return; + } + // A prior result may have completed entirely within this PipeReader grant. Its borrowed + // values were released by EndResultRetention, so the next tenure can move the retained + // origin to its own first message. Rebase offsets when the eventual read starts there. + var retainedOffset = checked(_pendingCursorOffset + + _messageContext.CaptureCurrentMessageOffset()); + if (_currentMessageLength > 0) + { + // Before the next grant these describe the cursor base and the message end relative to + // that base. Dropping preceding messages makes the retained message its own base, so + // preserve only its actual length and carry its current absolute offset until AdvanceTo. + _currentMessageLength = checked(_currentMessageLength + - (retainedOffset - _currentMessageOffset)); + _currentMessageOffset = retainedOffset; + _messageContext.RebaseCurrentMessageOffset(); + } + _retainedStart = _activeBuffer.GetPosition(retainedOffset); + _retainedOffset = retainedOffset; + _retainsResult = true; } - // Publishes the just-read batch as the current batch the message context iterates. - public void CommitBatch() => _messageContext.SetBatch(messageBatchEnumerator.Current); + public void EndResultRetention() + { + _retainsResult = false; + // Keep the active cursor until ordinary message advancement exhausts it. Besides avoiding + // an examined-position rewind (not supported by every PipeReader), this lets a successor + // consume messages already present in the current grant. PrepareRead advances the retained + // prefix when it eventually needs another grant, matching the former batch reader. + _messageContext.ReleaseContiguousProjections(); + } + + // The retiring zero-edge owner calls this before its activation turn is released. End the + // current PipeReader grant at the parsed successor boundary so an idle connection cannot retain + // completed result storage. Any buffered suffix is deliberately un-examined and will be + // reacquired by the next ordinary read. + public void ReleaseReadBufferAtIdle() + { + Debug.Assert(!_retainsResult); + if (!_hasActiveRead || _pendingRead is not PendingRead.None + || _currentMessageLength > 0 + || !_messageContext.TryGetCursorConsumedLength(out var cursorConsumedLength)) + return; + + var consumedOffset = checked(_pendingCursorOffset + cursorConsumedLength); + var consumed = _activeBuffer.GetPosition(consumedOffset); + _messageContext.RetireCursor(); + reader.AdvanceTo(consumed, consumed); + _activeBuffer = default; + _examined = default; + _retainedStart = default; + _retainedOffset = 0; + _currentMessageOffset = 0; + _currentMessageLength = -1; + _pendingCursorOffset = 0; + _minimumReadSize = 0; + _hasActiveRead = false; + } public void Dispose() { - _messageContext.RetireCurrentBatch(); - messageBatchEnumerator.Dispose(); + _messageContext.RetireCursor(); + if (ownsReader) + reader.Complete(); } public ValueTask DisposeAsync() { - _messageContext.RetireCurrentBatch(); - return messageBatchEnumerator.DisposeAsync(); + _messageContext.RetireCursor(); + return ownsReader ? reader.CompleteAsync() : default; } } diff --git a/Slon/Pg/Protocol/ProtocolWritePipe.cs b/Slon/Pg/Protocol/ProtocolWritePipe.cs index 1d6515f..153ab2f 100644 --- a/Slon/Pg/Protocol/ProtocolWritePipe.cs +++ b/Slon/Pg/Protocol/ProtocolWritePipe.cs @@ -68,6 +68,25 @@ sealed class ProtocolWritePipe(IOutputWriter writer, Encoding clientEncoding, Ac internal Span GetSpan(int sizeHint = 0) => _bufferingWriter.GetSpan(sizeHint); internal void Advance(int count) => _bufferingWriter.Advance(count); + // A trusted encoder which has computed and filled a complete sequence can publish it once. + // Validate any preceding incrementally-written message before granting the span; until Advance + // nothing from the new sequence is visible to the underlying writer. + internal Span GetCompleteMessagesSpan(int totalLength) + { + ArgumentOutOfRangeException.ThrowIfNegativeOrZero(totalLength); + var unflushed = checked((int)_bufferingWriter.UnflushedBytes); + if (_messageLength is { } previous && unflushed + _messageBytesFlushed != previous) + ThrowUnderwritten(previous, unflushed + _messageBytesFlushed); + return _bufferingWriter.GetSpan(totalLength); + } + + internal void AdvanceCompleteMessages(int totalLength) + { + _bufferingWriter.Advance(totalLength); + _messageLength = null; + _messageBytesFlushed = 0; + } + // Validates the previous message, arms length tracking for the new one, then writes its // five-byte header directly into the buffered span. Keeping these together avoids a second // shell traversal and a temporary header copy. Mid-message flushes are handled by @@ -112,12 +131,13 @@ internal void StartMessage(int totalLength) // re-validates and commits exactly the bytes that remain. void CheckMessageBytesFlushed(int count) { - if (_messageLength is null) + if (!_messageLength.HasValue) return; // Pre-startup raw writes (e.g. StartupMessage's CopyStartupBuffer). + ref readonly var messageLength = ref Nullable.GetValueRefOrDefaultRef(in _messageLength); if (count < 0) throw new ArgumentOutOfRangeException(nameof(count)); - if ((long)_messageBytesFlushed + count > _messageLength) - ThrowOverwritten(_messageLength.Value, _messageBytesFlushed + (long)count); + if ((long)_messageBytesFlushed + count > messageLength) + ThrowOverwritten(messageLength, _messageBytesFlushed + (long)count); } // Commit the counter for bytes that actually left the buffer (the drained delta: before-after @@ -149,10 +169,11 @@ internal int CurrentMessagePaddingLength { get { - if (_messageLength is null) + if (!_messageLength.HasValue) return 0; + ref readonly var messageLength = ref Nullable.GetValueRefOrDefaultRef(in _messageLength); var unflushed = checked((int)_bufferingWriter.UnflushedBytes); - return Math.Max(0, _messageLength.Value - (unflushed + _messageBytesFlushed)); + return Math.Max(0, messageLength - (unflushed + _messageBytesFlushed)); } } diff --git a/Slon/Pg/Protocol/ReadyForQueryMessage.cs b/Slon/Pg/Protocol/ReadyForQueryMessage.cs index 977e4ce..8a4703c 100644 --- a/Slon/Pg/Protocol/ReadyForQueryMessage.cs +++ b/Slon/Pg/Protocol/ReadyForQueryMessage.cs @@ -15,14 +15,8 @@ public static ReadyForQueryMessage Create(in BackendMessage message) message.EnsureExpected(PgTypes.BackendType.ReadyForQuery); message.EnsureBuffered(); - byte status; - if (message.TryGetFirstSpan(0, out var body) && !body.IsEmpty) + if (!message.TryGetFirstByte(0, out var status)) { - status = body[0]; - } - else - { - status = 0; message.BodyReader.TryCopyTo(new Span(ref status)); } var transactionStatus = (TransactionStatus)status; diff --git a/Slon/Pg/Row.cs b/Slon/Pg/Row.cs index 970a465..f7f85cd 100644 --- a/Slon/Pg/Row.cs +++ b/Slon/Pg/Row.cs @@ -2,6 +2,7 @@ using System.Buffers.Binary; using System.Diagnostics; using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; using System.Text; using Slon.Buffers; using Slon.Pg.Protocol; @@ -19,7 +20,10 @@ internal Row() { } RowDescription _rowDescription = null!; BackendMessageBodyReader? _bodyReader; IColumnLease? _columnLease; + byte[]? _bufferedArray; ReadOnlyMemory _bufferedBody; + int _bufferedOffset; + int _bufferedLength; int _leasedOrdinal; int _lastBufferedOrdinal = -1; int _lastBufferedOffset; @@ -29,6 +33,10 @@ internal Row() { } int _columnOffset; BackendMessage Message => _messageAccessor.Message; + ReadOnlyMemory BufferedBody + => _bufferedArray is { } array + ? array.AsMemory(_bufferedOffset, _bufferedLength) + : _bufferedBody; [MethodImpl(MethodImplOptions.AggressiveInlining)] SequenceReader GetColumnReader(int ordinal, out int columnIndex, out int columnOffset) @@ -51,6 +59,32 @@ SequenceReader GetColumnReader(int ordinal, out int columnIndex, out int c public T GetValue(int ordinal) => GetValueCore(ordinal, textEncoding: null); + /// + /// Borrows the field's raw PostgreSQL representation as contiguous memory. + /// + /// + /// + /// The returned memory is a view over storage owned by the command result. Its lifetime is not + /// enforced: after the row enumerator advances or is disposed, accessing it may observe storage + /// that has been reused for unrelated data rather than throw. + /// + /// + /// Calling before row enumeration extends + /// the borrow until that command result is released. Copy the memory when it must outlive the + /// applicable boundary. + /// + /// + public ReadOnlyMemory BorrowFieldMemory(int ordinal) + { + RevokeColumnLease(); + EnsureBuffered(); + if (TryGetFieldMemory(ordinal, out var field)) + return Message.GetContiguousMemory(field); + + var sequence = GetFieldSequence(ordinal); + return Message.GetContiguousMemory(sequence); + } + // Bootstrap consumers have no serializer but must still bind text decoding to one negotiated // encoding snapshot for the lifetime of their operation. internal T GetValue(int ordinal, Encoding textEncoding) @@ -419,18 +453,19 @@ bool TryGetFieldMemory(int ordinal, out ReadOnlyMemory field) { if (ordinal == _lastBufferedOrdinal) { - field = _bufferedBody.Slice(_lastBufferedOffset, _lastBufferedLength); + field = BufferedBody.Slice(_lastBufferedOffset, _lastBufferedLength); return true; } var columnIndex = _column <= ordinal ? _column : 0; var columnOffset = _column <= ordinal ? _columnOffset : sizeof(short); - if ((uint)columnOffset > (uint)_bufferedBody.Length) + var bufferedBody = BufferedBody; + if ((uint)columnOffset > (uint)bufferedBody.Length) { field = default; return false; } - var remainingMemory = _bufferedBody.Slice(columnOffset); + var remainingMemory = bufferedBody.Slice(columnOffset); var remaining = remainingMemory.Span; while (columnIndex++ < ordinal) @@ -497,6 +532,8 @@ internal ValueTask BufferAllAsync(CancellationToken cancellationToken = default) return default; return Core(cancellationToken); + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] async ValueTask Core(CancellationToken token) { await _bodyReader.BufferAllAsync(token).ConfigureAwait(false); @@ -552,6 +589,16 @@ public T Read() _remaining = _remaining.Slice(sizeof(int) + length); return BootstrapFieldDecoder.Read(field); } + + /// + /// Borrows the next field's raw PostgreSQL representation as contiguous memory. + /// + /// + /// The returned memory has the same unenforced lifetime as + /// . + /// + public ReadOnlyMemory ReadBorrowedMemory() + => _row.BorrowFieldMemory(_ordinal++); } internal void Initialize(RowDescription rowDescription) @@ -560,15 +607,21 @@ internal void Initialize(RowDescription rowDescription) _rowDescription = rowDescription; } - internal void InitializeRow(in BackendMessage row) + internal void InitializeRow(in BackendMessage.Accessor row) { if (_columnLease is not null) throw new InvalidOperationException("The previous column lease must be revoked before advancing the row."); - _bodyReader = row.Buffered ? null : row.OpenBodyReader(); + if (row.Buffered) + { + if (_bodyReader is not null) + _bodyReader = null; + } + else + _bodyReader = row.OpenBodyReader(); _column = 0; _columnOffset = sizeof(short); _lastBufferedOrdinal = -1; - BackendMessage.Accessor.WriteGranularly(ref _messageAccessor, row.GetAccessor()); + BackendMessage.Accessor.Assign(ref _messageAccessor, row); CaptureBufferedBody(row); } @@ -582,17 +635,47 @@ void EnsureBuffered() } void CaptureBufferedBody() - { - var message = Message; - CaptureBufferedBody(message); - } + => CaptureBufferedBody(_messageAccessor); - void CaptureBufferedBody(in BackendMessage message) + void CaptureBufferedBody(in BackendMessage.Accessor message) { - if (_bodyReader is null && message.TryGetBufferedFirstMemory(0, out var body)) - _bufferedBody = body; + if (_bodyReader is null + && message.TryGetBufferedArray(0, out var array, out var offset, out var length)) + { + if (!ReferenceEquals(_bufferedArray, array)) + _bufferedArray = array; + _bufferedOffset = offset; + _bufferedLength = length; + if (!_bufferedBody.IsEmpty) + _bufferedBody = default; + } + else if (_bodyReader is null && message.TryGetBufferedFirstMemory(0, out var body)) + { + if (MemoryMarshal.TryGetArray(body, out var segment)) + { + if (!ReferenceEquals(_bufferedArray, segment.Array)) + _bufferedArray = segment.Array; + _bufferedOffset = segment.Offset; + _bufferedLength = segment.Count; + if (!_bufferedBody.IsEmpty) + _bufferedBody = default; + } + else + { + if (_bufferedArray is not null) + _bufferedArray = null; + _bufferedBody = body; + } + } else - _bufferedBody = default; + { + if (_bufferedArray is not null) + _bufferedArray = null; + if (!_bufferedBody.IsEmpty) + _bufferedBody = default; + _bufferedOffset = 0; + _bufferedLength = 0; + } } // Returns false when the seek was exhausted, true if positioned correctly, and throws if the seek is invalid. diff --git a/Slon/Pg/Serialization/PgStreamingConverter.cs b/Slon/Pg/Serialization/PgStreamingConverter.cs index f6c4f3b..8272c75 100644 --- a/Slon/Pg/Serialization/PgStreamingConverter.cs +++ b/Slon/Pg/Serialization/PgStreamingConverter.cs @@ -67,6 +67,7 @@ static class PgStreamingConverterHelpers { // Split out from the generic class to amortize the huge size penalty per async state machine, which would otherwise be per // instantiation. + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] public static async ValueTask AwaitTask(Task task, Continuation continuation) { diff --git a/Slon/Pipelines/DelegatedPipelineScheduler.cs b/Slon/Pipelines/DelegatedPipelineScheduler.cs new file mode 100644 index 0000000..9642d0f --- /dev/null +++ b/Slon/Pipelines/DelegatedPipelineScheduler.cs @@ -0,0 +1,10 @@ +using Draghi.Pipelining; +using Slon.Threading; + +namespace Slon.Pipelines; + +sealed class DelegatedPipelineScheduler(Scheduler scheduler) : PipelineScheduler +{ + public override void SubmitDetached(Action action, object? state, bool preferLocal = true) + => scheduler.SubmitDetached(action, state, preferLocal); +} diff --git a/Slon/Pipelines/PipeOutputWriter.cs b/Slon/Pipelines/PipeOutputWriter.cs index 95e1b2b..2820065 100644 --- a/Slon/Pipelines/PipeOutputWriter.cs +++ b/Slon/Pipelines/PipeOutputWriter.cs @@ -42,6 +42,7 @@ ValueTask IOutputWriter.FlushAsync(CancellationToken cancellationToken) EnsureFlushed(flushTask.Result); return new(); + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] static async ValueTask Core(ValueTask flushTask) => EnsureFlushed(await flushTask.ConfigureAwait(false)); diff --git a/Slon/Pipelines/PipeSegmentEnumerator.cs b/Slon/Pipelines/PipeSegmentEnumerator.cs deleted file mode 100644 index 9f4ffbd..0000000 --- a/Slon/Pipelines/PipeSegmentEnumerator.cs +++ /dev/null @@ -1,592 +0,0 @@ -using System.Buffers; -using System.Collections; -using System.IO.Pipelines; -using System.Runtime.CompilerServices; - -namespace Slon.Pipelines; - -interface IPipeSegmenter -{ - /// - /// MinimumSize guarantees CreateSegment won't be called unless there is enough data to examine. - /// - int MinimumSize { get; } - - /// - /// Create the next segment from the given buffer, returning segment information on OperationStatus.Done. - /// - /// The buffer to try to read the next segment from. - /// The length of the segment, this may be larger than the amount buffered at the time of the call. - /// Segment to return to the caller. - /// Whether the call was successful, requires more data, or invalid data was found, DestinationTooSmall is not supported. - OperationStatus CreateSegment(in ReadOnlySequence buffer, out long segmentLength, out TSegment segment); -} - -readonly struct CurrentSegmentBuffer(ReadOnlySequence buffer, bool isComplete) -{ - public ReadOnlySequence Buffer { get; } = buffer; - public bool IsComplete { get; } = isComplete; -} - -sealed class PipeSegmentEnumerator(PipeReader reader, TSegmenter segmenter, bool ownsReader = false) - : IEnumerator, IAsyncEnumerator - where TSegmenter: IPipeSegmenter -{ - readonly StreamPipeReader? _directReader = reader as StreamPipeReader; - TSegmenter _segmenter = segmenter; - TSegment _current = default!; - - SequencePosition _examinedPosition; - SequencePosition _currentSegmentStart; - SequencePosition? _consumePosition; - long _currentLength = -1; - byte _currentSegmentReadPending; - - public PipeReader PipeReader => reader; - - public ValueTask ReadAsync(CancellationToken cancellationToken) - => reader.ReadAsync(cancellationToken); - - public bool TryBeginDirectRead(CancellationToken cancellationToken, out ValueTask task) - { - if (_directReader is { SupportsDirectRead: true } directReader) - { - task = directReader.BeginDirectRead(cancellationToken); - return true; - } - task = default; - return false; - } - - public bool CompleteDirectRead(int length, CancellationToken cancellationToken, out ValueTask next, out bool readFinished, out bool completed) - { - if (!_directReader!.CompleteDirectRead(length, cancellationToken, out next, out var result)) - { - readFinished = false; - completed = false; - return false; - } - readFinished = true; - return TryMoveNext(result, cancellationToken, out completed); - } - - public void AbortDirectRead() => _directReader!.AbortDirectRead(); - - public bool TryMoveNext(ReadResult result, CancellationToken cancellationToken, out bool completed) - { - if (result.IsCanceled) - ThrowHelper.ThrowOperationCanceled(cancellationToken); - return TryMoveNext(result, hasRead: true, out completed); - } - - // The underlying reader reported completion, so the wire is at EOF. Disarm the deferred advance by - // clearing the pending-segment sentinel: a re-drive past completion (recovery drain, or any caller - // that keeps pulling after false) must not re-apply a stale consume position, whose segment and - // backing array have since been consumed and pool-recycled, driving the buffer accounting negative. - // Returns false so completion sites read as return EndOfData(). - bool EndOfData() - { - _consumePosition = null; - _examinedPosition = default; - _currentSegmentStart = default; - _currentLength = -1; - _currentSegmentReadPending = 0; - return false; - } - - static bool IsEmptyCompletion(in ReadResult result) - => result.IsCompleted && result.Buffer.IsEmpty; - - static void ThrowIfTruncatedCompletion(in ReadResult result) - { - if (result.IsCompleted) - throw new EndOfStreamException("The pipe completed within a framed segment."); - } - - ValueTask IAsyncEnumerator.MoveNextAsync() => MoveNextAsync(CancellationToken.None); - - // Nonblocking counterpart to MoveNext/MoveNextAsync. It consumes every byte already available - // from the reader and returns false only when another physical read is required. completed - // distinguishes that would-block from EOF. This is the poll primitive used by read-wake drivers: - // after one leaf wake they re-enter here and synchronously descend framing again. - public bool TryMoveNext(out bool completed) - => TryMoveNext(default, hasRead: false, out completed); - - // Releases a consumed prefix of a partially buffered segment and polls for its next bytes. The - // returned buffer never crosses the segment boundary; normal MoveNext resumes at that boundary. - public bool TryContinueCurrentSegment(SequencePosition consumed, long consumedLength, out CurrentSegmentBuffer result) - { - PrepareCurrentSegmentRead(consumed, consumedLength, mode: 1); - if (!reader.TryRead(out var readResult)) - { - result = default; - return false; - } - - result = CompleteCurrentSegmentRead(readResult); - return true; - } - - public bool TryExtendCurrentSegment(out CurrentSegmentBuffer result) - { - PrepareCurrentSegmentRead(_currentSegmentStart, consumedLength: 0, mode: 2); - if (!reader.TryRead(out var readResult)) - { - result = default; - return false; - } - - result = CompleteCurrentSegmentRead(readResult); - return true; - } - - public ValueTask ContinueCurrentSegmentAsync( - SequencePosition consumed, long consumedLength, CancellationToken cancellationToken = default) - { - PrepareCurrentSegmentRead(consumed, consumedLength, mode: 1); - var task = reader.ReadAsync(cancellationToken); - return task.IsCompletedSuccessfully - ? new(CompleteCurrentSegmentRead(task.Result, cancellationToken)) - : Core(task, cancellationToken); - - [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] - async ValueTask Core(ValueTask task, CancellationToken cancellationToken) - => CompleteCurrentSegmentRead(await task.ConfigureAwait(false), cancellationToken); - } - - public ValueTask ExtendCurrentSegmentAsync(CancellationToken cancellationToken = default) - { - PrepareCurrentSegmentRead(_currentSegmentStart, consumedLength: 0, mode: 2); - var task = reader.ReadAtLeastAsync((int)Math.Min(_currentLength, int.MaxValue), cancellationToken); - return task.IsCompletedSuccessfully - ? new(CompleteCurrentSegmentRead(task.Result, cancellationToken)) - : Core(task, cancellationToken); - - [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] - async ValueTask Core(ValueTask task, CancellationToken cancellationToken) - => CompleteCurrentSegmentRead(await task.ConfigureAwait(false), cancellationToken); - } - - public CurrentSegmentBuffer ContinueCurrentSegment( - SequencePosition consumed, long consumedLength, TimeSpan timeout = default) - { - if (reader is not StreamPipeReader syncReader) - throw new NotSupportedException("Underlying pipe reader does not support synchronous reads."); - - PrepareCurrentSegmentRead(consumed, consumedLength, mode: 1); - return CompleteCurrentSegmentRead(syncReader.Read(timeout)); - } - - public CurrentSegmentBuffer ExtendCurrentSegment(TimeSpan timeout = default) - { - if (reader is not StreamPipeReader syncReader) - throw new NotSupportedException("Underlying pipe reader does not support synchronous reads."); - - PrepareCurrentSegmentRead(_currentSegmentStart, consumedLength: 0, mode: 2); - return CompleteCurrentSegmentRead(syncReader.ReadAtLeast((int)Math.Min(_currentLength, int.MaxValue), timeout)); - } - - void PrepareCurrentSegmentRead(SequencePosition consumed, long consumedLength, byte mode) - { - if (_currentSegmentReadPending != 0) - { - if (_currentSegmentReadPending != mode) - ThrowHelper.ThrowInvalidOperation("The pending segment read uses a different continuation mode."); - return; - } - if (_currentLength <= 0 || _consumePosition is not null) - ThrowHelper.ThrowInvalidOperation("The current segment is not awaiting more data."); - if (mode == 1 && (consumedLength <= 0 || consumedLength >= _currentLength)) - throw new ArgumentOutOfRangeException(nameof(consumedLength)); - - reader.AdvanceTo(consumed, _examinedPosition); - _currentLength -= consumedLength; - _currentSegmentReadPending = mode; - } - - CurrentSegmentBuffer CompleteCurrentSegmentRead(ReadResult result, - CancellationToken cancellationToken = default) - { - _currentSegmentReadPending = 0; - if (result.IsCanceled) - ThrowHelper.ThrowOperationCanceled(cancellationToken); - if (result.IsCompleted && result.Buffer.Length < _currentLength) - throw new EndOfStreamException("The pipe completed within a framed segment."); - - var isComplete = result.Buffer.Length >= _currentLength; - var buffer = isComplete ? result.Buffer.Slice(0, _currentLength) : result.Buffer; - _currentSegmentStart = result.Buffer.Start; - _consumePosition = isComplete ? buffer.End : null; - _examinedPosition = buffer.End; - return new(buffer, isComplete); - } - - bool TryMoveNext(ReadResult suppliedRead, bool hasRead, out bool completed) - { - completed = false; - - if (_currentLength is not -1) - { - var segmentReadPending = _currentSegmentReadPending != 0; - _currentSegmentReadPending = 0; - if (_consumePosition is null) - { - // A supplied read is already the result of the advance performed by the poll which - // returned false. Advancing again would retire that result before we inspect it. - if (!segmentReadPending && !hasRead) - reader.AdvanceTo(_currentSegmentStart, _examinedPosition); - if (!TryTakeRead(out var consumeResult)) - return false; - if (IsEmptyCompletion(consumeResult)) - ThrowIfTruncatedCompletion(consumeResult); - if (consumeResult.IsCanceled) - ThrowHelper.ThrowOperationCanceled(CancellationToken.None); - if (consumeResult.Buffer.Length < _currentLength) - { - ThrowIfTruncatedCompletion(consumeResult); - reader.AdvanceTo(consumeResult.Buffer.Start, consumeResult.Buffer.End); - return false; - } - reader.AdvanceTo(consumeResult.Buffer.GetPosition(_currentLength)); - _consumePosition = null; - _currentLength = -1; - } - else - { - reader.AdvanceTo(_consumePosition.GetValueOrDefault(), _examinedPosition); - _consumePosition = null; - _currentLength = -1; - } - } - - if (!TryTakeRead(out var result)) - return false; - if (IsEmptyCompletion(result)) - { - completed = true; - return EndOfData(); - } - if (result.IsCanceled) - ThrowHelper.ThrowOperationCanceled(CancellationToken.None); - - var status = _segmenter.CreateSegment(result.Buffer, out _currentLength, out _current); - switch (status) - { - case OperationStatus.NeedMoreData when _currentLength > 0: - case OperationStatus.Done: - ArgumentOutOfRangeException.ThrowIfNegativeOrZero(_currentLength, "segmentLength"); - _consumePosition = _currentLength <= result.Buffer.Length ? result.Buffer.GetPosition(_currentLength) : null; - _currentSegmentStart = result.Buffer.Start; - _examinedPosition = _consumePosition ?? result.Buffer.End; - return true; - case OperationStatus.NeedMoreData: - ThrowIfTruncatedCompletion(result); - reader.AdvanceTo(result.Buffer.Start, result.Buffer.End); - _currentLength = -1; - return false; - case OperationStatus.InvalidData: - reader.Complete(new Exception("Segmenter encountered invalid data.")); - completed = true; - return false; - case OperationStatus.DestinationTooSmall: - ThrowHelper.ThrowInvalidOperation(); - return default; - case var value: - ThrowHelper.ThrowUnhandledCase(value); - return default; - } - - bool TryTakeRead(out ReadResult result) - { - if (hasRead) - { - result = suppliedRead; - suppliedRead = default; - hasRead = false; - return true; - } - return reader.TryRead(out result); - } - } - - public ValueTask MoveNextAsync(CancellationToken cancellationToken = default) - { - ValueTask task; - ReadResult result; - - // Advance past current segment. - if (_currentLength is not -1) - { - var segmentReadPending = _currentSegmentReadPending != 0; - _currentSegmentReadPending = 0; - // Not everything was buffered when the segment was returned (e.g. with length prefixed segments). - if (_consumePosition is null) - { - if (!segmentReadPending) - reader.AdvanceTo(_currentSegmentStart, _examinedPosition); - task = reader.ReadAtLeastAsync((int)long.Min(_currentLength, int.MaxValue), cancellationToken); - if (!task.IsCompletedSuccessfully) - return Core(task, cancellationToken, consume: true); - result = task.Result; - if (IsEmptyCompletion(result)) - ThrowIfTruncatedCompletion(result); - if (result.IsCanceled) - return new(Task.FromException(new OperationCanceledException(cancellationToken))); - - if (result.Buffer.Length < _currentLength) - { - ThrowIfTruncatedCompletion(result); - return Core(new(result), cancellationToken, consume: true); - } - if (result.Buffer.Length > _currentLength) - return Core(new(result), cancellationToken, consume: true); - reader.AdvanceTo(result.Buffer.GetPosition(_currentLength)); - _consumePosition = null; - } - else - { - reader.AdvanceTo(_consumePosition.GetValueOrDefault(), _examinedPosition); - _consumePosition = null; - } - } - - task = reader.ReadAtLeastAsync(_segmenter.MinimumSize, cancellationToken); - if (!task.IsCompletedSuccessfully) - return Core(task, cancellationToken); - - result = task.Result; - if (IsEmptyCompletion(result)) - return new(EndOfData()); - if (result.IsCanceled) - return new(Task.FromException(new OperationCanceledException(cancellationToken))); - - var status = _segmenter.CreateSegment(result.Buffer, out _currentLength, out _current); - switch (status) - { - case OperationStatus.NeedMoreData when _currentLength > 0: - case OperationStatus.Done: - ArgumentOutOfRangeException.ThrowIfNegativeOrZero(_currentLength, "segmentLength"); - _consumePosition = _currentLength <= result.Buffer.Length ? result.Buffer.GetPosition(_currentLength) : null; - _currentSegmentStart = result.Buffer.Start; - // Stop examined at the segment boundary so trailing buffered bytes (next segment's data) stay visible to the next ReadAsync. - _examinedPosition = _consumePosition ?? result.Buffer.End; - return new(true); - case OperationStatus.DestinationTooSmall: - ThrowHelper.ThrowInvalidOperation(); - return default; - case OperationStatus.NeedMoreData: - ThrowIfTruncatedCompletion(result); - return Core(new(result), cancellationToken, needMoreData: true); - case OperationStatus.InvalidData: - return InvalidData(); - case var value: - ThrowHelper.ThrowUnhandledCase(value); - return default; - } - - - [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] - async ValueTask Core(ValueTask task, CancellationToken cancellationToken, bool consume = false, bool needMoreData = false) - { - while (true) - { - var result = await task.ConfigureAwait(false); - if (IsEmptyCompletion(result)) - { - if (consume || needMoreData) - ThrowIfTruncatedCompletion(result); - return EndOfData(); - } - if (result.IsCanceled) - ThrowHelper.ThrowOperationCanceled(cancellationToken); - - if (consume) - { - if (result.Buffer.Length < _currentLength) - { - ThrowIfTruncatedCompletion(result); - reader.AdvanceTo(result.Buffer.Start, result.Buffer.End); - task = reader.ReadAsync(cancellationToken); - continue; - } - reader.AdvanceTo(result.Buffer.GetPosition(_currentLength)); - task = reader.ReadAtLeastAsync(_segmenter.MinimumSize, cancellationToken); - consume = false; - continue; - } - if (needMoreData) - { - reader.AdvanceTo(result.Buffer.Start, result.Buffer.End); - _currentLength = -1; - task = reader.ReadAtLeastAsync(_segmenter.MinimumSize, cancellationToken); - needMoreData = false; - continue; - } - - var status = _segmenter.CreateSegment(result.Buffer, out _currentLength, out _current); - switch (status) - { - case OperationStatus.NeedMoreData when _currentLength > 0: - case OperationStatus.Done: - ArgumentOutOfRangeException.ThrowIfNegativeOrZero(_currentLength, "segmentLength"); - _consumePosition = _currentLength <= result.Buffer.Length ? result.Buffer.GetPosition(_currentLength) : null; - _currentSegmentStart = result.Buffer.Start; - // Stop examined at the segment boundary so trailing buffered bytes stay visible to the next ReadAsync. - _examinedPosition = _consumePosition ?? result.Buffer.End; - return true; - case OperationStatus.DestinationTooSmall: - ThrowHelper.ThrowInvalidOperation(); - return default; - case OperationStatus.NeedMoreData: - ThrowIfTruncatedCompletion(result); - needMoreData = true; - break; - case OperationStatus.InvalidData: - return await InvalidData().ConfigureAwait(false); - case var value: - ThrowHelper.ThrowUnhandledCase(value); - return default; - } - } - } - - async ValueTask InvalidData() - { - await reader.CompleteAsync(new Exception("Segmenter encountered invalid data.")).ConfigureAwait(false); - return false; - } - } - - bool IEnumerator.MoveNext() => MoveNext(default(TimeSpan)); - public bool MoveNext(TimeSpan timeout = default) - { - if (reader is not StreamPipeReader syncReader) - throw new NotSupportedException("Underlying pipe reader does not support synchronous reads."); - - ReadResult result; - var consume = false; - var needMoreData = false; - - // Advance past current segment. - if (_currentLength is not -1) - { - var segmentReadPending = _currentSegmentReadPending != 0; - _currentSegmentReadPending = 0; - if (_consumePosition is null) - { - if (!segmentReadPending) - reader.AdvanceTo(_currentSegmentStart, _examinedPosition); - result = syncReader.ReadAtLeast((int)long.Min(_currentLength, int.MaxValue), timeout); - if (IsEmptyCompletion(result)) - ThrowIfTruncatedCompletion(result); - if (result.IsCanceled) - ThrowHelper.ThrowOperationCanceled(CancellationToken.None); - - if (result.Buffer.Length < _currentLength) - { - ThrowIfTruncatedCompletion(result); - consume = true; - goto loop; - } - if (result.Buffer.Length > _currentLength) - { - consume = true; - goto loop; - } - reader.AdvanceTo(result.Buffer.GetPosition(_currentLength)); - _consumePosition = null; - } - else - { - reader.AdvanceTo(_consumePosition.GetValueOrDefault(), _examinedPosition); - _consumePosition = null; - } - } - - result = syncReader.ReadAtLeast(_segmenter.MinimumSize, timeout); - - loop: - while (true) - { - if (IsEmptyCompletion(result)) - { - if (consume || needMoreData) - ThrowIfTruncatedCompletion(result); - return EndOfData(); - } - if (result.IsCanceled) - ThrowHelper.ThrowOperationCanceled(CancellationToken.None); - - if (consume) - { - if (result.Buffer.Length < _currentLength) - { - ThrowIfTruncatedCompletion(result); - reader.AdvanceTo(result.Buffer.Start, result.Buffer.End); - result = syncReader.Read(timeout); - continue; - } - reader.AdvanceTo(result.Buffer.GetPosition(_currentLength)); - result = syncReader.ReadAtLeast(_segmenter.MinimumSize, timeout); - consume = false; - continue; - } - if (needMoreData) - { - reader.AdvanceTo(result.Buffer.Start, result.Buffer.End); - _currentLength = -1; - result = syncReader.ReadAtLeast(_segmenter.MinimumSize, timeout); - needMoreData = false; - continue; - } - - var status = _segmenter.CreateSegment(result.Buffer, out _currentLength, out _current); - switch (status) - { - case OperationStatus.NeedMoreData when _currentLength > 0: - case OperationStatus.Done: - ArgumentOutOfRangeException.ThrowIfNegativeOrZero(_currentLength, "segmentLength"); - _consumePosition = _currentLength <= result.Buffer.Length ? result.Buffer.GetPosition(_currentLength) : null; - _currentSegmentStart = result.Buffer.Start; - // Stop examined at the segment boundary so trailing buffered bytes stay visible to the next ReadAsync. - _examinedPosition = _consumePosition ?? result.Buffer.End; - return true; - case OperationStatus.DestinationTooSmall: - ThrowHelper.ThrowInvalidOperation(); - return default; - case OperationStatus.NeedMoreData: - ThrowIfTruncatedCompletion(result); - needMoreData = true; - break; - case OperationStatus.InvalidData: - reader.Complete(new Exception("Segmenter encountered invalid data.")); - return false; - case var value: - ThrowHelper.ThrowUnhandledCase(value); - return default; - } - } - } - - public TSegment Current - { - [MethodImpl(MethodImplOptions.AggressiveInlining)] - get => _current; - } - - public void Dispose() - { - if (ownsReader) - reader.Complete(); - } - - public ValueTask DisposeAsync() - { - if (ownsReader) - return reader.CompleteAsync(); - return new(); - } - - object? IEnumerator.Current => Current; - void IEnumerator.Reset() => throw new NotSupportedException(); -} diff --git a/Slon/Pipelines/StreamPipeReader.cs b/Slon/Pipelines/StreamPipeReader.cs index 2fcc725..59923da 100644 --- a/Slon/Pipelines/StreamPipeReader.cs +++ b/Slon/Pipelines/StreamPipeReader.cs @@ -12,12 +12,14 @@ interface IStreamOwner : IDisposable, IAsyncDisposable { } abstract class StreamPipeReader : PipeReader { +#if !NET11_0_OR_GREATER readonly ValueTaskSourcePromise _readAsyncCorePromise = new(); +#endif + bool _directReadAwaitingData; readonly IStreamOwner? _streamOwner; int _isReadActive; int _readerCompleted; int _readerDisposed; - bool _directReadAwaitingData; // Null in direct-read mode (CancelPendingRead unsupported): the caller's token threads straight to // the underlying stream read, so neither this source nor a per-read registration is allocated. @@ -73,8 +75,7 @@ public override void AdvanceTo(SequencePosition consumed, SequencePosition exami ThrowIfCompleted(); var bufferedBytes = Segments.BufferedBytes; var examinedBytes = Segments.AdvanceTo(consumed, examined); - if (examinedBytes == bufferedBytes) - ExaminedEverything = true; + ExaminedEverything = examinedBytes == bufferedBytes; } /// @@ -223,17 +224,18 @@ public override bool TryRead(out ReadResult result) return TryReadCore(out result); } - // Direct reads retain reader tenure until their continuation returns through CompleteDirectRead - // or AbortDirectRead. The protocol must interrupt and join that tenure before completing the - // reader and returning its destination buffer. - internal bool SupportsDirectRead => PendingReadTokenSource is null; - internal void EnsureCanUpgradeStream() { if (IsReaderCompleted || Volatile.Read(ref _isReadActive) is not 0 || Segments.BufferedBytes is not 0) throw new InvalidOperationException("The reader must be open, idle, and empty before its stream can be upgraded."); } + // Direct reads retain reader tenure until their continuation returns through CompleteDirectRead + // or AbortDirectRead. The protocol must interrupt and join that tenure before completing the + // reader and returning its destination buffer. + internal bool SupportsDirectRead => PendingReadTokenSource is null; + const int BufferedDirectRead = -1; + // Direct leaf handoff. The caller awaits the stream's ValueTask directly, then returns // the byte count through CompleteDirectRead. This keeps buffer ownership and PipeReader read tenure // here while removing the intermediate ReadAsyncCore completion frame. @@ -247,7 +249,9 @@ internal ValueTask BeginDirectRead(CancellationToken cancellationToken) try { - if (Segments.BufferedBytes is 0 && UseZeroByteReads) + if (Segments.BufferedBytes is not 0 && !ExaminedEverything) + return new(BufferedDirectRead); + if (UseZeroByteReads) { _directReadAwaitingData = true; return Stream.ReadAsync(Memory.Empty, cancellationToken); @@ -276,12 +280,13 @@ internal bool CompleteDirectRead(int length, CancellationToken cancellationToken length = next.Result; } - if (length is not 0) + if (length > 0) { ExaminedEverything = false; Segments.Grow(length); } - result = new ReadResult(Segments.GetReadOnlySequence(), isCanceled: false, isCompleted: length is 0); + result = new ReadResult(Segments.GetReadOnlySequence(), isCanceled: false, + isCompleted: length is 0); next = default; EndStartedRead(); return true; @@ -296,7 +301,6 @@ internal void AbortDirectRead() if (Volatile.Read(ref _isReadActive) is not 0) EndStartedRead(); } - ValueTask StartDataRead(CancellationToken cancellationToken) { var buffer = Segments.Reserve(0, enforceHint: false); @@ -332,13 +336,7 @@ protected ReadResult ReadCore(int minimumSize, TimeSpan timeout) do { // We know minimumSize must be null or larger than what we have buffered to get here. - var segmentSize = minimumSize; - if (minimumSize is not 0) - { - // We must request a segment that is minimumSize - BufferedBytes. - Debug.Assert(Segments.BufferedBytes <= minimumSize); - segmentSize -= (int)Segments.BufferedBytes; - } + var segmentSize = GetReadSizeHint(minimumSize); // We don't mind if we get smaller segments, we just want to make progress towards minimumSize. var buffer = Segments.Reserve(segmentSize, enforceHint: false); @@ -382,35 +380,66 @@ protected ReadResult ReadCore(int minimumSize, TimeSpan timeout) protected ValueTask ReadAsyncCore(int minimumSize, CancellationToken cancellationToken) { +#if !NET11_0_OR_GREATER PromiseAsyncValueTaskMethodBuilder.Promise = _readAsyncCorePromise; try { - return ReadAsyncCore(minimumSize, PendingReadTokenSource, cancellationToken); +#endif + return StartReadAsync(minimumSize, PendingReadTokenSource, cancellationToken); +#if !NET11_0_OR_GREATER } finally { PromiseAsyncValueTaskMethodBuilder.Promise = null; } +#endif - [RuntimeAsyncMethodGeneration(false)] - [AsyncMethodBuilder(typeof(PromiseAsyncValueTaskMethodBuilder<>))] - async ValueTask ReadAsyncCore(int minimumSize, + ValueTask StartReadAsync(int minimumSize, AutoResetCancellationTokenSource? tokenSource, CancellationToken cancellationToken) { - // Cancellation token was already checked before getting here. if (!TryStartRead()) ThrowAlreadyReading(); CancellationTokenRegistration reg = default; - CancellationToken token; - if (tokenSource is { } src) + var registered = false; + try { - if (cancellationToken.CanBeCanceled) - reg = src.UnsafeRegister(cancellationToken); - token = src.Token; + CancellationToken token; + if (tokenSource is { } src) + { + if (cancellationToken.CanBeCanceled) + { + registered = true; + reg = src.UnsafeRegister(cancellationToken); + } + token = src.Token; + } + else + { + token = cancellationToken; + } + + var read = ReadLoopAsync( + minimumSize, token, cancellationToken, endStartedRead: !registered); + return registered ? CompleteRegisteredReadAsync(read, reg) : read; } - else - token = cancellationToken; + catch + { + if (registered) + reg.Dispose(); + EndStartedRead(); + throw; + } + } + +#if !NET11_0_OR_GREATER + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PromiseAsyncValueTaskMethodBuilder<>))] +#endif + async ValueTask ReadLoopAsync( + int minimumSize, CancellationToken token, CancellationToken cancellationToken, + bool endStartedRead) + { try { if (Segments.BufferedBytes is 0 && UseZeroByteReads) @@ -422,12 +451,7 @@ async ValueTask ReadAsyncCore(int minimumSize, int length; do { - var segmentSize = minimumSize; - if (minimumSize is not 0) - { - Debug.Assert(Segments.BufferedBytes <= minimumSize); - segmentSize -= (int)Segments.BufferedBytes; - } + var segmentSize = GetReadSizeHint(minimumSize); var buffer = Segments.Reserve(segmentSize, enforceHint: false); length = await Stream.ReadAsync(buffer, token).ConfigureAwait(false); @@ -449,10 +473,28 @@ async ValueTask ReadAsyncCore(int minimumSize, throw; } finally + { + if (endStartedRead) + EndStartedRead(); + } + } + +#if !NET11_0_OR_GREATER + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] +#endif + async ValueTask CompleteRegisteredReadAsync( + ValueTask read, CancellationTokenRegistration registration) + { + try + { + return await read.ConfigureAwait(false); + } + finally { try { - await reg.DisposeAsync().ConfigureAwait(false); + await registration.DisposeAsync().ConfigureAwait(false); } finally { @@ -462,6 +504,17 @@ async ValueTask ReadAsyncCore(int minimumSize, } } + [MethodImpl(MethodImplOptions.AggressiveInlining)] + int GetReadSizeHint(int minimumSize) + { + if (minimumSize is 0) + return 0; + var bufferedBytes = Segments.BufferedBytes; + return bufferedBytes < minimumSize + ? minimumSize - (int)bufferedBytes + : 0; + } + protected async Task CopyToAsyncCore(PipeWriter destination, CancellationToken cancellationToken = default) { if (!TryStartRead()) diff --git a/Slon/Pipelines/StreamPipeWriter.cs b/Slon/Pipelines/StreamPipeWriter.cs index fc42a3e..4bb0a87 100644 --- a/Slon/Pipelines/StreamPipeWriter.cs +++ b/Slon/Pipelines/StreamPipeWriter.cs @@ -9,7 +9,9 @@ namespace Slon.Pipelines; abstract class StreamPipeWriter : PipeWriter, IOutputWriter { +#if !NET11_0_OR_GREATER readonly ValueTaskSourcePromise _flushAsyncCorePromise = new(); +#endif readonly IStreamOwner? _streamOwner; bool _isFlushActive; @@ -175,6 +177,7 @@ ValueTask IOutputWriter.FlushAsync(CancellationToken cancellationToken) EnsureFlushed(flushTask.Result); return default; + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] static async ValueTask Core(ValueTask flushTask) => EnsureFlushed(await flushTask.ConfigureAwait(false)); @@ -307,18 +310,24 @@ protected virtual FlushResult FlushCore(bool writeToStream, ReadOnlySpan d protected virtual ValueTask FlushAsyncCore(bool writeToStream, ReadOnlyMemory data, CancellationToken cancellationToken) { +#if !NET11_0_OR_GREATER PromiseAsyncValueTaskMethodBuilder.Promise = _flushAsyncCorePromise; try { +#endif return FlushAsyncCore(PendingFlushTokenSource, writeToStream, data, cancellationToken); +#if !NET11_0_OR_GREATER } finally { PromiseAsyncValueTaskMethodBuilder.Promise = null; } +#endif +#if !NET11_0_OR_GREATER [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PromiseAsyncValueTaskMethodBuilder<>))] +#endif async ValueTask FlushAsyncCore(AutoResetCancellationTokenSource? tokenSource, bool writeToStream, ReadOnlyMemory data, CancellationToken cancellationToken) { // Cancellation token was already checked before getting here. @@ -396,7 +405,11 @@ async ValueTask FlushAsyncCore(AutoResetCancellationTokenSource? to if (didWrite) { - await Stream.FlushAsync(token).ConfigureAwait(false); + var flush = Stream.FlushAsync(token); + if (!flush.IsCompletedSuccessfully) + await AwaitFlush(flush).ConfigureAwait(false); + else + flush.GetAwaiter().GetResult(); } } @@ -428,6 +441,11 @@ async ValueTask FlushAsyncCore(AutoResetCancellationTokenSource? to EndStartedFlush(); } } + + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] + static async ValueTask AwaitFlush(Task flush) + => await flush.ConfigureAwait(false); } } diff --git a/Slon/PublicAPI.Unshipped.txt b/Slon/PublicAPI.Unshipped.txt index 8622b53..d292969 100644 --- a/Slon/PublicAPI.Unshipped.txt +++ b/Slon/PublicAPI.Unshipped.txt @@ -366,7 +366,6 @@ Slon.SlonConnectionInitializerContext.Connection.get -> Slon.SlonConnection! Slon.SlonDataReader Slon.SlonDataReader.GetBytes(int ordinal) -> byte[]! Slon.SlonDataReader.GetColumnSchema() -> System.Collections.ObjectModel.ReadOnlyCollection! -Slon.SlonDataReader.GetData(int ordinal) -> Slon.SlonDataReader! Slon.SlonDataReader.LongRecordsAffected.get -> long Slon.SlonDataSource Slon.SlonDataSource.CreateBatch() -> Slon.SlonBatch! diff --git a/Slon/Runtime/CompilerServices/FieldRef.cs b/Slon/Runtime/CompilerServices/FieldRef.cs deleted file mode 100644 index 6e72c79..0000000 --- a/Slon/Runtime/CompilerServices/FieldRef.cs +++ /dev/null @@ -1,31 +0,0 @@ -using System.Runtime.CompilerServices; - -namespace Slon.Runtime.CompilerServices; - -// This is safe against unloading as long as instance is the same type as the getter function is defined on. -// There is unfortunately no static type safety to guarantee a user creates one correctly. -readonly struct FieldRef -{ - readonly unsafe delegate* _getter; - readonly object _instance; - - unsafe FieldRef(delegate* getter, object instance) - { - _getter = getter; - _instance = instance; - } - - public object Instance => _instance; - - [MethodImpl(MethodImplOptions.AggressiveInlining)] - public ref T Invoke() - { - unsafe - { - return ref _getter(_instance); - } - } - - public static unsafe FieldRef Create(delegate* getter, TInstance instance) where TInstance : class - => new((delegate*)getter, instance); -} diff --git a/Slon/Runtime/CompilerServices/IFieldRef.cs b/Slon/Runtime/CompilerServices/IFieldRef.cs new file mode 100644 index 0000000..9ef1423 --- /dev/null +++ b/Slon/Runtime/CompilerServices/IFieldRef.cs @@ -0,0 +1,9 @@ +namespace Slon.Runtime.CompilerServices; + +// Provides reusable, allocation-free access to a field retained by another object or struct. +// Generic consumers preserve the concrete implementation type so the ref-returning call can inline. +interface IFieldRef + where TOwner : struct, IFieldRef +{ + ref T GetField(); +} diff --git a/Slon/Runtime/CompilerServices/PromiseAsyncValueTaskMethodBuilder.cs b/Slon/Runtime/CompilerServices/PromiseAsyncValueTaskMethodBuilder.cs index c347898..edab8f0 100644 --- a/Slon/Runtime/CompilerServices/PromiseAsyncValueTaskMethodBuilder.cs +++ b/Slon/Runtime/CompilerServices/PromiseAsyncValueTaskMethodBuilder.cs @@ -252,7 +252,10 @@ public void AwaitUnsafeOnCompleted( } catch (Exception ex) { - ThreadPool.QueueUserWorkItem(state => ((ExceptionDispatchInfo)state!).Throw(), ExceptionDispatchInfo.Capture(ex)); + Slon.Threading.SchedulingContext.Submit( + static state => ((ExceptionDispatchInfo)state!).Throw(), + ExceptionDispatchInfo.Capture(ex), + preferLocal: false); } } diff --git a/Slon/Runtime/CompilerServices/StackValue.cs b/Slon/Runtime/CompilerServices/StackValue.cs new file mode 100644 index 0000000..3431dd0 --- /dev/null +++ b/Slon/Runtime/CompilerServices/StackValue.cs @@ -0,0 +1,8 @@ +namespace Slon.Runtime.CompilerServices; + +// A byref to this wrapper is proven stack-only, allowing non-inlined callees to store +// reference-bearing values without checked write barriers. +ref struct StackValue(T value) where T : allows ref struct +{ + public T Value = value; +} diff --git a/Slon/Slon.csproj b/Slon/Slon.csproj index f7301cc..19e1074 100644 --- a/Slon/Slon.csproj +++ b/Slon/Slon.csproj @@ -1,17 +1,21 @@ - net10.0 + net10.0;net11.0 + runtime-async=on enable enable true preview true + $(MSBuildThisFileDirectory)../../Draghi $(TargetsForTfmSpecificBuildOutput);IncludeProjectReferenceDlls $(TargetsForTfmSpecificContentInPackage);IncludeReferenceAssemblies $(NoWarn);NU5131;DRAGHI001;SLONPG001;SLONPOOL001 + + @@ -20,7 +24,7 @@ - + diff --git a/Slon/SlonBatch.cs b/Slon/SlonBatch.cs index 5766bdc..c382d26 100644 --- a/Slon/SlonBatch.cs +++ b/Slon/SlonBatch.cs @@ -9,40 +9,62 @@ namespace Slon; /// public sealed partial class SlonBatch { - AdoBatchCore _batchCore; + internal readonly struct BatchCoreRef(SlonBatch owner) + : IAdoBatchCoreRef + { + public IAdoCommandExecutionOwner Owner => owner; + public ref AdoBatchCore GetField() + => ref owner._batchCore; + } + + AdoBatchCore _batchCore; SlonBatchCommands? _batchCommands; - unsafe SlonBatch(SlonConnection? connection, SlonDataSource? dataSource) + SlonBatch(SlonConnection? connection, SlonDataSource? dataSource) { - var fieldRef = FieldRef>.Create(&GetBatchCore, this); if (connection is not null) { - _batchCore = new(connection, fieldRef); + _batchCore = new(new(this), connection); _batchCore.Timeout = connection.DefaultCommandTimeout; } else if (dataSource is not null) { - _batchCore = new(dataSource, fieldRef); + _batchCore = new(new(this), dataSource); _batchCore.Timeout = dataSource.DefaultCommandTimeout; } else - _batchCore = new(fieldRef); + _batchCore = new(new(this)); } - unsafe SlonBatchCommands CreateBatchCommandCollection() - => new(FieldRef>.Create(&GetBatchCore, this)); + SlonBatchCommands CreateBatchCommandCollection() => new(this); + + internal ref AdoBatchCore BatchCore => ref _batchCore; - internal void OnFlowStarted(CommandFlow flow) + internal void OnFlowStarted(AdoCommandExecutionFlow flow) => _batchCore.OnFlowStarted(flow); - internal void OnFlowCompleting(CommandFlow flow, Exception? exception) + internal void OnFlowCompleting(AdoCommandExecutionFlow flow, Exception? exception) => _batchCore.OnFlowCompleting(flow, exception); - static ref AdoBatchCore GetBatchCore(SlonBatch instance) => ref instance._batchCore; + AdoCommandFlowOptions IAdoCommandExecutionOwner.CreateExecutionOptions( + DbParameterCollection? parameters, CommandBehavior behavior, + SlonDataSource.PgDbDependencies dependencies, SlonConnection? connection, + PgConnection pgConnection, TimeSpan? pendingTimeout, bool preparing) + => _batchCore.CreateAdoCommandFlowOptions( + [parameters], behavior, dependencies, connection, pgConnection, + pendingTimeout, preparing); + + void IAdoCommandExecutionOwner.OnFlowStarted(AdoCommandExecutionFlow flow) + => OnFlowStarted(flow); + + void IAdoCommandExecutionOwner.OnFlowCompleting( + AdoCommandExecutionFlow flow, Exception? exception) + => OnFlowCompleting(flow, exception); + } // Public surface & ADO.NET -public sealed partial class SlonBatch : DbBatch +public sealed partial class SlonBatch : DbBatch, IAdoCommandExecutionOwner { /// Initializes an unbound batch. public SlonBatch() : this(null, null) {} diff --git a/Slon/SlonBatchCommands.cs b/Slon/SlonBatchCommands.cs index 880368a..2c88c33 100644 --- a/Slon/SlonBatchCommands.cs +++ b/Slon/SlonBatchCommands.cs @@ -1,22 +1,21 @@ using System.Data.Common; -using Slon.Runtime.CompilerServices; namespace Slon; /// public sealed class SlonBatchCommands : DbBatchCommandCollection, IList { - readonly FieldRef> _batchRef; + readonly SlonBatch _batch; - internal SlonBatchCommands(FieldRef> batchRef) => _batchRef = batchRef; + internal SlonBatchCommands(SlonBatch batch) => _batch = batch; - ref AdoCommandList List => ref _batchRef.Invoke().Commands; + ref AdoCommandList List => ref _batch.BatchCore.Commands; /// public override int Count => List.Count; /// - public override bool IsReadOnly => _batchRef.Invoke().IsReadOnly; + public override bool IsReadOnly => _batch.BatchCore.IsReadOnly; /// IEnumerator IEnumerable.GetEnumerator() @@ -142,6 +141,6 @@ static SlonBatchCommand ThrowInvalidCastException(DbBatchCommand? value) => void ThrowIfReadOnly() { - _batchRef.Invoke().ThrowIfDisposedOrReadOnly(); + _batch.BatchCore.ThrowIfDisposedOrReadOnly(); } } diff --git a/Slon/SlonCommand.cs b/Slon/SlonCommand.cs index 5226525..edbbd94 100644 --- a/Slon/SlonCommand.cs +++ b/Slon/SlonCommand.cs @@ -11,7 +11,15 @@ namespace Slon; /// public sealed partial class SlonCommand { - AdoBatchCore _batchCore; + readonly struct BatchCoreRef(SlonCommand owner) + : IAdoBatchCoreRef + { + public IAdoCommandExecutionOwner Owner => owner; + public ref AdoBatchCore GetField() + => ref owner._batchCore; + } + + AdoBatchCore _batchCore; // Supporting state for implicit batching through SQL parsing. string _overallCommandText; @@ -19,28 +27,24 @@ public sealed partial class SlonCommand SlonParameters? _overallParameterCollection; bool _isOverallStateDirty; - internal unsafe SlonCommand(SlonConnection? connection, SlonDataSource? dataSource, string? commandText) + internal SlonCommand(SlonConnection? connection, SlonDataSource? dataSource, string? commandText) { GC.SuppressFinalize(this); _isOverallStateDirty = true; _overallCommandText = commandText ?? string.Empty; _overallCommandType = CommandType.Text; - var fieldRef = FieldRef>.Create(&GetBatchCore, this); if (connection is not null) { - _batchCore = new(connection, fieldRef); + _batchCore = new(new(this), connection); _batchCore.Timeout = connection.DefaultCommandTimeout; } else if (dataSource is not null) { - _batchCore = new(dataSource, fieldRef); + _batchCore = new(new(this), dataSource); _batchCore.Timeout = dataSource.DefaultCommandTimeout; } else - _batchCore = new(fieldRef); - - // ReSharper disable once AddressOfMarshalByRefObject - static ref AdoBatchCore GetBatchCore(SlonCommand instance) => ref instance._batchCore; + _batchCore = new(new(this)); } void SetupCommands() @@ -66,12 +70,27 @@ void Rebuild() } } - internal void OnFlowStarted(CommandFlow flow) + internal void OnFlowStarted(AdoCommandExecutionFlow flow) => _batchCore.OnFlowStarted(flow); - internal void OnFlowCompleting(CommandFlow flow, Exception? exception) + internal void OnFlowCompleting(AdoCommandExecutionFlow flow, Exception? exception) => _batchCore.OnFlowCompleting(flow, exception); + AdoCommandFlowOptions IAdoCommandExecutionOwner.CreateExecutionOptions( + DbParameterCollection? parameters, CommandBehavior behavior, + SlonDataSource.PgDbDependencies dependencies, SlonConnection? connection, + PgConnection pgConnection, TimeSpan? pendingTimeout, bool preparing) + => _batchCore.CreateAdoCommandFlowOptions( + [parameters], behavior, dependencies, connection, pgConnection, + pendingTimeout, preparing); + + void IAdoCommandExecutionOwner.OnFlowStarted(AdoCommandExecutionFlow flow) + => OnFlowStarted(flow); + + void IAdoCommandExecutionOwner.OnFlowCompleting( + AdoCommandExecutionFlow flow, Exception? exception) + => OnFlowCompleting(flow, exception); + struct AdoCommand : IAdoCommand { public void MakeReadOnly() { } @@ -85,7 +104,7 @@ public void MakeReadOnly() { } } // Public surface & ADO.NET -public sealed partial class SlonCommand : DbCommand +public sealed partial class SlonCommand : DbCommand, IAdoCommandExecutionOwner { /// Initializes an unbound command. public SlonCommand() : this(null, null, null) {} diff --git a/Slon/SlonConnection.cs b/Slon/SlonConnection.cs index d991987..cad73f3 100644 --- a/Slon/SlonConnection.cs +++ b/Slon/SlonConnection.cs @@ -485,7 +485,7 @@ ValueTask CloseOwnedStatements(bool async, TrackedCommand[] tracked) _proxy.Enqueue(flow); return AwaitCompletion(completion); - static async ValueTask AwaitCompletion(ValueTask completion) + static async ValueTask AwaitCompletion(ValueTask completion) => _ = await completion.ConfigureAwait(false); } diff --git a/Slon/SlonDataReader.cs b/Slon/SlonDataReader.cs index 52a42dd..14eaf39 100644 --- a/Slon/SlonDataReader.cs +++ b/Slon/SlonDataReader.cs @@ -16,6 +16,26 @@ public sealed partial class SlonDataReader { const CommandBehavior EnumerateCommandResultsBehavior = (CommandBehavior)int.MinValue; + interface IReadPolicy + { + static abstract ValueTask MoveNext( + ref CommandResult.RowEnumerator rows, CancellationToken cancellationToken); + } + + readonly struct DefaultReadPolicy : IReadPolicy + { + public static ValueTask MoveNext( + ref CommandResult.RowEnumerator rows, CancellationToken cancellationToken) + => rows.MoveNextAsync(); + } + + readonly struct CancelableReadPolicy : IReadPolicy + { + public static ValueTask MoveNext( + ref CommandResult.RowEnumerator rows, CancellationToken cancellationToken) + => rows.MoveNextAsync(cancellationToken); + } + int _state; ReaderState State { @@ -39,12 +59,12 @@ ReaderState State CommandResult.RowEnumerator _rowEnumerator; PgSerializerFieldReader _fieldReader; int _remainingResults; - CommandFlow.Enumerator _enumerator; + AdoCommandExecutionFlow.Enumerator _enumerator; long? _recordsAffected; SlonDataReader() { } - void Initialize(CommandFlow.Enumerator enumerator, CommandBehavior behavior, int remainingResults, + void Initialize(AdoCommandExecutionFlow.Enumerator enumerator, CommandBehavior behavior, int remainingResults, PgSerializerOptions serializerOptions, SlonConnection? connectionToClose, long? recordsAffected, bool hasCurrent) { @@ -84,7 +104,7 @@ static int GetResultLimit(CommandBehavior behavior, int commandCount) static bool ShouldEnumerateCommandResults(CommandBehavior behavior) => behavior.HasFlag(EnumerateCommandResultsBehavior); - static SlonDataReader CreateReader(CommandFlow.Enumerator enumerator, CommandBehavior behavior, + static SlonDataReader CreateReader(AdoCommandExecutionFlow.Enumerator enumerator, CommandBehavior behavior, int remainingResults, PgSerializerOptions serializerOptions, SlonConnection? connectionToClose, long? recordsAffected, bool hasCurrent) { @@ -99,7 +119,7 @@ static SlonDataReader CreateReader(CommandFlow.Enumerator enumerator, CommandBeh CommandResult? Current => _enumerator.Current; bool IsSequential => _rowBuffering is CommandResult.RowBuffering.Streaming; - internal static SlonDataReader Create(CommandBehavior behavior, CommandFlow flow, + internal static SlonDataReader Create(CommandBehavior behavior, AdoCommandExecutionFlow flow, PgSerializerOptions serializerOptions, SlonConnection? connectionToClose = null) { @@ -126,14 +146,16 @@ internal static SlonDataReader Create(CommandBehavior behavior, CommandFlow flow } } + [RuntimeAsyncMethodGeneration(false)] + [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder<>))] internal static async ValueTask CreateAsync(CommandBehavior behavior, - ValueTask flowTask, PgSerializerOptions serializerOptions, + ValueTask flowTask, PgSerializerOptions serializerOptions, CancellationToken cancellationToken = default, SlonConnection? connectionToClose = null, Activity? activity = null) where TReader : DbDataReader { Debug.Assert(typeof(TReader) == typeof(SlonDataReader) || typeof(TReader) == typeof(DbDataReader)); - CommandFlow.Enumerator enumerator = default; + AdoCommandExecutionFlow.Enumerator enumerator = default; try { var flow = await flowTask.ConfigureAwait(false); @@ -196,7 +218,8 @@ static void ApplyCompletion(CommandResult current, ref long? recordsAffected) { current.TryGetCommandComplete(out _); if (current.Error is null) - recordsAffected += current.RecordsAffected; + AccumulateRecordsAffected( + ref recordsAffected, current.BatchRecordsAffected); } } @@ -312,10 +335,18 @@ void ApplyCompletion(CommandResult current) { SurfaceCompletion(current); if (current.Error is null) - _recordsAffected += current.RecordsAffected; + AccumulateRecordsAffected( + ref _recordsAffected, current.BatchRecordsAffected); _currentCompletion = ResultCompletionState.Applied; } + static void AccumulateRecordsAffected(ref long? total, long current) + { + if (current < 0) + return; + total = total is { } existing ? checked(existing + current) : current; + } + bool ReadCore() { try @@ -339,33 +370,78 @@ bool ReadCore() } } - async Task ReadAsyncCore(CancellationToken cancellationToken) + Task ReadAsyncCore(CancellationToken cancellationToken) + where TPolicy : struct, IReadPolicy { try { Debug.Assert(_singleRowBehavior && _remainingResults is 0 || !_singleRowBehavior); - bool hasRow; if (_rowPresence is RowPresence.Prefetched) { _rowPresence = RowPresence.Present; - return true; + return Task.FromResult(true); } - else if (_singleRowBehavior && _rowPresence is RowPresence.Present) - { + + bool hasRow; + if (_singleRowBehavior && _rowPresence is RowPresence.Present) hasRow = false; - } else { - hasRow = await _rowEnumerator.MoveNextAsync(cancellationToken) - .ConfigureAwait(false); + var moveNext = TPolicy.MoveNext(ref _rowEnumerator, cancellationToken); + if (!moveNext.IsCompletedSuccessfully) + return AwaitMoveNext(this, moveNext); + hasRow = moveNext.GetAwaiter().GetResult(); } - if (hasRow) - return ProcessReadResult(hasRow: true); + return CompleteRead(this, hasRow); + } + catch (Exception ex) + { + return Task.FromException(AdoException.Project(ex)); + } + } - if (Current is { IsComplete: false } current) + static Task CompleteRead(SlonDataReader reader, bool hasRow) + { + if (hasRow) + return Task.FromResult(reader.ProcessReadResult(hasRow: true)); + + if (reader.Current is { IsComplete: false } current) + { + var completion = current.CompleteAsync(); + if (!completion.IsCompletedSuccessfully) + return AwaitCompletion(reader, completion); + completion.GetAwaiter().GetResult(); + } + return Task.FromResult(reader.ProcessReadResult(hasRow: false)); + } + + static async Task AwaitMoveNext( + SlonDataReader reader, ValueTask moveNext) + { + try + { + var hasRow = await moveNext.ConfigureAwait(false); + if (hasRow) + return reader.ProcessReadResult(hasRow: true); + if (reader.Current is { IsComplete: false } current) await current.CompleteAsync().ConfigureAwait(false); - return ProcessReadResult(hasRow: false); + return reader.ProcessReadResult(hasRow: false); + } + catch (Exception ex) + { + AdoException.Throw(ex); + return default; + } + } + + static async Task AwaitCompletion( + SlonDataReader reader, ValueTask completion) + { + try + { + await completion.ConfigureAwait(false); + return reader.ProcessReadResult(hasRow: false); } catch (Exception ex) { @@ -451,6 +527,7 @@ void DisposeEnumerator(out bool ownsCleanup) } } + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] async ValueTask DisposeEnumeratorAsync() { @@ -472,7 +549,7 @@ async ValueTask DisposeEnumeratorAsync() } } - (CommandResult.RowEnumerator Rows, CommandFlow.Enumerator Results) BeginEnumeratorDisposal() + (CommandResult.RowEnumerator Rows, AdoCommandExecutionFlow.Enumerator Results) BeginEnumeratorDisposal() { if (_enumeratorDisposalActive) ThrowHelper.ThrowInvalidOperation("Invalid concurrent call."); @@ -572,6 +649,7 @@ void CloseCore(bool resetForReuse) } } + [RuntimeAsyncMethodGeneration(false)] [AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))] async ValueTask CloseAsyncCore(bool resetForReuse) { @@ -763,6 +841,8 @@ public override Task NextResultAsync(CancellationToken cancellationToken) { if (GetExceptionIfClosedOrDisposed() is { } exception) return Task.FromException(exception); + if (_remainingResults is 0 && _currentCompletion is ResultCompletionState.Applied) + return Task.FromResult(false); return NextResultAsyncCore(cancellationToken); } @@ -781,7 +861,9 @@ public override Task ReadAsync(CancellationToken cancellationToken) if (GetExceptionIfClosedOrDisposed() is { } exception) return Task.FromException(exception); - return ReadAsyncCore(cancellationToken); + return cancellationToken.CanBeCanceled + ? ReadAsyncCore(cancellationToken) + : ReadAsyncCore(default); } /// @@ -841,16 +923,9 @@ public override Task IsDBNullAsync(int ordinal, CancellationToken cancella : IsDBNullAsyncCore(ordinal, cancellationToken); } - /// Returns a nested data reader for the requested column. - /// The zero-based column ordinal. - /// Nested data readers are not supported. - /// A data reader. - public new SlonDataReader GetData(int ordinal) - => throw new NotSupportedException("Nested data readers are not supported."); - /// protected override DbDataReader GetDbDataReader(int ordinal) - => GetData(ordinal); + => throw new NotSupportedException("Nested data readers are not supported."); /// Reads the complete field at the specified ordinal as a byte array. /// The zero-based column ordinal. @@ -1012,6 +1087,14 @@ protected override void Dispose(bool disposing) return; State = ReaderState.Disposed; + if (_remainingResults is 0 + && _currentCompletion is ResultCompletionState.Applied + && _enumerator.IsDefault + && _connectionToClose is null) + { + Reset(); + return; + } CloseCore(resetForReuse: true); } @@ -1022,6 +1105,14 @@ public override ValueTask DisposeAsync() return new(); State = ReaderState.Disposed; + if (_remainingResults is 0 + && _currentCompletion is ResultCompletionState.Applied + && _enumerator.IsDefault + && _connectionToClose is null) + { + Reset(); + return default; + } return CloseAsyncCore(resetForReuse: true); } } diff --git a/Slon/SlonDataSource.cs b/Slon/SlonDataSource.cs index 3e77245..f6d94ff 100644 --- a/Slon/SlonDataSource.cs +++ b/Slon/SlonDataSource.cs @@ -81,7 +81,7 @@ internal void ReportTransactionDisposeRollbackFailure(Exception exception) // The multiplexed path lets the pool select a wire before materializing connection-local command // state. A rejected candidate rolls that attempt back and reuses the still-unqueued flow shell. - internal CommandFlow EnqueueCommands(CommandFlow flow, TimeSpan pendingTimeout) + internal AdoCommandExecutionFlow EnqueueCommands(AdoCommandExecutionFlow flow, TimeSpan pendingTimeout) { try { @@ -95,8 +95,8 @@ internal CommandFlow EnqueueCommands(CommandFlow flow, TimeSpan pendingTimeout) } } - internal async ValueTask EnqueueCommandsAsync( - CommandFlow flow, TimeSpan pendingTimeout, CancellationToken cancellationToken) + internal async ValueTask EnqueueCommandsAsync( + AdoCommandExecutionFlow flow, TimeSpan pendingTimeout, CancellationToken cancellationToken) { try { @@ -111,7 +111,7 @@ await _connectionPool.GetAsync(static (ctx, f) => TrySchedule(ctx, f), flow, pen } } - static bool TrySchedule(ConnectionCandidate context, CommandFlow flow) + static bool TrySchedule(ConnectionCandidate context, AdoCommandExecutionFlow flow) { var enqueueOptions = context.IsIdleCandidate ? FlowEnqueueOptions.AllowMigration diff --git a/Slon/SlonDataSourceOptions.cs b/Slon/SlonDataSourceOptions.cs index d1a99ad..0f34a1f 100644 --- a/Slon/SlonDataSourceOptions.cs +++ b/Slon/SlonDataSourceOptions.cs @@ -112,7 +112,7 @@ public int PoolSize /// /// DataRows larger than this may cross the decoder boundary before their complete body has arrived. /// - public int DataRowStreamingThreshold { get; init; } = BackendMessageBatch.Segmenter.DefaultDataRowStreamingThreshold; + public int DataRowStreamingThreshold { get; init; } = BackendMessageCursor.DefaultDataRowStreamingThreshold; /// Configures which connection state is reset when an exclusive scope is released. internal PgSessionResetOptions SessionReset { get; init; } = new(); /// diff --git a/Slon/Text/EncodedCString.cs b/Slon/Text/EncodedCString.cs index 1260feb..009248f 100644 --- a/Slon/Text/EncodedCString.cs +++ b/Slon/Text/EncodedCString.cs @@ -35,6 +35,12 @@ public bool ValueEquals(EncodedCString other) public static implicit operator EncodedCString(string value) => new(value); + internal static void Assign(ref EncodedCString destination, in EncodedCString value) + { + if (!ReferenceEquals(destination._core, value._core)) + Unsafe.AsRef(in destination._core) = value._core; + } + // Used for long lived strings that may have to be re-encoded (but usually wont), thread-safe. [DebuggerDisplay("{_value,nq}")] sealed class Core(string value) diff --git a/Slon/Threading/Scheduler.cs b/Slon/Threading/Scheduler.cs new file mode 100644 index 0000000..4016bbd --- /dev/null +++ b/Slon/Threading/Scheduler.cs @@ -0,0 +1,20 @@ +namespace Slon.Threading; + +abstract class Scheduler +{ + public abstract void SubmitDetached(Action action, object? state, bool preferLocal = true); +} + +sealed class DelegatedScheduler : Scheduler +{ + readonly Action, object?, bool> _submitDetached; + + public DelegatedScheduler(Action, object?, bool> submitDetached) + { + ArgumentNullException.ThrowIfNull(submitDetached); + _submitDetached = submitDetached; + } + + public override void SubmitDetached(Action action, object? state, bool preferLocal = true) + => _submitDetached(action, state, preferLocal); +} diff --git a/Slon/Threading/SchedulingContext.cs b/Slon/Threading/SchedulingContext.cs new file mode 100644 index 0000000..5ee0517 --- /dev/null +++ b/Slon/Threading/SchedulingContext.cs @@ -0,0 +1,38 @@ +namespace Slon.Threading; + +// Temporary host seam until the runtime provides the ambient scheduler lookup directly. +static class SchedulingContext +{ + static Action, object?, bool>? _submit; + static Action, object?, bool>? _submitDetached; + + public static void SetDispatch( + Action, object?, bool>? submit, + Action, object?, bool>? submitDetached) + { + _submit = submit; + _submitDetached = submitDetached; + } + + internal static void Submit(Action action, object? state, bool preferLocal) + { + var submit = Volatile.Read(ref _submit); + if (submit is null) + ThreadPool.QueueUserWorkItem(action, state, preferLocal); + else + submit(action, state, preferLocal); + } + + internal static void SubmitDetached(Action action, object? state, bool preferLocal) + => _ = TrySubmitDetached(action, state, preferLocal); + + internal static bool TrySubmitDetached(Action action, object? state, bool preferLocal) + { + var submit = Volatile.Read(ref _submitDetached); + if (submit is null) + return ThreadPool.UnsafeQueueUserWorkItem(action, state, preferLocal); + + submit(action, state, preferLocal); + return true; + } +} diff --git a/Slon/Threading/Tasks/Sources/ContinuationDispatcher.cs b/Slon/Threading/Tasks/Sources/ContinuationDispatcher.cs index a33309a..a84049a 100644 --- a/Slon/Threading/Tasks/Sources/ContinuationDispatcher.cs +++ b/Slon/Threading/Tasks/Sources/ContinuationDispatcher.cs @@ -53,7 +53,7 @@ public void SignalCompletion(bool runContinuationsAsynchronously) { if (runContinuationsAsynchronously) { - ThreadPool.UnsafeQueueUserWorkItem(continuation, _continuationState, preferLocal: true); + SchedulingContext.SubmitDetached(continuation, _continuationState, preferLocal: true); } else { @@ -136,11 +136,11 @@ public void OnCompleted(Action continuation, object? state, ValueTaskSo switch (capturedContext) { case null: - ThreadPool.UnsafeQueueUserWorkItem(continuation, state, preferLocal: true); + SchedulingContext.SubmitDetached(continuation, state, preferLocal: true); break; case ExecutionContext: - ThreadPool.QueueUserWorkItem(continuation, state, preferLocal: true); + SchedulingContext.Submit(continuation, state, preferLocal: true); break; default: @@ -204,7 +204,7 @@ static void InvokeWithContext(object capturedContext, Action continuati { try { - ThreadPool.QueueUserWorkItem(continuation, continuationState, preferLocal: true); + SchedulingContext.Submit(continuation, continuationState, preferLocal: true); } finally { diff --git a/Slon/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs b/Slon/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs index 4b0e11c..25ca297 100644 --- a/Slon/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs +++ b/Slon/Threading/Tasks/Sources/ManualResetValueTaskSourceCore.cs @@ -37,18 +37,8 @@ struct ManualResetValueTaskSourceCore TResult? _result; /// The current version of this value, used to help prevent misuse. short _version; - /// Whether the current operation has completed (this can mean it's still completing, which is why _continuation is also checked). + /// Whether the current operation has completed. bool _completed; - /// Whether concurrent completions are handled correctly. - bool _canCompleteConcurrently; - - /// Gets or sets whether concurrent completions are handled correctly. - /// Enabling this allows the use of TrySet methods and makes the Set methods thread safe. - public bool CanCompleteConcurrently - { - get => _canCompleteConcurrently; - set => _canCompleteConcurrently = value; - } /// Resets to prepare for the next operation. public void Reset() @@ -60,8 +50,7 @@ public void Reset() _capturedContext = null; _error = null; _result = default; - // Release, written last: a concurrent completer's CAS (acquire) observing false then sees - // the resets above. Without it a racing TrySet could complete a half-reset source. + // Release, written last so the next tenure observes all reset state. Volatile.Write(ref _completed, false); } @@ -69,12 +58,10 @@ public void Reset() /// The result. public void SetResult(TResult result) { - var canCompleteConcurrently = _canCompleteConcurrently; - if (_completed || (canCompleteConcurrently && Interlocked.CompareExchange(ref _completed, true, false) is not false)) + if (_completed) ThrowInvalidOperationException(); - if (!canCompleteConcurrently) - _completed = true; _result = result; + Volatile.Write(ref _completed, true); new ContinuationDispatcher(ref _continuation, ref _continuationState, ref _capturedContext) .SignalCompletion(runContinuationsAsynchronously: false); } @@ -84,12 +71,10 @@ public void SetResult(TResult result) /// whether to force continuations to run asynchronously this call. public void SetResult(TResult result, bool runContinuationsAsynchronously) { - var canCompleteConcurrently = _canCompleteConcurrently; - if (_completed || (canCompleteConcurrently && Interlocked.CompareExchange(ref _completed, true, false) is not false)) + if (_completed) ThrowInvalidOperationException(); - if (!canCompleteConcurrently) - _completed = true; _result = result; + Volatile.Write(ref _completed, true); new ContinuationDispatcher(ref _continuation, ref _continuationState, ref _capturedContext) .SignalCompletion(runContinuationsAsynchronously); } @@ -98,12 +83,10 @@ public void SetResult(TResult result, bool runContinuationsAsynchronously) /// The exception. public void SetException(Exception error) { - var canCompleteConcurrently = _canCompleteConcurrently; - if (_completed || (canCompleteConcurrently && Interlocked.CompareExchange(ref _completed, true, false) is not false)) + if (_completed) ThrowInvalidOperationException(); - if (!canCompleteConcurrently) - _completed = true; _error = ExceptionDispatchInfo.Capture(error); + Volatile.Write(ref _completed, true); new ContinuationDispatcher(ref _continuation, ref _continuationState, ref _capturedContext) .SignalCompletion(false); } @@ -113,72 +96,12 @@ public void SetException(Exception error) /// whether to force continuations to run asynchronously this call. public void SetException(Exception error, bool runContinuationsAsynchronously) { - var canCompleteConcurrently = _canCompleteConcurrently; - if (_completed || (canCompleteConcurrently && Interlocked.CompareExchange(ref _completed, true, false) is not false)) - ThrowInvalidOperationException(); - if (!canCompleteConcurrently) - _completed = true; - _error = ExceptionDispatchInfo.Capture(error); - new ContinuationDispatcher(ref _continuation, ref _continuationState, ref _capturedContext) - .SignalCompletion(runContinuationsAsynchronously); - } - - /// Completes with a successful result. - /// The result. - public bool TrySetResult(TResult result) - { - if (!_canCompleteConcurrently) - ThrowInvalidOperationException(); - if (_completed || Interlocked.CompareExchange(ref _completed, true, false) is not false) - return false; - _result = result; - new ContinuationDispatcher(ref _continuation, ref _continuationState, ref _capturedContext) - .SignalCompletion(false); - return true; - } - - /// Completes with a successful result. - /// The result. - /// whether to force continuations to run asynchronously this call. - public bool TrySetResult(TResult result, bool runContinuationsAsynchronously) - { - if (!_canCompleteConcurrently) - ThrowInvalidOperationException(); - if (_completed || Interlocked.CompareExchange(ref _completed, true, false) is not false) - return false; - _result = result; - new ContinuationDispatcher(ref _continuation, ref _continuationState, ref _capturedContext) - .SignalCompletion(runContinuationsAsynchronously); - return true; - } - - /// Completes with an error. - /// The exception. - public bool TrySetException(Exception error) - { - if (!_canCompleteConcurrently) - ThrowInvalidOperationException(); - if (_completed || Interlocked.CompareExchange(ref _completed, true, false) is not false) - return false; - _error = ExceptionDispatchInfo.Capture(error); - new ContinuationDispatcher(ref _continuation, ref _continuationState, ref _capturedContext) - .SignalCompletion(false); - return true; - } - - /// Completes with an error. - /// The exception. - /// whether to force continuations to run asynchronously this call. - public bool TrySetException(Exception error, bool runContinuationsAsynchronously) - { - if (!_canCompleteConcurrently) + if (_completed) ThrowInvalidOperationException(); - if (_completed || Interlocked.CompareExchange(ref _completed, true, false) is not false) - return false; _error = ExceptionDispatchInfo.Capture(error); + Volatile.Write(ref _completed, true); new ContinuationDispatcher(ref _continuation, ref _continuationState, ref _capturedContext) .SignalCompletion(runContinuationsAsynchronously); - return true; } /// Gets the operation version. @@ -193,7 +116,7 @@ public ValueTaskSourceStatus GetStatus(short token) ThrowInvalidOperationException(); } return - _continuation is null || !_completed ? ValueTaskSourceStatus.Pending : + _continuation is null || !Volatile.Read(ref _completed) ? ValueTaskSourceStatus.Pending : _error is null ? ValueTaskSourceStatus.Succeeded : _error.SourceException is OperationCanceledException ? ValueTaskSourceStatus.Canceled : ValueTaskSourceStatus.Faulted; diff --git a/Slon/Transport/TransportConnection.cs b/Slon/Transport/TransportConnection.cs index 4a05ae8..5904825 100644 --- a/Slon/Transport/TransportConnection.cs +++ b/Slon/Transport/TransportConnection.cs @@ -1,4 +1,5 @@ using System.IO.Pipelines; +using Slon.Threading; namespace Slon.Transport; @@ -30,6 +31,7 @@ public readonly struct ResumableWrite(ResumeSignal? signal, TimeSpan timeout) // The protocol completes both endpoints after all borrowed buffers have been returned. public abstract PipeReader Reader { get; } public abstract PipeWriter Writer { get; } + internal virtual Scheduler? Scheduler => null; // Classifies transport-specific exceptions that mean the established byte stream was lost. public virtual bool IsConnectionLost(Exception exception) => false; diff --git a/eng/scheduler-chaos b/eng/scheduler-chaos new file mode 100755 index 0000000..a557010 --- /dev/null +++ b/eng/scheduler-chaos @@ -0,0 +1,8 @@ +#!/bin/sh +set -eu + +here=$(CDPATH= cd -- "$(dirname -- "$0")" && pwd) +binary="${TMPDIR:-/tmp}/slon-scheduler-chaos" + +cc -O2 -pthread "$here/scheduler-chaos.c" -o "$binary" +exec "$binary" "$@" diff --git a/eng/scheduler-chaos.c b/eng/scheduler-chaos.c new file mode 100644 index 0000000..ace2c1a --- /dev/null +++ b/eng/scheduler-chaos.c @@ -0,0 +1,99 @@ +#include +#include +#include +#include +#include +#include +#include +#include + +static _Atomic int stopping; + +static void stop(int signal) +{ + (void)signal; + atomic_store_explicit(&stopping, 1, memory_order_relaxed); +} + +static uint32_t next(uint32_t *state) +{ + uint32_t value = *state; + value ^= value << 13; + value ^= value >> 17; + value ^= value << 5; + return *state = value; +} + +static void pause_ns(long nanoseconds) +{ + struct timespec duration = { .tv_nsec = nanoseconds }; + nanosleep(&duration, NULL); +} + +static void *churn(void *argument) +{ + uint32_t state = (uint32_t)(uintptr_t)argument ^ (uint32_t)clock(); + volatile uint64_t work = state; + + while (!atomic_load_explicit(&stopping, memory_order_relaxed)) + { + switch (next(&state) % 10) + { + case 0: + case 1: + sched_yield(); + break; + case 2: + case 3: + case 4: + pause_ns(1000L * (50 + next(&state) % 1950)); + break; + default: + { + // Keep a couple of cores busy in aggregate. macOS affinity tags are + // task-local hints, so light pressure could simply be placed away + // from the independently running test process. + const uint32_t iterations = 1000 + next(&state) % 200000; + for (uint32_t i = 0; i < iterations; i++) + work = work * 6364136223846793005ULL + 1; + break; + } + } + } + + return (void *)(uintptr_t)work; +} + +int main(int argc, char **argv) +{ + long cores = sysconf(_SC_NPROCESSORS_ONLN); + int threads = argc > 1 ? atoi(argv[1]) : (int)cores; + if (threads < 1) + { + fprintf(stderr, "usage: %s [threads]\n", argv[0]); + return 2; + } + + signal(SIGINT, stop); + signal(SIGTERM, stop); + + pthread_t *workers = calloc((size_t)threads, sizeof(*workers)); + if (workers == NULL) + return 1; + + fprintf(stderr, "scheduler chaos: %d threads, distributed pressure; Ctrl-C to stop\n", threads); + for (int i = 0; i < threads; i++) + { + if (pthread_create(&workers[i], NULL, churn, (void *)(uintptr_t)(i + 1)) != 0) + { + atomic_store(&stopping, 1); + threads = i; + break; + } + } + + for (int i = 0; i < threads; i++) + pthread_join(workers[i], NULL); + free(workers); + return 0; +} diff --git a/provenance/initial-publication.tag.ots b/provenance/initial-publication.tag.ots index ab2ccbf..6957b0d 100644 Binary files a/provenance/initial-publication.tag.ots and b/provenance/initial-publication.tag.ots differ