Skip to content

Commit 94ba180

Browse files
committed
fix(grpc): preserve graceful streaming shutdown
Signed-off-by: Yordis Prieto <yordis.prieto@gmail.com>
1 parent c002bc9 commit 94ba180

14 files changed

Lines changed: 496 additions & 14 deletions

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

Lines changed: 17 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ public partial class EnumeratorTests
2020
private static EnumeratorWrapper CreateAllSubscription(
2121
IPublisher publisher,
2222
Position? checkpoint,
23-
ClaimsPrincipal user = null)
23+
ClaimsPrincipal user = null,
24+
CancellationToken cancellationToken = default)
2425
{
2526

2627
return new EnumeratorWrapper(new Enumerator.AllSubscription(
@@ -30,7 +31,7 @@ private static EnumeratorWrapper CreateAllSubscription(
3031
resolveLinks: false,
3132
user: user ?? SystemAccounts.System,
3233
requiresLeader: false,
33-
cancellationToken: CancellationToken.None));
34+
cancellationToken: cancellationToken));
3435
}
3536

3637
[TestFixture(typeof(LogFormat.V2), typeof(string))]
@@ -83,6 +84,20 @@ public async Task should_receive_live_caught_up_message_immediately()
8384
Assert.True(await sub.GetNext() is SubscriptionConfirmation);
8485
Assert.True(await sub.GetNext() is CaughtUp);
8586
}
87+
88+
[Test]
89+
public async Task cancellation_preserves_the_callers_token()
90+
{
91+
using var cancellation = new CancellationTokenSource();
92+
await using var sub = CreateAllSubscription(_publisher, Position.End, cancellationToken: cancellation.Token);
93+
94+
Assert.That(await sub.GetNext(), Is.InstanceOf<SubscriptionConfirmation>());
95+
Assert.That(await sub.GetNext(), Is.InstanceOf<CaughtUp>());
96+
cancellation.Cancel();
97+
98+
var exception = Assert.ThrowsAsync<OperationCanceledException>(() => sub.GetNext());
99+
Assert.That(exception!.CancellationToken, Is.EqualTo(cancellation.Token));
100+
}
86101
}
87102

88103
[TestFixture(typeof(LogFormat.V2), typeof(string))]

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

Lines changed: 21 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,8 @@ private static EnumeratorWrapper CreateAllSubscriptionFiltered(
2323
IEventFilter eventFilter = null,
2424
uint? maxSearchWindow = null,
2525
uint checkpointIntervalMultiplier = 1,
26-
ClaimsPrincipal user = null)
26+
ClaimsPrincipal user = null,
27+
CancellationToken cancellationToken = default)
2728
{
2829

2930
return new EnumeratorWrapper(new Enumerator.AllSubscriptionFiltered(
@@ -36,7 +37,7 @@ private static EnumeratorWrapper CreateAllSubscriptionFiltered(
3637
requiresLeader: false,
3738
maxSearchWindow: maxSearchWindow,
3839
checkpointIntervalMultiplier: checkpointIntervalMultiplier,
39-
cancellationToken: CancellationToken.None));
40+
cancellationToken: cancellationToken));
4041
}
4142

4243

@@ -101,6 +102,24 @@ public async Task should_receive_live_caught_up_message_immediately()
101102
Assert.True(await sub.GetNext() is SubscriptionConfirmation);
102103
Assert.True(await sub.GetNext() is CaughtUp);
103104
}
105+
106+
[Test]
107+
public async Task cancellation_preserves_the_callers_token()
108+
{
109+
using var cancellation = new CancellationTokenSource();
110+
await using var sub = CreateAllSubscriptionFiltered(
111+
_publisher,
112+
Position.End,
113+
EventFilter.EventType.Prefixes(false, "type1"),
114+
cancellationToken: cancellation.Token);
115+
116+
Assert.That(await sub.GetNext(), Is.InstanceOf<SubscriptionConfirmation>());
117+
Assert.That(await sub.GetNext(), Is.InstanceOf<CaughtUp>());
118+
cancellation.Cancel();
119+
120+
var exception = Assert.ThrowsAsync<OperationCanceledException>(() => sub.GetNext());
121+
Assert.That(exception!.CancellationToken, Is.EqualTo(cancellation.Token));
122+
}
104123
}
105124

106125
[TestFixture(typeof(LogFormat.V2), typeof(string))]

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

