[cc] Unblock shared drainer on CDC shutdown

clients/da-vinci-client/src/main/java/com/linkedin/davinci/consumer/VeniceChangelogConsumerDaVinciRecordTransformerImpl.java

@@ -81,15 +81,17 @@ public class VeniceChangelogConsumerDaVinciRecordTransformerImpl<K, V>
   private final CachingDaVinciClientFactory daVinciClientFactory;
   private final SeekableDaVinciClient<K, V> daVinciClient;
   private final AtomicBoolean isStarted = new AtomicBoolean(false);
+  private final AtomicBoolean isClosed = new AtomicBoolean(false);
   private final CountDownLatch startLatch = new CountDownLatch(1);
   // Using a dedicated thread pool for CompletableFutures created by this class to avoid potential thread starvation
   // issues in the default ForkJoinPool
   private final ExecutorService completableFutureThreadPool;
 
   private final Set<Integer> subscribedPartitions = VeniceConcurrentHashMap.newKeySet();
   private final Set<Integer> eopReceivedPartitions = VeniceConcurrentHashMap.newKeySet();
-  private final ReentrantLock bufferLock = new ReentrantLock();
-  private final Condition bufferIsFullCondition = bufferLock.newCondition();
+  // Coordinates poll waiting/draining with shutdown. ArrayBlockingQueue manages insertion locking.
+  private final ReentrantLock pubSubMessagesStateLock = new ReentrantLock();
+  private final Condition pubSubMessagesFullCondition = pubSubMessagesStateLock.newCondition();
   private volatile BackgroundReporterThread backgroundReporterThread;
   private final BasicConsumerStats changeCaptureStats;
   private final AtomicBoolean isCaughtUp = new AtomicBoolean(false);
@@ -197,13 +199,23 @@ public VeniceChangelogConsumerDaVinciRecordTransformerImpl(
   }
 
   private synchronized void startDaVinciClient() {
+    throwIfClosed();
     // Start daVinci client if not already started
     if (!isStarted.get()) {
       daVinciClient.start();
       isStarted.set(true);
     }
   }
 
+  private void throwIfClosed() {
+    if (isClosed.get()) {
+      throw new VeniceClientException(
+          "This VeniceChangelogConsumer instance is closed and cannot be restarted. Create a new consumer from "
+              + "VeniceChangelogConsumerClientFactory for store: " + storeName + ", consumer name: "
+              + changelogClientConfig.getConsumerName());
+    }
+  }
+
   /**
    * Helper method to initialize client, update subscribed partitions, and execute subscription.
    * This consolidates common logic across start, seekToCheckpoint, and seekToTimestamps.
@@ -215,6 +227,7 @@ private synchronized void startDaVinciClient() {
   private synchronized CompletableFuture<Void> initializeAndSubscribe(
       Set<Integer> partitions,
       Function<Set<Integer>, CompletableFuture<Void>> subscriptionCall) {
+    throwIfClosed();
     Set<Integer> targetPartitions = new HashSet<>();
     boolean partitionsAdded = false;
     try {
@@ -321,18 +334,33 @@ public CompletableFuture<Void> start() {
   }
 
   @Override
-  public void stop() throws Exception {
-    LOGGER.info("Closing Changelog Consumer with name: {}", changelogClientConfig.getConsumerName());
+  public synchronized void stop() throws Exception {
+    if (isClosed.get()) {
+      return;
+    }
+
+    LOGGER.info("Closing VeniceChangelogConsumer with name: {}", changelogClientConfig.getConsumerName());
+    pubSubMessagesStateLock.lock();
+    try {
+      isClosed.set(true);
+      isStarted.set(false);
+      partitionToVersionToServe.clear();
+      pubSubMessages.clear();
+      pubSubMessagesFullCondition.signalAll();
+    } finally {
+      pubSubMessagesStateLock.unlock();
+    }
     try {
       if (backgroundReporterThread != null) {
         backgroundReporterThread.interrupt();
       }
       daVinciClient.close();
     } finally {
       isStarted.set(false);
+      pubSubMessages.clear();
       veniceChangelogConsumerClientFactory.deregisterClient(changelogClientConfig.getConsumerName());
       clearPartitionState(Collections.emptySet());
-      LOGGER.info("Closed Changelog Consumer with name: {}", changelogClientConfig.getConsumerName());
+      LOGGER.info("Closed VeniceChangelogConsumer with name: {}", changelogClientConfig.getConsumerName());
     }
   }
 
@@ -453,23 +481,26 @@ public void close() {
   @Override
   public Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> poll(long timeoutInMs) {
     try {
+      Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> drainedPubSubMessages = new ArrayList<>();
       try {
-        bufferLock.lock();
+        pubSubMessagesStateLock.lock();
 
         // Wait until pubSubMessages becomes full, or until the timeout is reached
-        if (pubSubMessages.remainingCapacity() > 0) {
-          bufferIsFullCondition.await(timeoutInMs, TimeUnit.MILLISECONDS);
+        if (!isClosed.get() && pubSubMessages.remainingCapacity() > 0) {
+          pubSubMessagesFullCondition.await(timeoutInMs, TimeUnit.MILLISECONDS);
+        }
+        if (isClosed.get()) {
+          return Collections.emptyList();
         }
+        pubSubMessages.drainTo(drainedPubSubMessages);
       } catch (InterruptedException exception) {
         LOGGER.info("Thread was interrupted", exception);
         // Restore the interrupt status
         Thread.currentThread().interrupt();
       } finally {
-        bufferLock.unlock();
+        pubSubMessagesStateLock.unlock();
       }
 
-      Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> drainedPubSubMessages = new ArrayList<>();
-      pubSubMessages.drainTo(drainedPubSubMessages);
       int messagesPolled = drainedPubSubMessages.size();
 
       if (changelogClientConfig.shouldCompactMessages()) {
@@ -784,12 +815,16 @@ public void addControlMessageToBuffer(
     private void internalAddMessageToBuffer(
         int partitionId,
         ImmutableChangeCapturePubSubMessage<K, ChangeEvent<V>> pubSubMessage) {
-      if (!isServedByThisVersion(partitionId)) {
+      if (isClosed.get() || !isStarted.get() || !isServedByThisVersion(partitionId)) {
         return;
       }
 
       try {
         pubSubMessages.put(pubSubMessage);
+        if (isClosed.get() || !isStarted.get()) {
+          pubSubMessages.remove(pubSubMessage);
+          return;
+        }
         /*
          * pubSubMessages is full, signal to a poll thread awaiting on bufferFullCondition.
          * Not signaling to all threads, because if multiple poll threads try to read pubSubMessages at
@@ -801,11 +836,11 @@ private void internalAddMessageToBuffer(
          * beginning of this comment block.
          */
         if (pubSubMessages.remainingCapacity() == 0) {
-          bufferLock.lock();
+          pubSubMessagesStateLock.lock();
           try {
-            bufferIsFullCondition.signal();
+            pubSubMessagesFullCondition.signal();
           } finally {
-            bufferLock.unlock();
+            pubSubMessagesStateLock.unlock();
           }
         }
 
@@ -970,6 +1005,11 @@ public boolean isStarted() {
     return isStarted.get();
   }
 
+  @VisibleForTesting
+  public boolean isClosed() {
+    return isClosed.get();
+  }
+
   @VisibleForTesting
   public Set<Integer> getSubscribedPartitions() {
     return subscribedPartitions;