Skip to content

Commit d24c9b2

Browse files
svametcalfsvametcalf
committed
Fixes Reactor context propagation, including tests
1 parent 13df614 commit d24c9b2

File tree

4 files changed

+234
-29
lines changed

4 files changed

+234
-29
lines changed

‎reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ClientCalls.java‎

Lines changed: 10 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@
88
package com.salesforce.reactorgrpc.stub;
99

1010
import io.grpc.CallOptions;
11-
import io.grpc.stub.CallStreamObserver;
1211
import io.grpc.stub.StreamObserver;
1312
import java.util.function.BiConsumer;
1413
import java.util.function.Function;
@@ -34,9 +33,8 @@ public static <TRequest, TResponse> Mono<TResponse> oneToOne(
3433
BiConsumer<TRequest, StreamObserver<TResponse>> delegate,
3534
CallOptions options) {
3635
try {
37-
return Mono
38-
.<TResponse>create(emitter -> monoSource.subscribe(
39-
request -> delegate.accept(request, new StreamObserver<TResponse>() {
36+
return monoSource.flatMap(r ->
37+
Mono.<TResponse>create(emitter -> delegate.accept(r, new StreamObserver<TResponse>() {
4038
@Override
4139
public void onNext(TResponse tResponse) {
4240
emitter.success(tResponse);
@@ -51,10 +49,10 @@ public void onError(Throwable throwable) {
5149
public void onCompleted() {
5250
// Do nothing
5351
}
54-
}),
55-
emitter::error
56-
))
57-
.transform(Operators.lift(new SubscribeOnlyOnceLifter<TResponse>()));
52+
})
53+
)
54+
.transform(Operators.lift(new SubscribeOnlyOnceLifter<TResponse>()))
55+
);
5856
} catch (Throwable throwable) {
5957
return Mono.error(throwable);
6058
}
@@ -97,17 +95,8 @@ public static <TRequest, TResponse> Mono<TResponse> manyToOne(
9795
Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate,
9896
CallOptions options) {
9997
try {
100-
ReactorSubscriberAndClientProducer<TRequest> subscriberAndGRPCProducer =
101-
fluxSource.subscribeWith(new ReactorSubscriberAndClientProducer<>());
102-
ReactorClientStreamObserverAndPublisher<TResponse> observerAndPublisher =
103-
new ReactorClientStreamObserverAndPublisher<>(
104-
s -> subscriberAndGRPCProducer.subscribe((CallStreamObserver<TRequest>) s),
105-
subscriberAndGRPCProducer::cancel
106-
);
107-
108-
return Flux.from(observerAndPublisher)
109-
.doOnSubscribe(s -> delegate.apply(observerAndPublisher))
110-
.singleOrEmpty();
98+
ReactorGrpcClientCallFlux<TRequest, TResponse> operator = new ReactorGrpcClientCallFlux<>(fluxSource, delegate);
99+
return operator.doOnSubscribe(operator.onSubscribeHook()).singleOrEmpty();
111100
} catch (Throwable throwable) {
112101
return Mono.error(throwable);
113102
}
@@ -123,19 +112,11 @@ public static <TRequest, TResponse> Flux<TResponse> manyToMany(
123112
Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate,
124113
CallOptions options) {
125114
try {
126-
127115
final int prefetch = ReactorCallOptions.getPrefetch(options);
128116
final int lowTide = ReactorCallOptions.getLowTide(options);
129117

130-
ReactorSubscriberAndClientProducer<TRequest> subscriberAndGRPCProducer =
131-
fluxSource.subscribeWith(new ReactorSubscriberAndClientProducer<>());
132-
ReactorClientStreamObserverAndPublisher<TResponse> observerAndPublisher =
133-
new ReactorClientStreamObserverAndPublisher<>(
134-
s -> subscriberAndGRPCProducer.subscribe((CallStreamObserver<TRequest>) s),
135-
subscriberAndGRPCProducer::cancel, prefetch, lowTide
136-
);
137-
138-
return Flux.from(observerAndPublisher).doOnSubscribe(s -> delegate.apply(observerAndPublisher));
118+
ReactorGrpcClientCallFlux<TRequest, TResponse> operator = new ReactorGrpcClientCallFlux<>(fluxSource, delegate, prefetch, lowTide);
119+
return operator.doOnSubscribe(operator.onSubscribeHook());
139120
} catch (Throwable throwable) {
140121
return Flux.error(throwable);
141122
}

‎reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ReactorGrpcClientCallFlux.java‎

Lines changed: 80 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,80 @@
1+
/*
2+
* Copyright (c) 2019, Salesforce.com, Inc.
3+
* All rights reserved.
4+
* Licensed under the BSD 3-Clause license.
5+
* For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6+
*/
7+
package com.salesforce.reactorgrpc.stub;
8+
9+
import io.grpc.stub.CallStreamObserver;
10+
import io.grpc.stub.StreamObserver;
11+
import org.reactivestreams.Subscription;
12+
import reactor.core.CoreSubscriber;
13+
import reactor.core.publisher.Flux;
14+
import reactor.core.publisher.FluxOperator;
15+
import reactor.util.context.Context;
16+
17+
import java.util.function.Consumer;
18+
import java.util.function.Function;
19+
20+
/**
21+
* Create a {@link Flux} that allows for correct context propagation in client calls
22+
*
23+
* @param <TRequest>
24+
* @param <TResponse>
25+
*/
26+
public class ReactorGrpcClientCallFlux<TRequest, TResponse> extends FluxOperator<TRequest, TResponse> {
27+
28+
private final ReactorSubscriberAndClientProducer<TRequest> requestConsumer;
29+
private final ReactorClientStreamObserverAndPublisher<TResponse> responsePublisher;
30+
private final Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate;
31+
32+
ReactorGrpcClientCallFlux(Flux<TRequest> in, Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate) {
33+
super(in);
34+
this.delegate = delegate;
35+
this.requestConsumer = new ReactorSubscriberAndClientProducer<>();
36+
this.responsePublisher = new ReactorClientStreamObserverAndPublisher<>(s -> requestConsumer.subscribe((CallStreamObserver<TRequest>) s), requestConsumer::cancel);
37+
}
38+
39+
public ReactorGrpcClientCallFlux(Flux<TRequest> in, Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate, int prefetch, int lowTide) {
40+
super(in);
41+
this.delegate = delegate;
42+
this.requestConsumer = new ReactorSubscriberAndClientProducer<>();
43+
this.responsePublisher = new ReactorClientStreamObserverAndPublisher<>(s -> requestConsumer.subscribe((CallStreamObserver<TRequest>) s), requestConsumer::cancel, prefetch, lowTide);
44+
}
45+
46+
public Consumer<? super Subscription> onSubscribeHook() {
47+
return s -> this.delegate.apply(responsePublisher);
48+
}
49+
50+
@Override
51+
public void subscribe(CoreSubscriber<? super TResponse> actual) {
52+
responsePublisher.subscribe(actual);
53+
source.subscribe(new CoreSubscriber<TRequest>() {
54+
@Override
55+
public void onSubscribe(Subscription s) {
56+
requestConsumer.onSubscribe(s);
57+
}
58+
59+
@Override
60+
public void onNext(TRequest tRequest) {
61+
requestConsumer.onNext(tRequest);
62+
}
63+
64+
@Override
65+
public void onError(Throwable throwable) {
66+
requestConsumer.onError(throwable);
67+
}
68+
69+
@Override
70+
public void onComplete() {
71+
requestConsumer.onComplete();
72+
}
73+
74+
@Override
75+
public Context currentContext() {
76+
return actual.currentContext();
77+
}
78+
});
79+
}
80+
}

‎reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/SubscribeOnlyOnceLifter.java‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import org.reactivestreams.Subscription;
1111
import reactor.core.CoreSubscriber;
1212
import reactor.core.Scannable;
13+
import reactor.util.context.Context;
1314

1415
import java.util.concurrent.atomic.AtomicBoolean;
1516
import java.util.function.BiFunction;
@@ -25,6 +26,11 @@ public class SubscribeOnlyOnceLifter<T> extends AtomicBoolean implements BiFunct
2526
@Override
2627
public CoreSubscriber<? super T> apply(Scannable scannable, CoreSubscriber<? super T> coreSubscriber) {
2728
return new CoreSubscriber<T>() {
29+
@Override
30+
public Context currentContext() {
31+
return coreSubscriber.currentContext();
32+
}
33+
2834
@Override
2935
public void onSubscribe(Subscription subscription) {
3036
if (!compareAndSet(false, true)) {

‎reactor/reactor-grpc-test/src/test/java/com/salesforce/reactorgrpc/ReactorContextPropagationTest.java‎

Lines changed: 138 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,138 @@
1+
/*
2+
* Copyright (c) 2019, Salesforce.com, Inc.
3+
* All rights reserved.
4+
* Licensed under the BSD 3-Clause license.
5+
* For full license text, see LICENSE.txt file in the repo root or https://opensource.org/licenses/BSD-3-Clause
6+
*/
7+
8+
package com.salesforce.reactorgrpc;
9+
10+
import io.grpc.testing.GrpcServerRule;
11+
import org.junit.*;
12+
import reactor.core.publisher.Flux;
13+
import reactor.core.publisher.Hooks;
14+
import reactor.core.publisher.Mono;
15+
import reactor.test.StepVerifier;
16+
17+
import java.util.stream.Collectors;
18+
19+
import static org.assertj.core.api.Assertions.assertThat;
20+
21+
public class ReactorContextPropagationTest {
22+
23+
@Rule
24+
public GrpcServerRule serverRule = new GrpcServerRule();
25+
26+
private static class SimpleGreeter extends ReactorGreeterGrpc.GreeterImplBase {
27+
@Override
28+
public Mono<HelloResponse> sayHello(Mono<HelloRequest> request) {
29+
return request.map(HelloRequest::getName)
30+
.map(name -> HelloResponse.newBuilder().setMessage("Hello " + name).build());
31+
}
32+
33+
@Override
34+
public Mono<HelloResponse> sayHelloReqStream(Flux<HelloRequest> request) {
35+
return request.transformDeferredContextual((f, ctx) -> f.map(HelloRequest::getName))
36+
.collect(Collectors.joining("and"))
37+
.map(names -> HelloResponse.newBuilder().setMessage("Hello " + names).build());
38+
}
39+
40+
@Override
41+
public Flux<HelloResponse> sayHelloRespStream(Mono<HelloRequest> request) {
42+
return request.repeat(2)
43+
.map(HelloRequest::getName)
44+
.zipWith(Flux.just("Hello ", "Hi ", "Greetings "), String::join)
45+
.map(greeting -> HelloResponse.newBuilder().setMessage(greeting).build());
46+
}
47+
48+
@Override
49+
public Flux<HelloResponse> sayHelloBothStream(Flux<HelloRequest> request) {
50+
return request.map(HelloRequest::getName)
51+
.map(name -> HelloResponse.newBuilder().setMessage("Hello " + name).build());
52+
}
53+
}
54+
55+
@BeforeClass
56+
public static void beforeAll(){
57+
Hooks.enableContextLossTracking();
58+
Hooks.onOperatorDebug();
59+
}
60+
61+
@AfterClass
62+
public static void afterAll(){
63+
Hooks.disableContextLossTracking();
64+
Hooks.resetOnOperatorDebug();
65+
}
66+
67+
@Before
68+
public void setup() {
69+
serverRule.getServiceRegistry().addService(new SimpleGreeter());
70+
}
71+
72+
@Test
73+
public void oneToOne() {
74+
ReactorGreeterGrpc.ReactorGreeterStub stub = ReactorGreeterGrpc.newReactorStub(serverRule.getChannel());
75+
Mono<HelloRequest> req = Mono.just(HelloRequest.newBuilder().setName("reactor").build());
76+
77+
Mono<HelloResponse> resp = req
78+
.doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
79+
.transform(stub::sayHello)
80+
.doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
81+
.contextWrite(ctx -> ctx.put("name", "context"));
82+
83+
StepVerifier.create(resp)
84+
.expectNextCount(1)
85+
.verifyComplete();
86+
}
87+
88+
@Test
89+
public void oneToMany() {
90+
ReactorGreeterGrpc.ReactorGreeterStub stub = ReactorGreeterGrpc.newReactorStub(serverRule.getChannel());
91+
Mono<HelloRequest> req = Mono.just(HelloRequest.newBuilder().setName("reactor").build());
92+
93+
Flux<HelloResponse> resp = req
94+
.doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
95+
.as(stub::sayHelloRespStream)
96+
.doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
97+
.contextWrite(ctx -> ctx.put("name", "context"));
98+
99+
StepVerifier.create(resp)
100+
.expectNextCount(3)
101+
.verifyComplete();
102+
}
103+
104+
@Test
105+
public void manyToOne() {
106+
ReactorGreeterGrpc.ReactorGreeterStub stub = ReactorGreeterGrpc.newReactorStub(serverRule.getChannel());
107+
Flux<HelloRequest> req = Mono.deferContextual(ctx -> Mono.just(HelloRequest.newBuilder().setName(ctx.get("name")).build())).repeat(2);
108+
109+
Mono<HelloResponse> resp = req
110+
.doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
111+
.as(stub::sayHelloReqStream)
112+
.doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
113+
.contextWrite(ctx -> ctx.put("name", "context"));
114+
115+
StepVerifier.create(resp)
116+
.expectAccessibleContext()
117+
.contains("name", "context")
118+
.then()
119+
.expectNextCount(1)
120+
.verifyComplete();
121+
}
122+
123+
@Test
124+
public void manyToMany() {
125+
ReactorGreeterGrpc.ReactorGreeterStub stub = ReactorGreeterGrpc.newReactorStub(serverRule.getChannel());
126+
Flux<HelloRequest> req = Mono.just(HelloRequest.newBuilder().setName("reactor").build()).repeat(2).contextWrite(c -> c.put("name", "boom"));
127+
128+
Flux<HelloResponse> resp = req
129+
.doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
130+
.transform(stub::sayHelloBothStream)
131+
.doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
132+
.contextWrite(ctx -> ctx.put("name", "context"));
133+
134+
StepVerifier.create(resp)
135+
.expectNextCount(3)
136+
.verifyComplete();
137+
}
138+
}

0 commit comments

Comments
 (0)