From 94ba18059cb5778cbf945a270bf28f0c88c09b98 Mon Sep 17 00:00:00 2001 From: Yordis Prieto Date: Sat, 18 Jul 2026 22:24:01 -0400 Subject: [PATCH] fix(grpc): preserve graceful streaming shutdown Signed-off-by: Yordis Prieto --- .../Enumerator.AllSubscription.Tests.cs | 19 ++- ...numerator.AllSubscriptionFiltered.Tests.cs | 23 ++- .../Enumerator.StreamSubscription.Tests.cs | 20 ++- .../GrpcServerShutdownInterceptorTests.cs | 148 ++++++++++++++++++ .../Grpc/GrpcStreamLifetimeMiddlewareTests.cs | 135 ++++++++++++++++ src/EventStore.Core/ClusterVNodeStartup.cs | 3 + .../Enumerators/Enumerator.AllSubscription.cs | 15 ++ .../Enumerator.AllSubscriptionFiltered.cs | 15 ++ .../Enumerator.StreamSubscription.cs | 15 ++ .../Services/Transport/Grpc/Constants.cs | 1 + .../Grpc/GrpcServerShutdownInterceptor.cs | 61 ++++++++ .../Grpc/GrpcStreamLifetimeMiddleware.cs | 39 +++++ .../Services/Transport/Grpc/RpcExceptions.cs | 5 + .../Services/Transport/Grpc/Streams.Read.cs | 11 +- 14 files changed, 496 insertions(+), 14 deletions(-) create mode 100644 src/EventStore.Core.Tests/Services/Transport/Grpc/GrpcServerShutdownInterceptorTests.cs create mode 100644 src/EventStore.Core.Tests/Services/Transport/Grpc/GrpcStreamLifetimeMiddlewareTests.cs create mode 100644 src/EventStore.Core/Services/Transport/Grpc/GrpcServerShutdownInterceptor.cs create mode 100644 src/EventStore.Core/Services/Transport/Grpc/GrpcStreamLifetimeMiddleware.cs diff --git a/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.AllSubscription.Tests.cs b/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.AllSubscription.Tests.cs index 5625440754..94f3dc52a9 100644 --- a/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.AllSubscription.Tests.cs +++ b/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.AllSubscription.Tests.cs @@ -20,7 +20,8 @@ public partial class EnumeratorTests private static EnumeratorWrapper CreateAllSubscription( IPublisher publisher, Position? checkpoint, - ClaimsPrincipal user = null) + ClaimsPrincipal user = null, + CancellationToken cancellationToken = default) { return new EnumeratorWrapper(new Enumerator.AllSubscription( @@ -30,7 +31,7 @@ private static EnumeratorWrapper CreateAllSubscription( resolveLinks: false, user: user ?? SystemAccounts.System, requiresLeader: false, - cancellationToken: CancellationToken.None)); + cancellationToken: cancellationToken)); } [TestFixture(typeof(LogFormat.V2), typeof(string))] @@ -83,6 +84,20 @@ public async Task should_receive_live_caught_up_message_immediately() Assert.True(await sub.GetNext() is SubscriptionConfirmation); Assert.True(await sub.GetNext() is CaughtUp); } + + [Test] + public async Task cancellation_preserves_the_callers_token() + { + using var cancellation = new CancellationTokenSource(); + await using var sub = CreateAllSubscription(_publisher, Position.End, cancellationToken: cancellation.Token); + + Assert.That(await sub.GetNext(), Is.InstanceOf()); + Assert.That(await sub.GetNext(), Is.InstanceOf()); + cancellation.Cancel(); + + var exception = Assert.ThrowsAsync(() => sub.GetNext()); + Assert.That(exception!.CancellationToken, Is.EqualTo(cancellation.Token)); + } } [TestFixture(typeof(LogFormat.V2), typeof(string))] diff --git a/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.AllSubscriptionFiltered.Tests.cs b/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.AllSubscriptionFiltered.Tests.cs index dd7d389d4e..1c9ee48d41 100644 --- a/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.AllSubscriptionFiltered.Tests.cs +++ b/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.AllSubscriptionFiltered.Tests.cs @@ -23,7 +23,8 @@ private static EnumeratorWrapper CreateAllSubscriptionFiltered( IEventFilter eventFilter = null, uint? maxSearchWindow = null, uint checkpointIntervalMultiplier = 1, - ClaimsPrincipal user = null) + ClaimsPrincipal user = null, + CancellationToken cancellationToken = default) { return new EnumeratorWrapper(new Enumerator.AllSubscriptionFiltered( @@ -36,7 +37,7 @@ private static EnumeratorWrapper CreateAllSubscriptionFiltered( requiresLeader: false, maxSearchWindow: maxSearchWindow, checkpointIntervalMultiplier: checkpointIntervalMultiplier, - cancellationToken: CancellationToken.None)); + cancellationToken: cancellationToken)); } @@ -101,6 +102,24 @@ public async Task should_receive_live_caught_up_message_immediately() Assert.True(await sub.GetNext() is SubscriptionConfirmation); Assert.True(await sub.GetNext() is CaughtUp); } + + [Test] + public async Task cancellation_preserves_the_callers_token() + { + using var cancellation = new CancellationTokenSource(); + await using var sub = CreateAllSubscriptionFiltered( + _publisher, + Position.End, + EventFilter.EventType.Prefixes(false, "type1"), + cancellationToken: cancellation.Token); + + Assert.That(await sub.GetNext(), Is.InstanceOf()); + Assert.That(await sub.GetNext(), Is.InstanceOf()); + cancellation.Cancel(); + + var exception = Assert.ThrowsAsync(() => sub.GetNext()); + Assert.That(exception!.CancellationToken, Is.EqualTo(cancellation.Token)); + } } [TestFixture(typeof(LogFormat.V2), typeof(string))] diff --git a/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.StreamSubscription.Tests.cs b/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.StreamSubscription.Tests.cs index eaf8561285..b81e245e1c 100644 --- a/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.StreamSubscription.Tests.cs +++ b/src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.StreamSubscription.Tests.cs @@ -20,7 +20,8 @@ private static EnumeratorWrapper CreateStreamSubscription( IPublisher publisher, string streamName, StreamRevision? checkpoint = null, - ClaimsPrincipal user = null) + ClaimsPrincipal user = null, + CancellationToken cancellationToken = default) { return new EnumeratorWrapper(new Enumerator.StreamSubscription( @@ -31,7 +32,7 @@ private static EnumeratorWrapper CreateStreamSubscription( resolveLinks: false, user: user ?? SystemAccounts.System, requiresLeader: false, - cancellationToken: CancellationToken.None)); + cancellationToken: cancellationToken)); } [TestFixture(typeof(LogFormat.V2), typeof(string))] @@ -84,5 +85,20 @@ public async Task should_receive_live_caught_up_message_immediately() Assert.True(await enumerator.GetNext() is SubscriptionConfirmation); Assert.True(await enumerator.GetNext() is CaughtUp); } + + [Test] + public async Task cancellation_preserves_the_callers_token() + { + using var cancellation = new CancellationTokenSource(); + await using var enumerator = CreateStreamSubscription( + _publisher, streamName: "test-stream1", StreamRevision.End, cancellationToken: cancellation.Token); + + Assert.That(await enumerator.GetNext(), Is.InstanceOf()); + Assert.That(await enumerator.GetNext(), Is.InstanceOf()); + cancellation.Cancel(); + + var exception = Assert.ThrowsAsync(() => enumerator.GetNext()); + Assert.That(exception!.CancellationToken, Is.EqualTo(cancellation.Token)); + } } } diff --git a/src/EventStore.Core.Tests/Services/Transport/Grpc/GrpcServerShutdownInterceptorTests.cs b/src/EventStore.Core.Tests/Services/Transport/Grpc/GrpcServerShutdownInterceptorTests.cs new file mode 100644 index 0000000000..c70392df6b --- /dev/null +++ b/src/EventStore.Core.Tests/Services/Transport/Grpc/GrpcServerShutdownInterceptorTests.cs @@ -0,0 +1,148 @@ +using System; +using System.Collections.Generic; +using System.Linq; +using System.Threading; +using System.Threading.Tasks; +using EventStore.Core.Services.Transport.Grpc; +using Grpc.Core; +using Microsoft.Extensions.Hosting; +using NUnit.Framework; + +namespace EventStore.Core.Tests.Services.Transport.Grpc; + +[TestFixture] +public class GrpcServerShutdownInterceptorTests +{ + [TestCase(MethodType.ClientStreaming)] + [TestCase(MethodType.ServerStreaming)] + [TestCase(MethodType.DuplexStreaming)] + public void shutdown_cancellation_is_retryable_for_every_streaming_shape(MethodType methodType) + { + using var lifetime = new TestHostLifetime(); + using var callCancellation = CancellationTokenSource.CreateLinkedTokenSource(lifetime.ApplicationStopping); + lifetime.StopApplication(); + var context = new TestServerCallContext(callCancellation.Token); + var interceptor = new GrpcServerShutdownInterceptor(lifetime); + + var exception = Assert.ThrowsAsync(() => Invoke(interceptor, methodType, context)); + + Assert.That(exception!.StatusCode, Is.EqualTo(StatusCode.Unavailable)); + Assert.That( + exception.Trailers.Select(entry => (entry.Key, entry.Value)), + Does.Contain(( + EventStore.Core.Services.Transport.Grpc.Constants.Exceptions.ExceptionKey, + EventStore.Core.Services.Transport.Grpc.Constants.Exceptions.ServerShuttingDown))); + } + + [Test] + public void client_cancellation_is_not_changed_when_server_is_running() + { + using var lifetime = new TestHostLifetime(); + using var callCancellation = new CancellationTokenSource(); + callCancellation.Cancel(); + var context = new TestServerCallContext(callCancellation.Token); + var interceptor = new GrpcServerShutdownInterceptor(lifetime); + + var exception = Assert.CatchAsync(() => + Invoke(interceptor, MethodType.ServerStreaming, context)); + + Assert.That(exception!.CancellationToken, Is.EqualTo(callCancellation.Token)); + } + + [Test] + public async Task successful_streaming_call_passes_through() + { + using var lifetime = new TestHostLifetime(); + var context = new TestServerCallContext(CancellationToken.None); + var interceptor = new GrpcServerShutdownInterceptor(lifetime); + var continuationCalled = false; + + await interceptor.ServerStreamingServerHandler( + "request", + new TestServerStreamWriter(), + context, + (_, _, _) => + { + continuationCalled = true; + return Task.CompletedTask; + }); + + Assert.That(continuationCalled, Is.True); + } + + private static async Task Invoke( + GrpcServerShutdownInterceptor interceptor, + MethodType methodType, + ServerCallContext context) + { + switch (methodType) + { + case MethodType.ClientStreaming: + await interceptor.ClientStreamingServerHandler( + new TestAsyncStreamReader(), + context, + (_, callContext) => Task.FromCanceled(callContext.CancellationToken)); + break; + case MethodType.ServerStreaming: + await interceptor.ServerStreamingServerHandler( + "request", + new TestServerStreamWriter(), + context, + (_, _, callContext) => Task.FromCanceled(callContext.CancellationToken)); + break; + case MethodType.DuplexStreaming: + await interceptor.DuplexStreamingServerHandler( + new TestAsyncStreamReader(), + new TestServerStreamWriter(), + context, + (_, _, callContext) => Task.FromCanceled(callContext.CancellationToken)); + break; + default: + throw new ArgumentOutOfRangeException(nameof(methodType), methodType, null); + } + } + + private sealed class TestAsyncStreamReader : IAsyncStreamReader + { + public T Current => default!; + public Task MoveNext(CancellationToken cancellationToken) => Task.FromResult(false); + } + + private sealed class TestServerStreamWriter : IServerStreamWriter + { + public WriteOptions WriteOptions { get; set; } + public Task WriteAsync(T message) => Task.CompletedTask; + } + + private sealed class TestServerCallContext(CancellationToken cancellationToken) : ServerCallContext + { + protected override string MethodCore => "/service/method"; + protected override string HostCore => "host"; + protected override string PeerCore => "peer"; + protected override DateTime DeadlineCore => DateTime.MaxValue; + protected override Metadata RequestHeadersCore { get; } = new(); + protected override CancellationToken CancellationTokenCore => cancellationToken; + protected override Metadata ResponseTrailersCore { get; } = new(); + protected override Status StatusCore { get; set; } + protected override WriteOptions WriteOptionsCore { get; set; } + protected override AuthContext AuthContextCore { get; } = + new(string.Empty, new Dictionary>()); + protected override IDictionary UserStateCore { get; } = + new Dictionary(); + protected override Task WriteResponseHeadersAsyncCore(Metadata responseHeaders) => Task.CompletedTask; + protected override ContextPropagationToken CreatePropagationTokenCore(ContextPropagationOptions options) => + throw new NotSupportedException(); + } + + private sealed class TestHostLifetime : IHostApplicationLifetime, IDisposable + { + private readonly CancellationTokenSource _stopping = new(); + + public CancellationToken ApplicationStarted => CancellationToken.None; + public CancellationToken ApplicationStopping => _stopping.Token; + public CancellationToken ApplicationStopped => CancellationToken.None; + + public void StopApplication() => _stopping.Cancel(); + public void Dispose() => _stopping.Dispose(); + } +} diff --git a/src/EventStore.Core.Tests/Services/Transport/Grpc/GrpcStreamLifetimeMiddlewareTests.cs b/src/EventStore.Core.Tests/Services/Transport/Grpc/GrpcStreamLifetimeMiddlewareTests.cs new file mode 100644 index 0000000000..3ba9fd8e42 --- /dev/null +++ b/src/EventStore.Core.Tests/Services/Transport/Grpc/GrpcStreamLifetimeMiddlewareTests.cs @@ -0,0 +1,135 @@ +using System; +using System.Threading; +using System.Threading.Tasks; +using EventStore.Core.Services.Transport.Grpc; +using Grpc.AspNetCore.Server; +using Grpc.Core; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Hosting; +using NUnit.Framework; + +namespace EventStore.Core.Tests.Services.Transport.Grpc; + +[TestFixture] +public class GrpcStreamLifetimeMiddlewareTests +{ + [TestCase(MethodType.ClientStreaming)] + [TestCase(MethodType.ServerStreaming)] + [TestCase(MethodType.DuplexStreaming)] + public async Task streaming_calls_observe_host_shutdown(MethodType methodType) + { + using var lifetime = new TestHostLifetime(); + var context = CreateContext(methodType); + var observedCancellation = false; + var middleware = new GrpcStreamLifetimeMiddleware(nextContext => + { + lifetime.StopApplication(); + observedCancellation = nextContext.RequestAborted.IsCancellationRequested; + return Task.CompletedTask; + }, lifetime); + + await middleware.InvokeAsync(context); + + Assert.That(observedCancellation, Is.True); + } + + [Test] + public async Task unary_calls_are_not_cancelled_by_host_shutdown() + { + using var lifetime = new TestHostLifetime(); + var context = CreateContext(MethodType.Unary); + var observedCancellation = true; + var middleware = new GrpcStreamLifetimeMiddleware(nextContext => + { + lifetime.StopApplication(); + observedCancellation = nextContext.RequestAborted.IsCancellationRequested; + return Task.CompletedTask; + }, lifetime); + + await middleware.InvokeAsync(context); + + Assert.That(observedCancellation, Is.False); + } + + [Test] + public async Task streaming_calls_keep_client_abort_linked() + { + using var lifetime = new TestHostLifetime(); + using var clientAbort = new CancellationTokenSource(); + var context = CreateContext(MethodType.ServerStreaming); + context.RequestAborted = clientAbort.Token; + var observedCancellation = false; + var middleware = new GrpcStreamLifetimeMiddleware(nextContext => + { + clientAbort.Cancel(); + observedCancellation = nextContext.RequestAborted.IsCancellationRequested; + return Task.CompletedTask; + }, lifetime); + + await middleware.InvokeAsync(context); + + Assert.That(observedCancellation, Is.True); + } + + [Test] + public async Task non_grpc_requests_are_not_linked_to_host_shutdown() + { + using var lifetime = new TestHostLifetime(); + var context = new DefaultHttpContext(); + var observedCancellation = true; + var middleware = new GrpcStreamLifetimeMiddleware(nextContext => + { + lifetime.StopApplication(); + observedCancellation = nextContext.RequestAborted.IsCancellationRequested; + return Task.CompletedTask; + }, lifetime); + + await middleware.InvokeAsync(context); + + Assert.That(observedCancellation, Is.False); + } + + [Test] + public async Task original_request_token_is_restored_after_stream_completes() + { + using var lifetime = new TestHostLifetime(); + using var clientAbort = new CancellationTokenSource(); + var context = CreateContext(MethodType.ServerStreaming); + context.RequestAborted = clientAbort.Token; + var middleware = new GrpcStreamLifetimeMiddleware(_ => Task.CompletedTask, lifetime); + + await middleware.InvokeAsync(context); + + Assert.That(context.RequestAborted, Is.EqualTo(clientAbort.Token)); + } + + private static DefaultHttpContext CreateContext(MethodType methodType) + { + var context = new DefaultHttpContext(); + context.SetEndpoint(new Endpoint( + _ => Task.CompletedTask, + new EndpointMetadataCollection(new GrpcMethodMetadata(typeof(object), new TestMethod(methodType))), + "test")); + return context; + } + + private sealed class TestMethod(MethodType type) : IMethod + { + public MethodType Type => type; + public string ServiceName => "service"; + public string Name => "method"; + public string FullName => "/service/method"; + } + + private sealed class TestHostLifetime : IHostApplicationLifetime, IDisposable + { + private readonly CancellationTokenSource _stopping = new(); + + public CancellationToken ApplicationStarted => CancellationToken.None; + public CancellationToken ApplicationStopping => _stopping.Token; + public CancellationToken ApplicationStopped => CancellationToken.None; + + public void StopApplication() => _stopping.Cancel(); + public void Dispose() => _stopping.Dispose(); + } +} diff --git a/src/EventStore.Core/ClusterVNodeStartup.cs b/src/EventStore.Core/ClusterVNodeStartup.cs index c750a3e23d..427c26f890 100644 --- a/src/EventStore.Core/ClusterVNodeStartup.cs +++ b/src/EventStore.Core/ClusterVNodeStartup.cs @@ -131,6 +131,7 @@ public void Configure(IApplicationBuilder app) // is driven by the HttpContext.User established above .UseAuthentication() .UseRouting() + .UseMiddleware() .UseAuthorization() .UseAntiforgery(); @@ -237,9 +238,11 @@ public void ConfigureServicesOnly(IServiceCollection services) // gRPC .AddSingleton() + .AddSingleton() .AddGrpc(options => { options.Interceptors.Add(); + options.Interceptors.Add(); var eventStoreConfiguration = _configuration.GetSection("EventStore"); var compressionLevel = (eventStoreConfiguration.Exists() diff --git a/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.AllSubscription.cs b/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.AllSubscription.cs index eaaf2f3af9..7827155086 100644 --- a/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.AllSubscription.cs +++ b/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.AllSubscription.cs @@ -27,6 +27,7 @@ public class AllSubscription : IAsyncEnumerator private readonly ClaimsPrincipal _user; private readonly bool _requiresLeader; private readonly CancellationTokenSource _cts; + private readonly CancellationToken _cancellationToken; private readonly Channel _channel; private readonly Channel<(ulong SequenceNumber, ResolvedEvent ResolvedEvent)> _liveEvents; @@ -54,6 +55,7 @@ public AllSubscription( _resolveLinks = resolveLinks; _user = user; _requiresLeader = requiresLeader; + _cancellationToken = cancellationToken; _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); _channel = Channel.CreateBounded(BoundedChannelOptions); _liveEvents = Channel.CreateBounded<(ulong, ResolvedEvent)>(LiveChannelOptions); @@ -83,6 +85,19 @@ public ValueTask DisposeAsync() public async ValueTask MoveNextAsync() { + try + { + return await MoveNextCoreAsync(); + } + catch (OperationCanceledException exception) when ( + exception.CancellationToken == _cts.Token && _cancellationToken.IsCancellationRequested) + { + throw new OperationCanceledException(exception.Message, exception, _cancellationToken); + } + } + + private async ValueTask MoveNextCoreAsync() + { ReadLoop: if (!await _channel.Reader.WaitToReadAsync(_cts.Token)) diff --git a/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.AllSubscriptionFiltered.cs b/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.AllSubscriptionFiltered.cs index a3df614fb1..5bb1b23bfe 100644 --- a/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.AllSubscriptionFiltered.cs +++ b/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.AllSubscriptionFiltered.cs @@ -30,6 +30,7 @@ public class AllSubscriptionFiltered : IAsyncEnumerator private readonly uint _maxSearchWindow; private readonly uint _checkpointInterval; private readonly CancellationTokenSource _cts; + private readonly CancellationToken _cancellationToken; private readonly Channel _channel; private readonly Channel<(ulong SequenceNumber, ResolvedEvent? ResolvedEvent, TFPos? Checkpoint)> _liveEvents; @@ -65,6 +66,7 @@ public AllSubscriptionFiltered(IPublisher bus, _requiresLeader = requiresLeader; _maxSearchWindow = maxSearchWindow ?? ReadBatchSize; _checkpointInterval = checkpointIntervalMultiplier * _maxSearchWindow; + _cancellationToken = cancellationToken; _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); _channel = Channel.CreateBounded(BoundedChannelOptions); _liveEvents = Channel.CreateBounded<(ulong, ResolvedEvent?, TFPos?)>(LiveChannelOptions); @@ -94,6 +96,19 @@ public ValueTask DisposeAsync() public async ValueTask MoveNextAsync() { + try + { + return await MoveNextCoreAsync(); + } + catch (OperationCanceledException exception) when ( + exception.CancellationToken == _cts.Token && _cancellationToken.IsCancellationRequested) + { + throw new OperationCanceledException(exception.Message, exception, _cancellationToken); + } + } + + private async ValueTask MoveNextCoreAsync() + { ReadLoop: if (!await _channel.Reader.WaitToReadAsync(_cts.Token)) diff --git a/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.StreamSubscription.cs b/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.StreamSubscription.cs index 9e90010176..3162da1338 100644 --- a/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.StreamSubscription.cs +++ b/src/EventStore.Core/Services/Transport/Enumerators/Enumerator.StreamSubscription.cs @@ -36,6 +36,7 @@ public class StreamSubscription : StreamSubscription private readonly ClaimsPrincipal _user; private readonly bool _requiresLeader; private readonly CancellationTokenSource _cts; + private readonly CancellationToken _cancellationToken; private readonly Channel _channel; private readonly Channel<(ulong SequenceNumber, ResolvedEvent ResolvedEvent)> _liveEvents; @@ -66,6 +67,7 @@ public StreamSubscription( _resolveLinks = resolveLinks; _user = user; _requiresLeader = requiresLeader; + _cancellationToken = cancellationToken; _cts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken); _channel = Channel.CreateBounded(BoundedChannelOptions); _liveEvents = Channel.CreateBounded<(ulong, ResolvedEvent)>(LiveChannelOptions); @@ -97,6 +99,19 @@ public override ValueTask DisposeAsync() public override async ValueTask MoveNextAsync() { + try + { + return await MoveNextCoreAsync(); + } + catch (OperationCanceledException exception) when ( + exception.CancellationToken == _cts.Token && _cancellationToken.IsCancellationRequested) + { + throw new OperationCanceledException(exception.Message, exception, _cancellationToken); + } + } + + private async ValueTask MoveNextCoreAsync() + { ReadLoop: if (!await _channel.Reader.WaitToReadAsync(_cts.Token)) diff --git a/src/EventStore.Core/Services/Transport/Grpc/Constants.cs b/src/EventStore.Core/Services/Transport/Grpc/Constants.cs index 61bc03d265..0905fecb3d 100644 --- a/src/EventStore.Core/Services/Transport/Grpc/Constants.cs +++ b/src/EventStore.Core/Services/Transport/Grpc/Constants.cs @@ -14,6 +14,7 @@ public static class Exceptions public const string MaximumAppendSizeExceeded = "maximum-append-size-exceeded"; public const string MissingRequiredMetadataProperty = "missing-required-metadata-property"; public const string NotLeader = "not-leader"; + public const string ServerShuttingDown = "server-shutting-down"; public const string PersistentSubscriptionFailed = "persistent-subscription-failed"; public const string PersistentSubscriptionDoesNotExist = "persistent-subscription-does-not-exist"; diff --git a/src/EventStore.Core/Services/Transport/Grpc/GrpcServerShutdownInterceptor.cs b/src/EventStore.Core/Services/Transport/Grpc/GrpcServerShutdownInterceptor.cs new file mode 100644 index 0000000000..5dd60e89ff --- /dev/null +++ b/src/EventStore.Core/Services/Transport/Grpc/GrpcServerShutdownInterceptor.cs @@ -0,0 +1,61 @@ +using System; +using System.Threading.Tasks; +using Grpc.Core; +using Grpc.Core.Interceptors; +using Microsoft.Extensions.Hosting; + +namespace EventStore.Core.Services.Transport.Grpc; + +public sealed class GrpcServerShutdownInterceptor(IHostApplicationLifetime applicationLifetime) : Interceptor +{ + public override Task ClientStreamingServerHandler( + IAsyncStreamReader requestStream, + ServerCallContext context, + ClientStreamingServerMethod continuation) => + TranslateShutdownAsync(context, () => continuation(requestStream, context)); + + public override Task ServerStreamingServerHandler( + TRequest request, + IServerStreamWriter responseStream, + ServerCallContext context, + ServerStreamingServerMethod continuation) => + TranslateShutdownAsync(context, () => continuation(request, responseStream, context)); + + public override Task DuplexStreamingServerHandler( + IAsyncStreamReader requestStream, + IServerStreamWriter responseStream, + ServerCallContext context, + DuplexStreamingServerMethod continuation) => + TranslateShutdownAsync(context, () => continuation(requestStream, responseStream, context)); + + private async Task TranslateShutdownAsync(ServerCallContext context, Func continuation) + { + try + { + await continuation(); + } + catch (OperationCanceledException exception) when (IsShutdownCancellation(exception, context)) + { + throw RpcExceptions.ServerShuttingDown(); + } + } + + private async Task TranslateShutdownAsync( + ServerCallContext context, + Func> continuation) + { + try + { + return await continuation(); + } + catch (OperationCanceledException exception) when (IsShutdownCancellation(exception, context)) + { + throw RpcExceptions.ServerShuttingDown(); + } + } + + private bool IsShutdownCancellation(OperationCanceledException exception, ServerCallContext context) => + applicationLifetime.ApplicationStopping.IsCancellationRequested && + context.CancellationToken.IsCancellationRequested && + exception.CancellationToken == context.CancellationToken; +} diff --git a/src/EventStore.Core/Services/Transport/Grpc/GrpcStreamLifetimeMiddleware.cs b/src/EventStore.Core/Services/Transport/Grpc/GrpcStreamLifetimeMiddleware.cs new file mode 100644 index 0000000000..3157760cd9 --- /dev/null +++ b/src/EventStore.Core/Services/Transport/Grpc/GrpcStreamLifetimeMiddleware.cs @@ -0,0 +1,39 @@ +using System.Threading; +using System.Threading.Tasks; +using Grpc.AspNetCore.Server; +using Grpc.Core; +using Microsoft.AspNetCore.Http; +using Microsoft.Extensions.Hosting; + +namespace EventStore.Core.Services.Transport.Grpc; + +public sealed class GrpcStreamLifetimeMiddleware( + RequestDelegate next, + IHostApplicationLifetime applicationLifetime) +{ + public Task InvokeAsync(HttpContext context) + { + var methodType = context.GetEndpoint()?.Metadata.GetMetadata()?.Method.Type; + return methodType is MethodType.ClientStreaming or MethodType.ServerStreaming or MethodType.DuplexStreaming + ? InvokeStreamingAsync(context) + : next(context); + } + + private async Task InvokeStreamingAsync(HttpContext context) + { + var requestAborted = context.RequestAborted; + using var streamLifetime = CancellationTokenSource.CreateLinkedTokenSource( + requestAborted, + applicationLifetime.ApplicationStopping); + context.RequestAborted = streamLifetime.Token; + + try + { + await next(context); + } + finally + { + context.RequestAborted = requestAborted; + } + } +} diff --git a/src/EventStore.Core/Services/Transport/Grpc/RpcExceptions.cs b/src/EventStore.Core/Services/Transport/Grpc/RpcExceptions.cs index 792adbdcdb..0130a2be28 100644 --- a/src/EventStore.Core/Services/Transport/Grpc/RpcExceptions.cs +++ b/src/EventStore.Core/Services/Transport/Grpc/RpcExceptions.cs @@ -18,6 +18,11 @@ public static RpcException ServerNotReady() => public static RpcException ServerBusy() => new(new Status(StatusCode.Unavailable, "Server Is Too Busy")); + public static RpcException ServerShuttingDown() => + new(new Status(StatusCode.Unavailable, "Server Is Shutting Down"), new Metadata { + {Constants.Exceptions.ExceptionKey, Constants.Exceptions.ServerShuttingDown} + }); + public static Exception NoLeaderInfo() => new RpcException(new Status(StatusCode.Unknown, "No leader info available in response")); diff --git a/src/EventStore.Core/Services/Transport/Grpc/Streams.Read.cs b/src/EventStore.Core/Services/Transport/Grpc/Streams.Read.cs index 4230d62ad6..703d66464e 100644 --- a/src/EventStore.Core/Services/Transport/Grpc/Streams.Read.cs +++ b/src/EventStore.Core/Services/Transport/Grpc/Streams.Read.cs @@ -76,18 +76,13 @@ public override async Task Read( filterOptionsCase, context.CancellationToken); - async void DisposeEnumerator() => await enumerator.DisposeAsync(); - await using (enumerator) { - await using (context.CancellationToken.Register(DisposeEnumerator)) + while (await enumerator.MoveNextAsync()) { - while (await enumerator.MoveNextAsync()) + if (TryConvertReadResponse(enumerator.Current, uuidOption, out var readResponse)) { - if (TryConvertReadResponse(enumerator.Current, uuidOption, out var readResponse)) - { - await responseStream.WriteAsync(readResponse); - } + await responseStream.WriteAsync(readResponse); } } }