[changelog] Clean up consumer state on changelog client (#1398)

* [chanelog] Clean up consumer state on changelog client

And other protections like subscribing to an empty set of partitions.

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

@@ -81,11 +81,17 @@ public Collection<PubSubMessage<K, ChangeEvent<V>, VeniceChangeCoordinate>> poll
 
   @Override
   public CompletableFuture<Void> seekToTimestamps(Map<Integer, Long> timestamps) {
+    if (timestamps.isEmpty()) {
+      return CompletableFuture.completedFuture(null);
+    }
     return internalSeekToTimestamps(timestamps, "");
   }
 
   @Override
   public CompletableFuture<Void> subscribe(Set<Integer> partitions) {
+    if (partitions.isEmpty()) {
+      return CompletableFuture.completedFuture(null);
+    }
     if (!versionSwapThreadScheduled.get()) {
       // schedule the version swap thread and set up the callback listener
       this.storeRepository.registerStoreDataChangedListener(versionSwapListener);
@@ -101,18 +107,25 @@ public CompletableFuture<Void> subscribe(Set<Integer> partitions) {
 
   @Override
   public CompletableFuture<Void> seekToTail(Set<Integer> partitions) {
+    if (partitions.isEmpty()) {
+      return CompletableFuture.completedFuture(null);
+    }
     return internalSeekToTail(partitions, "");
   }
 
   @Override
   public CompletableFuture<Void> seekToEndOfPush(Set<Integer> partitions) {
+    if (partitions.isEmpty()) {
+      return CompletableFuture.completedFuture(null);
+    }
     return CompletableFuture.supplyAsync(() -> {
       synchronized (internalSeekConsumer) {
         try {
           // TODO: This implementation basically just scans the version topic until it finds the EOP message. The
           // approach
           // we'd like to do is instead add the offset of the EOP message in the VT, and then just seek to that offset.
           // We'll do that in a future patch.
+          internalSeekConsumer.get().unsubscribeAll();
           internalSeekConsumer.get().subscribe(partitions).get();
 
           // We need to get the internal consumer as we have to intercept the control messages that we would normally

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

@@ -239,7 +239,7 @@ protected CompletableFuture<Void> internalSubscribe(Set<Integer> partitions, Pub
 
       synchronized (pubSubConsumer) {
         Set<PubSubTopicPartition> topicPartitionSet = getTopicAssignment();
-        for (PubSubTopicPartition topicPartition: pubSubConsumer.getAssignment()) {
+        for (PubSubTopicPartition topicPartition: topicPartitionSet) {
           if (partitions.contains(topicPartition.getPartitionNumber())) {
             pubSubConsumer.unSubscribe(topicPartition);
           }
@@ -549,6 +549,9 @@ private List<PubSubTopicPartition> getPartitionListToSubscribe(
 
   @Override
   public void unsubscribe(Set<Integer> partitions) {
+    if (partitions.isEmpty()) {
+      return;
+    }
     synchronized (pubSubConsumer) {
       Set<PubSubTopicPartition> topicPartitionSet = getTopicAssignment();
       Set<PubSubTopicPartition> topicPartitionsToUnsub = new HashSet<>();

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

@@ -56,6 +56,10 @@ public void handleStoreChanged(Store store) {
           }
         }
 
+        if (partitions.isEmpty()) {
+          return;
+        }
+
         LOGGER.info(
             "New Version detected!  Seeking consumer to version: " + currentVersion + " in consumer: " + consumerName);
         this.consumer.seekToEndOfPush(partitions).get();

clients/da-vinci-client/src/test/java/com/linkedin/davinci/consumer/VeniceChangelogConsumerImplTest.java

@@ -229,6 +229,11 @@ public void testAfterImageConsumerSeek() throws ExecutionException, InterruptedE
 
     Set<Integer> partitionSet = new HashSet<>();
     partitionSet.add(0);
+
+    // try and do it with an empty set, this should complete immediately
+    veniceChangelogConsumer.seekToEndOfPush(Collections.emptySet()).get();
+    Mockito.verifyNoInteractions(mockInternalSeekConsumer);
+
     veniceChangelogConsumer.seekToEndOfPush(partitionSet).get();
 
     Mockito.verify(mockInternalSeekConsumer).subscribe(partitionSet);