Lines changed: 18 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,7 +20,8 @@ private static EnumeratorWrapper CreateStreamSubscription<TStreamId>(
2020
IPublisher publisher,
2121
string streamName,
2222
StreamRevision? checkpoint = null,
23-
ClaimsPrincipal user = null)
23+
ClaimsPrincipal user = null,
24+
CancellationToken cancellationToken = default)
2425
{
2526

2627
return new EnumeratorWrapper(new Enumerator.StreamSubscription<TStreamId>(
@@ -31,7 +32,7 @@ private static EnumeratorWrapper CreateStreamSubscription<TStreamId>(
3132
resolveLinks: false,
3233
user: user ?? SystemAccounts.System,
3334
requiresLeader: false,
34-
cancellationToken: CancellationToken.None));
35+
cancellationToken: cancellationToken));
3536
}
3637

3738
[TestFixture(typeof(LogFormat.V2), typeof(string))]
@@ -84,5 +85,20 @@ public async Task should_receive_live_caught_up_message_immediately()
8485
Assert.True(await enumerator.GetNext() is SubscriptionConfirmation);
8586
Assert.True(await enumerator.GetNext() is CaughtUp);
8687
}
88+
89+
[Test]
90+
public async Task cancellation_preserves_the_callers_token()
91+
{
92+
using var cancellation = new CancellationTokenSource();
93+
await using var enumerator = CreateStreamSubscription<TStreamId>(
94+
_publisher, streamName: "test-stream1", StreamRevision.End, cancellationToken: cancellation.Token);
95+
96+
Assert.That(await enumerator.GetNext(), Is.InstanceOf<SubscriptionConfirmation>());
97+
Assert.That(await enumerator.GetNext(), Is.InstanceOf<CaughtUp>());
98+
cancellation.Cancel();
99+
100+
var exception = Assert.ThrowsAsync<OperationCanceledException>(() => enumerator.GetNext());
101+
Assert.That(exception!.CancellationToken, Is.EqualTo(cancellation.Token));
102+
}
87103
}
88104
}
Lines changed: 148 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,148 @@
1+
using System;
2+
using System.Collections.Generic;
3+
using System.Linq;
4+
using System.Threading;
5+
using System.Threading.Tasks;
6+
using EventStore.Core.Services.Transport.Grpc;
7+
using Grpc.Core;
8+
using Microsoft.Extensions.Hosting;
9+
using NUnit.Framework;
10+
11+
namespace EventStore.Core.Tests.Services.Transport.Grpc;
12+
13+
[TestFixture]
14+
public class GrpcServerShutdownInterceptorTests
15+
{
16+
[TestCase(MethodType.ClientStreaming)]
17+
[TestCase(MethodType.ServerStreaming)]
18+
[TestCase(MethodType.DuplexStreaming)]
19+
public void shutdown_cancellation_is_retryable_for_every_streaming_shape(MethodType methodType)
20+
{
21+
using var lifetime = new TestHostLifetime();
22+
using var callCancellation = CancellationTokenSource.CreateLinkedTokenSource(lifetime.ApplicationStopping);
23+
lifetime.StopApplication();
24+
var context = new TestServerCallContext(callCancellation.Token);
25+
var interceptor = new GrpcServerShutdownInterceptor(lifetime);
26+
27+
var exception = Assert.ThrowsAsync<RpcException>(() => Invoke(interceptor, methodType, context));
28+
29+
Assert.That(exception!.StatusCode, Is.EqualTo(StatusCode.Unavailable));
30+
Assert.That(
31+
exception.Trailers.Select(entry => (entry.Key, entry.Value)),
32+
Does.Contain((
33+
EventStore.Core.Services.Transport.Grpc.Constants.Exceptions.ExceptionKey,
34+
EventStore.Core.Services.Transport.Grpc.Constants.Exceptions.ServerShuttingDown)));
35+
}
36+
37+
[Test]
38+
public void client_cancellation_is_not_changed_when_server_is_running()
39+
{
40+
using var lifetime = new TestHostLifetime();
41+
using var callCancellation = new CancellationTokenSource();
42+
callCancellation.Cancel();
43+
var context = new TestServerCallContext(callCancellation.Token);
44+
var interceptor = new GrpcServerShutdownInterceptor(lifetime);
45+
46+
var exception = Assert.CatchAsync<OperationCanceledException>(() =>
47+
Invoke(interceptor, MethodType.ServerStreaming, context));
48+
49+
Assert.That(exception!.CancellationToken, Is.EqualTo(callCancellation.Token));
50+
}
51+
52+
[Test]
53+
public async Task successful_streaming_call_passes_through()
54+
{
55+
using var lifetime = new TestHostLifetime();
56+
var context = new TestServerCallContext(CancellationToken.None);
57+
var interceptor = new GrpcServerShutdownInterceptor(lifetime);
58+
var continuationCalled = false;
59+
60+
await interceptor.ServerStreamingServerHandler(
61+
"request",
62+
new TestServerStreamWriter<string>(),
63+
context,
64+
(_, _, _) =>
65+
{
66+
continuationCalled = true;
67+
return Task.CompletedTask;
68+
});
69+
70+
Assert.That(continuationCalled, Is.True);
71+
}
72+
73+
private static async Task Invoke(
74+
GrpcServerShutdownInterceptor interceptor,
75+
MethodType methodType,
76+
ServerCallContext context)
77+
{
78+
switch (methodType)
79+
{
80+
case MethodType.ClientStreaming:
81+
await interceptor.ClientStreamingServerHandler(
82+
new TestAsyncStreamReader<string>(),
83+
context,
84+
(_, callContext) => Task.FromCanceled<string>(callContext.CancellationToken));
85+
break;
86+
case MethodType.ServerStreaming:
87+
await interceptor.ServerStreamingServerHandler(
88+
"request",
89+
new TestServerStreamWriter<string>(),
90+
context,
91+
(_, _, callContext) => Task.FromCanceled(callContext.CancellationToken));
92+
break;
93+
case MethodType.DuplexStreaming:
94+
await interceptor.DuplexStreamingServerHandler(
95+
new TestAsyncStreamReader<string>(),
96+
new TestServerStreamWriter<string>(),
97+
context,
98+
(_, _, callContext) => Task.FromCanceled(callContext.CancellationToken));
99+
break;
100+
default:
101+
throw new ArgumentOutOfRangeException(nameof(methodType), methodType, null);
102+
}
103+
}
104+
105+
private sealed class TestAsyncStreamReader<T> : IAsyncStreamReader<T>
106+
{
107+
public T Current => default!;
108+
public Task<bool> MoveNext(CancellationToken cancellationToken) => Task.FromResult(false);
109+
}
110+
111+
private sealed class TestServerStreamWriter<T> : IServerStreamWriter<T>
112+
{
113+
public WriteOptions WriteOptions { get; set; }
114+
public Task WriteAsync(T message) => Task.CompletedTask;
115+
}
116+
117+
private sealed class TestServerCallContext(CancellationToken cancellationToken) : ServerCallContext
118+
{
119+
protected override string MethodCore => "/service/method";
120+
protected override string HostCore => "host";
121+
protected override string PeerCore => "peer";
122+
protected override DateTime DeadlineCore => DateTime.MaxValue;
123+
protected override Metadata RequestHeadersCore { get; } = new();
124+
protected override CancellationToken CancellationTokenCore => cancellationToken;
125+
protected override Metadata ResponseTrailersCore { get; } = new();
126+
protected override Status StatusCore { get; set; }
127+
protected override WriteOptions WriteOptionsCore { get; set; }
128+
protected override AuthContext AuthContextCore { get; } =
129+
new(string.Empty, new Dictionary<string, List<AuthProperty>>());
130+
protected override IDictionary<object, object> UserStateCore { get; } =
131+
new Dictionary<object, object>();
132+
protected override Task WriteResponseHeadersAsyncCore(Metadata responseHeaders) => Task.CompletedTask;
133+
protected override ContextPropagationToken CreatePropagationTokenCore(ContextPropagationOptions options) =>
134+
throw new NotSupportedException();
135+
}
136+
137+
private sealed class TestHostLifetime : IHostApplicationLifetime, IDisposable
138+
{
139+
private readonly CancellationTokenSource _stopping = new();
140+
141+
public CancellationToken ApplicationStarted => CancellationToken.None;
142+
public CancellationToken ApplicationStopping => _stopping.Token;
143+
public CancellationToken ApplicationStopped => CancellationToken.None;
144+
145+
public void StopApplication() => _stopping.Cancel();
146+
public void Dispose() => _stopping.Dispose();
147+
}
148+
}

0 commit comments

Comments
 (0)