Address review comments

src/main/java/io/projectriff/invoker/server/GrpcServerAdapter.java

@@ -21,6 +21,7 @@
 import reactor.core.publisher.GroupedFlux;
 import reactor.core.publisher.Operators;
 import reactor.core.publisher.Signal;
+import reactor.util.context.Context;
 import reactor.util.function.Tuple2;
 import reactor.util.function.Tuples;
 
@@ -149,8 +150,8 @@ private Function<Flux<Tuple2<Integer, Message<byte[]>>>, Flux<Tuple2<Integer, Me
     // Used to transform the publisher chain into one that doesn't forward cancel() calls once it has complete()d.
     private Function<? super Publisher<Tuple2<Integer, Message<byte[]>>>, ? extends Publisher<Tuple2<Integer, Message<byte[]>>>> ignoreCancelsAfterComplete() {
         return Operators.lift((f, actual) ->
-                new CoreSubscriber<>() {
-                    AtomicBoolean completed = new AtomicBoolean();
+                new CoreSubscriber<Tuple2<Integer, Message<byte[]>>>() {
+                    private volatile boolean completed;
 
                     @Override
                     public void onSubscribe(Subscription s) {
@@ -162,7 +163,7 @@ public void request(long n) {
 
                             @Override
                             public void cancel() {
-                                if (!completed.get()) {
+                                if (!completed) {
                                     s.cancel();
                                 }
                             }
@@ -181,9 +182,14 @@ public void onError(Throwable t) {
 
                     @Override
                     public void onComplete() {
-                        completed.compareAndSet(false, true);
+                        completed = true;
                         actual.onComplete();
                     }
+
+                    @Override
+                    public Context currentContext() {
+                        return actual.currentContext();
+                    }
                 });
     }