wip: added awaitTimeout

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

Author: mattp-swirldslabs

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

@@ -19,4 +19,5 @@
 import io.grpc.stub.StreamObserver;
 
 public interface BlockEventHandler<U, V> extends StreamObserver<V>, EventHandler<U> {
+    void awaitShutdown() throws InterruptedException;
 }

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

@@ -23,6 +23,7 @@
 
 import java.time.Clock;
 import java.time.InstantSource;
+import java.util.concurrent.CountDownLatch;
 
 /**
  * The LiveStreamObserverImpl class implements the LiveStreamObserver interface to pass blocks to the downstream consumer
@@ -42,6 +43,8 @@ public class ConsumerBlockStreamObserver implements BlockEventHandler<ObjectEven
     private final InstantSource consumerLivenessClock;
     private long consumerLivenessMillis;
 
+    private final CountDownLatch shutdownLatch = new CountDownLatch(1);
+
     private final StreamMediator<ObjectEvent<BlockStreamServiceGrpcProto.Block>, BlockStreamServiceGrpcProto.BlockResponse> streamMediator;
 
     /**
@@ -135,4 +138,9 @@ private boolean isThresholdExceeded(long livenessMillis) {
 
         return false;
     }
+
+    @Override
+    public void awaitShutdown() throws InterruptedException {
+        shutdownLatch.await();
+    }
 }

server/src/main/java/com/hedera/block/server/mediator/LiveStreamMediatorImpl.java

@@ -64,15 +64,19 @@ public LiveStreamMediatorImpl(final BlockPersistenceHandler<BlockStreamServiceGr
 
     @Override
     public void publishEvent(BlockStreamServiceGrpcProto.Block block) {
+
+        // Publish the block for all subscribers to receive
         LOGGER.log(System.Logger.Level.INFO, "Publishing block: {0}", block);
         ringBuffer.publishEvent((event, sequence) -> event.set(block));
-        
+
+        // Block persistence
         blockPersistenceHandler.persist(block);
     }
 
     @Override
     public void subscribe(final BlockEventHandler<ObjectEvent<BlockStreamServiceGrpcProto.Block>, BlockStreamServiceGrpcProto.BlockResponse> handler) {
 
+        // Initialize the batch event processor and set it on the ring buffer
         final var batchEventProcessor = new BatchEventProcessorBuilder()
                 .build(ringBuffer, ringBuffer.newBarrier(), handler);
 
@@ -90,7 +94,11 @@ public void unsubscribe(final BlockEventHandler<ObjectEvent<BlockStreamServiceGr
         batchEventProcessor.halt();
 
         // Wait for shutdown the complete
-//        handler.awaitShutdown();
+        try {
+            handler.awaitShutdown();
+        } catch (InterruptedException e) {
+            LOGGER.log(System.Logger.Level.ERROR, "Error occurred while waiting for shutdown", e);
+        }
 
         // Remove the gating sequence from the ring buffer
         ringBuffer.removeGatingSequence(batchEventProcessor.getSequence());