Repository: salesforce/reactive-grpc
Pull request #293 (2022-02-14 14:49:56)

issue #292 - discard support
please review, @rmichela 

salesforce-cla[bot]: Thanks for the contribution!  Unfortunately we can't verify the commit author(s): Kirill Rodionov <k***@g***.com>.  One possible solution is to [add that email to your GitHub account](https://github.com/settings/emails).  Alternatively you can [change your commits to another email](https://help.github.com/articles/why-are-my-commits-linked-to-the-wrong-user/#commits-are-not-linked-to-any-user) and force push the change.  After getting your commits associated with your GitHub account, [sign the Salesforce.com Contributor License Agreement](https://cla.salesforce.com/sign-cla?prUrl=https%3A%2F%2Fgithub.com%2Fsalesforce%2Freactive-grpc%2Fpull%2F293) and this Pull Request will be revalidated.
salesforce-cla[bot]: Thanks for the contribution!  Before we can merge this, we need @bruto1 to [sign the Salesforce.com Contributor License Agreement](https://cla.salesforce.com/sign-cla?prUrl=https%3A%2F%2Fgithub.com%2Fsalesforce%2Freactive-grpc%2Fpull%2F293).

Patch:
diff --git a/common/reactive-grpc-common/src/main/java/com/salesforce/reactivegrpc/common/AbstractStreamObserverAndPublisher.java b/common/reactive-grpc-common/src/main/java/com/salesforce/reactivegrpc/common/AbstractStreamObserverAndPublisher.java
index 00348fe..b2afa64 100644
--- a/common/reactive-grpc-common/src/main/java/com/salesforce/reactivegrpc/common/AbstractStreamObserverAndPublisher.java
+++ b/common/reactive-grpc-common/src/main/java/com/salesforce/reactivegrpc/common/AbstractStreamObserverAndPublisher.java
@@ -85,7 +85,7 @@ public abstract class AbstractStreamObserverAndPublisher<T> extends AbstractUnim
     private volatile boolean done;
     private Throwable error;
 
-    private volatile Subscriber<? super T> downstream;
+    protected volatile Subscriber<? super T> downstream;
 
     private volatile boolean cancelled;
 
@@ -226,7 +226,7 @@ public abstract class AbstractStreamObserverAndPublisher<T> extends AbstractUnim
 
         for (;;) {
             if (cancelled) {
-                queue.clear();
+                discardQueue(queue);
                 downstream = null;
                 return;
             }
@@ -283,7 +283,7 @@ public abstract class AbstractStreamObserverAndPublisher<T> extends AbstractUnim
 
     private boolean checkTerminated(boolean d, boolean empty, Subscriber<? super T> subscriber, Queue<T> q) {
         if (cancelled) {
-            q.clear();
+            discardQueue(q);
             downstream = null;
             return true;
         }
@@ -305,6 +305,7 @@ public abstract class AbstractStreamObserverAndPublisher<T> extends AbstractUnim
     @Override
     public void onNext(T t) {
         if (done || cancelled) {
+            discardElement(t);
             return;
         }
 
@@ -419,7 +420,7 @@ public abstract class AbstractStreamObserverAndPublisher<T> extends AbstractUnim
 
         if (!outputFused) {
             if (WIP.getAndIncrement(this) == 0) {
-                queue.clear();
+                discardQueue(queue);
                 downstream = null;
             }
         }
@@ -456,4 +457,10 @@ public abstract class AbstractStreamObserverAndPublisher<T> extends AbstractUnim
     public void clear() {
         queue.clear();
     }
+
+    protected void discardQueue(Queue<T> q) {
+        q.clear();
+    }
+
+    protected void discardElement(T t) { }
 }
\ No newline at end of file
diff --git a/reactor/reactor-grpc-stub/pom.xml b/reactor/reactor-grpc-stub/pom.xml
index a4a6f75..92a09fa 100644
--- a/reactor/reactor-grpc-stub/pom.xml
+++ b/reactor/reactor-grpc-stub/pom.xml
@@ -67,6 +67,19 @@
             <artifactId>mockito-core</artifactId>
             <scope>test</scope>
         </dependency>
+        <!-- https://mvnrepository.com/artifact/org.slf4j/slf4j-api -->
+        <dependency>
+            <groupId>org.slf4j</groupId>
+            <artifactId>slf4j-api</artifactId>
+            <version>1.7.36</version>
+            <scope>test</scope>
+        </dependency>
+        <dependency>
+            <groupId>ch.qos.logback</groupId>
+            <artifactId>logback-classic</artifactId>
+            <version>1.2.10</version>
+            <scope>test</scope>
+        </dependency>
     </dependencies>
 
     <build>
diff --git a/reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ReactorClientStreamObserverAndPublisher.java b/reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ReactorClientStreamObserverAndPublisher.java
index 749fd30..12fb1c3 100644
--- a/reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ReactorClientStreamObserverAndPublisher.java
+++ b/reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ReactorClientStreamObserverAndPublisher.java
@@ -10,9 +10,13 @@ package com.salesforce.reactorgrpc.stub;
 import com.salesforce.reactivegrpc.common.AbstractClientStreamObserverAndPublisher;
 import com.salesforce.reactivegrpc.common.Consumer;
 import io.grpc.stub.CallStreamObserver;
+import reactor.core.CoreSubscriber;
 import reactor.core.Fuseable;
+import reactor.core.publisher.Operators;
 import reactor.util.concurrent.Queues;
 
+import java.util.Queue;
+
 /**
  * TODO: Explain what this class does.
  * @param <T> T
@@ -46,4 +50,20 @@ class ReactorClientStreamObserverAndPublisher<T>
         }
         return Fuseable.NONE;
     }
+
+    @Override
+    protected void discardQueue(Queue<T> q) {
+        if (downstream instanceof CoreSubscriber) {
+            Operators.onDiscardQueueWithClear(q, ((CoreSubscriber) downstream).currentContext(), null);
+        } else {
+            q.clear();
+        }
+    }
+
+    @Override
+    protected void discardElement(T t) {
+        if (downstream instanceof CoreSubscriber) {
+            Operators.onDiscard(t, ((CoreSubscriber) downstream).currentContext());
+        }
+    }
 }
diff --git a/reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ReactorServerStreamObserverAndPublisher.java b/reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ReactorServerStreamObserverAndPublisher.java
index 1619532..33bac86 100644
--- a/reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ReactorServerStreamObserverAndPublisher.java
+++ b/reactor/reactor-grpc-stub/src/main/java/com/salesforce/reactorgrpc/stub/ReactorServerStreamObserverAndPublisher.java
@@ -11,9 +11,13 @@ import com.salesforce.reactivegrpc.common.AbstractServerStreamObserverAndPublish
 import com.salesforce.reactivegrpc.common.Consumer;
 import io.grpc.stub.CallStreamObserver;
 import io.grpc.stub.ServerCallStreamObserver;
+import reactor.core.CoreSubscriber;
 import reactor.core.Fuseable;
+import reactor.core.publisher.Operators;
 import reactor.util.concurrent.Queues;
 
+import java.util.Queue;
+
 /**
  * TODO: Explain what this class does.
  * @param <T> T
@@ -38,4 +42,20 @@ class ReactorServerStreamObserverAndPublisher<T>
         }
         return Fuseable.NONE;
     }
+
+    @Override
+    protected void discardQueue(Queue<T> q) {
+        if (downstream instanceof CoreSubscriber) {
+            Operators.onDiscardQueueWithClear(q, ((CoreSubscriber) downstream).currentContext(), null);
+        } else {
+            q.clear();
+        }
+    }
+
+    @Override
+    protected void discardElement(T t) {
+        if (downstream instanceof CoreSubscriber) {
+            Operators.onDiscard(t, ((CoreSubscriber) downstream).currentContext());
+        }
+    }
 }
\ No newline at end of file


Test patch:
diff --git a/reactor/reactor-grpc-stub/src/test/java/com/salesforce/reactorgrpc/stub/ReactorClientStreamObserverAndPublisherTest.java b/reactor/reactor-grpc-stub/src/test/java/com/salesforce/reactorgrpc/stub/ReactorClientStreamObserverAndPublisherTest.java
index dd44527..12c0a08 100644
--- a/reactor/reactor-grpc-stub/src/test/java/com/salesforce/reactorgrpc/stub/ReactorClientStreamObserverAndPublisherTest.java
+++ b/reactor/reactor-grpc-stub/src/test/java/com/salesforce/reactorgrpc/stub/ReactorClientStreamObserverAndPublisherTest.java
@@ -6,17 +6,23 @@
 
 package com.salesforce.reactorgrpc.stub;
 
-import java.util.concurrent.ForkJoinPool;
-
-import org.assertj.core.api.Assertions;
 import org.junit.jupiter.api.Test;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
 import reactor.core.Fuseable;
 import reactor.core.publisher.Flux;
 import reactor.core.scheduler.Schedulers;
 import reactor.test.StepVerifier;
 
+import java.util.concurrent.ConcurrentLinkedQueue;
+import java.util.concurrent.ForkJoinPool;
+import java.util.concurrent.atomic.AtomicBoolean;
+
+import static org.assertj.core.api.Assertions.assertThat;
+
 
 public class ReactorClientStreamObserverAndPublisherTest {
+    private static final Logger log = LoggerFactory.getLogger(ReactorClientStreamObserverAndPublisherTest.class.getName());
 
     private static final int DEFAULT_CHUNK_SIZE = 512;
     private static final int PART_OF_CHUNK = DEFAULT_CHUNK_SIZE * 2 / 3;
@@ -24,7 +30,7 @@ public class ReactorClientStreamObserverAndPublisherTest {
     @Test
     public void multiThreadedProducerTest() {
         ReactorClientStreamObserverAndPublisher<Integer> processor =
-                new ReactorClientStreamObserverAndPublisher<>(null);
+            new ReactorClientStreamObserverAndPublisher<>(null);
         int countPerThread = 100000;
         TestCallStreamObserverProducer observer =
             new TestCallStreamObserverProducer(ForkJoinPool.commonPool(), processor, countPerThread);
@@ -34,21 +40,76 @@ public class ReactorClientStreamObserverAndPublisherTest {
                     .expectNextCount(countPerThread)
                     .verifyComplete();
 
-        Assertions.assertThat(observer.requestsQueue.size()).isBetween((countPerThread - DEFAULT_CHUNK_SIZE) / PART_OF_CHUNK + 1, (countPerThread - DEFAULT_CHUNK_SIZE) / PART_OF_CHUNK + 3);
+        assertThat(observer.requestsQueue.size()).isBetween((countPerThread - DEFAULT_CHUNK_SIZE) / PART_OF_CHUNK + 1,
+            (countPerThread - DEFAULT_CHUNK_SIZE) / PART_OF_CHUNK + 3);
     }
 
     @Test
     public void producerFusedTest() {
         ReactorClientStreamObserverAndPublisher<Integer> processor =
-                new ReactorClientStreamObserverAndPublisher<>(null);
+            new ReactorClientStreamObserverAndPublisher<>(null);
         int countPerThread = 100000;
-        TestCallStreamObserverProducer observer = new TestCallStreamObserverProducer(ForkJoinPool.commonPool(), processor, countPerThread);
+        TestCallStreamObserverProducer observer = new TestCallStreamObserverProducer(ForkJoinPool.commonPool(),
+            processor, countPerThread);
         processor.beforeStart(observer);
         StepVerifier.create(Flux.from(processor))
                     .expectFusion(Fuseable.ANY, Fuseable.ASYNC)
                     .expectNextCount(countPerThread)
                     .verifyComplete();
 
-        Assertions.assertThat(observer.requestsQueue.size()).isBetween((countPerThread - DEFAULT_CHUNK_SIZE) / PART_OF_CHUNK + 1, (countPerThread - DEFAULT_CHUNK_SIZE) / PART_OF_CHUNK + 3);
+        assertThat(observer.requestsQueue.size()).isBetween((countPerThread - DEFAULT_CHUNK_SIZE) / PART_OF_CHUNK + 1,
+            (countPerThread - DEFAULT_CHUNK_SIZE) / PART_OF_CHUNK + 3);
+    }
+
+    @Test
+    public void discardQueueTest() {
+        ReactorClientStreamObserverAndPublisher<Integer> processor =
+            new ReactorClientStreamObserverAndPublisher<>(null);
+        int countPerThread = 5;
+        TestCallStreamObserverProducer observer = new TestCallStreamObserverProducer(ForkJoinPool.commonPool(),
+            processor, countPerThread);
+        processor.beforeStart(observer);
+
+        ConcurrentLinkedQueue<Integer> discardedByObserverAndPublisher = new ConcurrentLinkedQueue<>();
+        ConcurrentLinkedQueue<Integer> discardedByPublishOn = new ConcurrentLinkedQueue<>();
+
+        AtomicBoolean firstHandled = new AtomicBoolean();
+        Flux<Integer> consumer =
+            Flux.from(processor)
+                .doOnDiscard(Integer.class, i -> {
+                    log.info("Processor: discarding {}", i);
+                    discardedByObserverAndPublisher.add(i);
+                })
+                .log("processor")
+                .limitRate(1)
+                .publishOn(Schedulers.parallel())
+                .limitRate(1)
+                .doOnDiscard(Integer.class, i -> {
+                    log.info("publishOn: discarding {}", i);
+                    discardedByPublishOn.add(i);
+                })
+                .<Integer>handle((i, sink) -> {
+                    if (firstHandled.compareAndSet(false, true)) {
+                        try {
+                            Thread.sleep(100);
+                        } catch (Exception e) {
+                            // noop
+                        }
+                        sink.next(i);
+                    } else {
+                        sink.complete();
+                    }
+                })
+                .log("handled");
+
+        StepVerifier.create(consumer)
+                    .expectNext(0)
+                    .verifyComplete();
+
+        // 1 is dropped in handle without invoking the discard hook,
+        assertThat(discardedByObserverAndPublisher).containsExactly(3, 4);
+        // impl details: processor is able to schedule 2 before it's cancelled
+        // also, discard hooks are cumulative, so not using containsExactly
+        assertThat(discardedByPublishOn).contains(2);
     }
 }
\ No newline at end of file