fix: found a callback to register if a downstream consumer cancels

Signed-off-by: Matt Peterson <[email protected]>

Author: mattp-swirldslabs

server/src/main/java/com/hedera/block/server/BlockStreamService.java

@@ -118,7 +118,6 @@ private StreamObserver<BlockStreamServiceGrpcProto.BlockItemResponse> streamSour
         final var streamObserver = new ConsumerBlockItemObserver(
                 timeoutThresholdMillis,
                 Clock.systemDefaultZone(),
-                Clock.systemDefaultZone(),
                 streamMediator,
                 responseStreamObserver);
 

server/src/main/java/com/hedera/block/server/consumer/ConsumerBlockItemObserver.java

@@ -19,6 +19,7 @@
 import com.hedera.block.protos.BlockStreamServiceGrpcProto;
 import com.hedera.block.server.data.ObjectEvent;
 import com.hedera.block.server.mediator.StreamMediator;
+import io.grpc.stub.ServerCallStreamObserver;
 import io.grpc.stub.StreamObserver;
 
 import java.time.Clock;
@@ -40,9 +41,6 @@ public class ConsumerBlockItemObserver implements BlockItemEventHandler<ObjectEv
     private final InstantSource producerLivenessClock;
     private long producerLivenessMillis;
 
-    private final InstantSource consumerLivenessClock;
-    private long consumerLivenessMillis;
-
     private final CountDownLatch shutdownLatch = new CountDownLatch(1);
 
     private final StreamMediator<ObjectEvent<BlockStreamServiceGrpcProto.BlockItem>, BlockStreamServiceGrpcProto.BlockItemResponse> streamMediator;
@@ -55,17 +53,29 @@ public class ConsumerBlockItemObserver implements BlockItemEventHandler<ObjectEv
     public ConsumerBlockItemObserver(
             final long timeoutThresholdMillis,
             final InstantSource producerLivenessClock,
-            final InstantSource consumerLivenessClock,
             final StreamMediator<ObjectEvent<BlockStreamServiceGrpcProto.BlockItem>, BlockStreamServiceGrpcProto.BlockItemResponse> streamMediator,
             final StreamObserver<BlockStreamServiceGrpcProto.BlockItem> responseStreamObserver) {
 
         this.timeoutThresholdMillis = timeoutThresholdMillis;
         this.producerLivenessClock = producerLivenessClock;
-        this.consumerLivenessClock = consumerLivenessClock;
-        this.responseStreamObserver = responseStreamObserver;
 
+        // The ServerCallStreamObserver can be configured with a Runnable to
+        // be executed if a downstream consumer cancels the stream without
+        // sending an HTTP/2 End Stream DATA frame.  If triggered, unsubscribe
+        // this observer to avoid orphaning subscribed resources.
+        if (responseStreamObserver instanceof ServerCallStreamObserver) {
+
+            // Unfortunately we have to cast the responseStreamObserver to a ServerCallStreamObserver
+            // to register the onCancelHandler.
+            ((ServerCallStreamObserver<BlockStreamServiceGrpcProto.BlockItem>)responseStreamObserver)
+                    .setOnCancelHandler(() -> {
+                        LOGGER.log(System.Logger.Level.DEBUG, "Consumer cancelled stream.  Unsubscribing observer.");
+                        streamMediator.unsubscribe(this);
+                    });
+        }
+
+        this.responseStreamObserver = responseStreamObserver;
         this.producerLivenessMillis = producerLivenessClock.millis();
-        this.consumerLivenessMillis = consumerLivenessClock.millis();
 
         this.streamMediator = streamMediator;
     }
@@ -77,15 +87,9 @@ public ConsumerBlockItemObserver(
     @Override
     public void onEvent(final ObjectEvent<BlockStreamServiceGrpcProto.BlockItem> event, final long l, final boolean b) throws Exception {
 
-        // Check if the consumer has timed out.  If so, unsubscribe the observer from the mediator.
-        if (isThresholdExceeded(consumerLivenessMillis)) {
-            LOGGER.log(System.Logger.Level.DEBUG, "Consumer timeout threshold exceeded.");
-            streamMediator.unsubscribe(this);
-        } else {
-            // Refresh the producer liveness and pass the block to the observer.
-            producerLivenessMillis = producerLivenessClock.millis();
-            responseStreamObserver.onNext(event.get());
-        }
+        // Refresh the producer liveness and pass the block to the observer.
+        producerLivenessMillis = producerLivenessClock.millis();
+        responseStreamObserver.onNext(event.get());
     }
 
     /**
@@ -100,10 +104,6 @@ public void onNext(final BlockStreamServiceGrpcProto.BlockItemResponse blockItem
         if (isThresholdExceeded(producerLivenessMillis)) {
             LOGGER.log(System.Logger.Level.DEBUG, "Producer timeout threshold exceeded.  Unsubscribing observer.");
             streamMediator.unsubscribe(this);
-        } else {
-            // Refresh the consumer liveness
-            LOGGER.log(System.Logger.Level.DEBUG, "Received response block " + blockItemResponse);
-            consumerLivenessMillis = consumerLivenessClock.millis();
         }
     }