fix(grpc): preserve graceful streaming shutdown

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<SubscriptionConfirmation>());
+			Assert.That(await sub.GetNext(), Is.InstanceOf<CaughtUp>());
+			cancellation.Cancel();
+
+			var exception = Assert.ThrowsAsync<OperationCanceledException>(() => sub.GetNext());
+			Assert.That(exception!.CancellationToken, Is.EqualTo(cancellation.Token));
+		}
 	}
 
 	[TestFixture(typeof(LogFormat.V2), typeof(string))]

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<SubscriptionConfirmation>());
+			Assert.That(await sub.GetNext(), Is.InstanceOf<CaughtUp>());
+			cancellation.Cancel();
+
+			var exception = Assert.ThrowsAsync<OperationCanceledException>(() => sub.GetNext());
+			Assert.That(exception!.CancellationToken, Is.EqualTo(cancellation.Token));
+		}
 	}
 
 	[TestFixture(typeof(LogFormat.V2), typeof(string))]

src/EventStore.Core.Tests/Services/Transport/Enumerators/Enumerator.StreamSubscription.Tests.cs

@@ -20,7 +20,8 @@ private static EnumeratorWrapper CreateStreamSubscription<TStreamId>(
 		IPublisher publisher,
 		string streamName,
 		StreamRevision? checkpoint = null,
-		ClaimsPrincipal user = null)
+		ClaimsPrincipal user = null,
+		CancellationToken cancellationToken = default)
 	{
 
 		return new EnumeratorWrapper(new Enumerator.StreamSubscription<TStreamId>(
@@ -31,7 +32,7 @@ private static EnumeratorWrapper CreateStreamSubscription<TStreamId>(
 			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<TStreamId>(
+				_publisher, streamName: "test-stream1", StreamRevision.End, cancellationToken: cancellation.Token);
+
+			Assert.That(await enumerator.GetNext(), Is.InstanceOf<SubscriptionConfirmation>());
+			Assert.That(await enumerator.GetNext(), Is.InstanceOf<CaughtUp>());
+			cancellation.Cancel();
+
+			var exception = Assert.ThrowsAsync<OperationCanceledException>(() => enumerator.GetNext());
+			Assert.That(exception!.CancellationToken, Is.EqualTo(cancellation.Token));
+		}
 	}
 }

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<RpcException>(() => 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<OperationCanceledException>(() =>
+			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<string>(),
+			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<string>(),
+					context,
+					(_, callContext) => Task.FromCanceled<string>(callContext.CancellationToken));
+				break;
+			case MethodType.ServerStreaming:
+				await interceptor.ServerStreamingServerHandler(
+					"request",
+					new TestServerStreamWriter<string>(),
+					context,
+					(_, _, callContext) => Task.FromCanceled(callContext.CancellationToken));
+				break;
+			case MethodType.DuplexStreaming:
+				await interceptor.DuplexStreamingServerHandler(
+					new TestAsyncStreamReader<string>(),
+					new TestServerStreamWriter<string>(),
+					context,
+					(_, _, callContext) => Task.FromCanceled(callContext.CancellationToken));
+				break;
+			default:
+				throw new ArgumentOutOfRangeException(nameof(methodType), methodType, null);
+		}
+	}
+
+	private sealed class TestAsyncStreamReader<T> : IAsyncStreamReader<T>
+	{
+		public T Current => default!;
+		public Task<bool> MoveNext(CancellationToken cancellationToken) => Task.FromResult(false);
+	}
+
+	private sealed class TestServerStreamWriter<T> : IServerStreamWriter<T>
+	{
+		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<string, List<AuthProperty>>());
+		protected override IDictionary<object, object> UserStateCore { get; } =
+			new Dictionary<object, object>();
+		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();
+	}
+}