Fixes Reactor context propagation, including tests

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

@@ -8,7 +8,6 @@
 package com.salesforce.reactorgrpc.stub;
 
 import io.grpc.CallOptions;
-import io.grpc.stub.CallStreamObserver;
 import io.grpc.stub.StreamObserver;
 import java.util.function.BiConsumer;
 import java.util.function.Function;
@@ -34,9 +33,8 @@ public static <TRequest, TResponse> Mono<TResponse> oneToOne(
             BiConsumer<TRequest, StreamObserver<TResponse>> delegate,
             CallOptions options) {
         try {
-            return Mono
-                    .<TResponse>create(emitter -> monoSource.subscribe(
-                        request -> delegate.accept(request, new StreamObserver<TResponse>() {
+            return monoSource.flatMap(r ->
+                    Mono.<TResponse>create(emitter -> delegate.accept(r, new StreamObserver<TResponse>() {
                             @Override
                             public void onNext(TResponse tResponse) {
                                 emitter.success(tResponse);
@@ -51,10 +49,10 @@ public void onError(Throwable throwable) {
                             public void onCompleted() {
                                 // Do nothing
                             }
-                        }),
-                        emitter::error
-                    ))
-                    .transform(Operators.lift(new SubscribeOnlyOnceLifter<TResponse>()));
+                        })
+                    )
+                    .transform(Operators.lift(new SubscribeOnlyOnceLifter<TResponse>()))
+            );
         } catch (Throwable throwable) {
             return Mono.error(throwable);
         }
@@ -97,17 +95,8 @@ public static <TRequest, TResponse> Mono<TResponse> manyToOne(
             Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate,
             CallOptions options) {
         try {
-            ReactorSubscriberAndClientProducer<TRequest> subscriberAndGRPCProducer =
-                    fluxSource.subscribeWith(new ReactorSubscriberAndClientProducer<>());
-            ReactorClientStreamObserverAndPublisher<TResponse> observerAndPublisher =
-                    new ReactorClientStreamObserverAndPublisher<>(
-                        s -> subscriberAndGRPCProducer.subscribe((CallStreamObserver<TRequest>) s),
-                        subscriberAndGRPCProducer::cancel
-                    );
-
-            return Flux.from(observerAndPublisher)
-                    .doOnSubscribe(s -> delegate.apply(observerAndPublisher))
-                    .singleOrEmpty();
+            ReactorGrpcClientCallFlux<TRequest, TResponse> operator = new ReactorGrpcClientCallFlux<>(fluxSource, delegate);
+            return operator.doOnSubscribe(operator.onSubscribeHook()).singleOrEmpty();
         } catch (Throwable throwable) {
             return Mono.error(throwable);
         }
@@ -123,19 +112,11 @@ public static <TRequest, TResponse> Flux<TResponse> manyToMany(
             Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate,
             CallOptions options) {
         try {
-
             final int prefetch = ReactorCallOptions.getPrefetch(options);
             final int lowTide = ReactorCallOptions.getLowTide(options);
 
-            ReactorSubscriberAndClientProducer<TRequest> subscriberAndGRPCProducer =
-                fluxSource.subscribeWith(new ReactorSubscriberAndClientProducer<>());
-            ReactorClientStreamObserverAndPublisher<TResponse> observerAndPublisher =
-                new ReactorClientStreamObserverAndPublisher<>(
-                    s -> subscriberAndGRPCProducer.subscribe((CallStreamObserver<TRequest>) s),
-                    subscriberAndGRPCProducer::cancel, prefetch, lowTide
-                );
-
-            return Flux.from(observerAndPublisher).doOnSubscribe(s -> delegate.apply(observerAndPublisher));
+            ReactorGrpcClientCallFlux<TRequest, TResponse> operator = new ReactorGrpcClientCallFlux<>(fluxSource, delegate, prefetch, lowTide);
+            return operator.doOnSubscribe(operator.onSubscribeHook());
         } catch (Throwable throwable) {
             return Flux.error(throwable);
         }

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

@@ -0,0 +1,80 @@
+/*
+ *  Copyright (c) 2019, Salesforce.com, Inc.
+ *  All rights reserved.
+ *  Licensed under the BSD 3-Clause license.
+ *  For full license text, see LICENSE.txt file in the repo root  or https://opensource.org/licenses/BSD-3-Clause
+ */
+package com.salesforce.reactorgrpc.stub;
+
+import io.grpc.stub.CallStreamObserver;
+import io.grpc.stub.StreamObserver;
+import org.reactivestreams.Subscription;
+import reactor.core.CoreSubscriber;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.FluxOperator;
+import reactor.util.context.Context;
+
+import java.util.function.Consumer;
+import java.util.function.Function;
+
+/**
+ * Create a {@link Flux} that allows for correct context propagation in client calls.
+ *
+ * @param <TRequest>
+ * @param <TResponse>
+ */
+public class ReactorGrpcClientCallFlux<TRequest, TResponse> extends FluxOperator<TRequest, TResponse> {
+
+    private final ReactorSubscriberAndClientProducer<TRequest> requestConsumer;
+    private final ReactorClientStreamObserverAndPublisher<TResponse> responsePublisher;
+    private final Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate;
+
+    ReactorGrpcClientCallFlux(Flux<TRequest> in, Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate) {
+        super(in);
+        this.delegate = delegate;
+        this.requestConsumer = new ReactorSubscriberAndClientProducer<>();
+        this.responsePublisher = new ReactorClientStreamObserverAndPublisher<>(s -> requestConsumer.subscribe((CallStreamObserver<TRequest>) s), requestConsumer::cancel);
+    }
+
+    public ReactorGrpcClientCallFlux(Flux<TRequest> in, Function<StreamObserver<TResponse>, StreamObserver<TRequest>> delegate, int prefetch, int lowTide) {
+        super(in);
+        this.delegate = delegate;
+        this.requestConsumer = new ReactorSubscriberAndClientProducer<>();
+        this.responsePublisher = new ReactorClientStreamObserverAndPublisher<>(s -> requestConsumer.subscribe((CallStreamObserver<TRequest>) s), requestConsumer::cancel, prefetch, lowTide);
+    }
+
+    public Consumer<? super Subscription> onSubscribeHook() {
+        return s -> this.delegate.apply(responsePublisher);
+    }
+
+    @Override
+    public void subscribe(CoreSubscriber<? super TResponse> actual) {
+        responsePublisher.subscribe(actual);
+        source.subscribe(new CoreSubscriber<TRequest>() {
+            @Override
+            public void onSubscribe(Subscription s) {
+                requestConsumer.onSubscribe(s);
+            }
+
+            @Override
+            public void onNext(TRequest tRequest) {
+                requestConsumer.onNext(tRequest);
+            }
+
+            @Override
+            public void onError(Throwable throwable) {
+                requestConsumer.onError(throwable);
+            }
+
+            @Override
+            public void onComplete() {
+                requestConsumer.onComplete();
+            }
+
+            @Override
+            public Context currentContext() {
+                return actual.currentContext();
+            }
+        });
+    }
+}

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

@@ -10,6 +10,7 @@
 import org.reactivestreams.Subscription;
 import reactor.core.CoreSubscriber;
 import reactor.core.Scannable;
+import reactor.util.context.Context;
 
 import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.function.BiFunction;
@@ -25,6 +26,11 @@ public class SubscribeOnlyOnceLifter<T> extends AtomicBoolean implements BiFunct
     @Override
     public CoreSubscriber<? super T> apply(Scannable scannable, CoreSubscriber<? super T> coreSubscriber) {
         return new CoreSubscriber<T>() {
+            @Override
+            public Context currentContext() {
+                return coreSubscriber.currentContext();
+            }
+
             @Override
             public void onSubscribe(Subscription subscription) {
                 if (!compareAndSet(false, true)) {

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

@@ -0,0 +1,138 @@
+/*
+ *  Copyright (c) 2019, Salesforce.com, Inc.
+ *  All rights reserved.
+ *  Licensed under the BSD 3-Clause license.
+ *  For full license text, see LICENSE.txt file in the repo root  or https://opensource.org/licenses/BSD-3-Clause
+ */
+
+package com.salesforce.reactorgrpc;
+
+import io.grpc.testing.GrpcServerRule;
+import org.junit.*;
+import reactor.core.publisher.Flux;
+import reactor.core.publisher.Hooks;
+import reactor.core.publisher.Mono;
+import reactor.test.StepVerifier;
+
+import java.util.stream.Collectors;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
+public class ReactorContextPropagationTest {
+
+    @Rule
+    public GrpcServerRule serverRule = new GrpcServerRule();
+
+    private static class SimpleGreeter extends ReactorGreeterGrpc.GreeterImplBase {
+        @Override
+        public Mono<HelloResponse> sayHello(Mono<HelloRequest> request) {
+            return request.map(HelloRequest::getName)
+                    .map(name -> HelloResponse.newBuilder().setMessage("Hello " + name).build());
+        }
+
+        @Override
+        public Mono<HelloResponse> sayHelloReqStream(Flux<HelloRequest> request) {
+            return request.transformDeferredContextual((f, ctx) -> f.map(HelloRequest::getName))
+                    .collect(Collectors.joining("and"))
+                    .map(names -> HelloResponse.newBuilder().setMessage("Hello " + names).build());
+        }
+
+        @Override
+        public Flux<HelloResponse> sayHelloRespStream(Mono<HelloRequest> request) {
+            return request.repeat(2)
+                    .map(HelloRequest::getName)
+                    .zipWith(Flux.just("Hello ", "Hi ", "Greetings "), String::join)
+                    .map(greeting -> HelloResponse.newBuilder().setMessage(greeting).build());
+        }
+
+        @Override
+        public Flux<HelloResponse> sayHelloBothStream(Flux<HelloRequest> request) {
+            return request.map(HelloRequest::getName)
+                    .map(name -> HelloResponse.newBuilder().setMessage("Hello " + name).build());
+        }
+    }
+
+    @BeforeClass
+    public static void beforeAll(){
+        Hooks.enableContextLossTracking();
+        Hooks.onOperatorDebug();
+    }
+
+    @AfterClass
+    public static void afterAll(){
+        Hooks.disableContextLossTracking();
+        Hooks.resetOnOperatorDebug();
+    }
+
+    @Before
+    public void setup() {
+        serverRule.getServiceRegistry().addService(new SimpleGreeter());
+    }
+
+    @Test
+    public void oneToOne() {
+        ReactorGreeterGrpc.ReactorGreeterStub stub = ReactorGreeterGrpc.newReactorStub(serverRule.getChannel());
+        Mono<HelloRequest> req = Mono.just(HelloRequest.newBuilder().setName("reactor").build());
+
+        Mono<HelloResponse> resp = req
+                .doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
+                .transform(stub::sayHello)
+                .doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
+                .contextWrite(ctx -> ctx.put("name", "context"));
+
+        StepVerifier.create(resp)
+                .expectNextCount(1)
+                .verifyComplete();
+    }
+
+    @Test
+    public void oneToMany() {
+        ReactorGreeterGrpc.ReactorGreeterStub stub = ReactorGreeterGrpc.newReactorStub(serverRule.getChannel());
+        Mono<HelloRequest> req = Mono.just(HelloRequest.newBuilder().setName("reactor").build());
+
+        Flux<HelloResponse> resp = req
+                .doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
+                .as(stub::sayHelloRespStream)
+                .doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
+                .contextWrite(ctx -> ctx.put("name", "context"));
+
+        StepVerifier.create(resp)
+                .expectNextCount(3)
+                .verifyComplete();
+    }
+
+    @Test
+    public void manyToOne() {
+        ReactorGreeterGrpc.ReactorGreeterStub stub = ReactorGreeterGrpc.newReactorStub(serverRule.getChannel());
+        Flux<HelloRequest> req = Mono.deferContextual(ctx -> Mono.just(HelloRequest.newBuilder().setName(ctx.get("name")).build())).repeat(2);
+
+        Mono<HelloResponse> resp = req
+                .doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
+                .as(stub::sayHelloReqStream)
+                .doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
+                .contextWrite(ctx -> ctx.put("name", "context"));
+
+        StepVerifier.create(resp)
+                .expectAccessibleContext()
+                .contains("name", "context")
+                .then()
+                .expectNextCount(1)
+                .verifyComplete();
+    }
+
+    @Test
+    public void manyToMany() {
+        ReactorGreeterGrpc.ReactorGreeterStub stub = ReactorGreeterGrpc.newReactorStub(serverRule.getChannel());
+        Flux<HelloRequest> req = Mono.just(HelloRequest.newBuilder().setName("reactor").build()).repeat(2).contextWrite(c -> c.put("name", "boom"));
+
+        Flux<HelloResponse> resp = req
+                .doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
+                .transform(stub::sayHelloBothStream)
+                .doOnEach(signal -> assertThat(signal.getContextView().getOrEmpty("name")).isNotEmpty())
+                .contextWrite(ctx -> ctx.put("name", "context"));
+
+        StepVerifier.create(resp)
+                .expectNextCount(3)
+                .verifyComplete();
+    }
+}