[fix][client] Fix multi-topics consumer could receive old messages after seek

Author: BewareMyPower

pulsar-broker/src/test/java/org/apache/pulsar/client/api/MultiTopicsConsumerTest.java

@@ -29,16 +29,22 @@
 
 import com.google.common.collect.Lists;
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.Collections;
+import java.util.Comparator;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.TreeSet;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
 import lombok.Cleanup;
@@ -371,4 +377,66 @@ public void testMultipleIOThreads() throws PulsarAdminException, PulsarClientExc
         assertTrue(consumer instanceof MultiTopicsConsumerImpl);
         assertTrue(consumer.isConnected());
     }
+
+    @DataProvider
+    public static Object[][] seekByFunction() {
+        return new Object[][] {
+                { true }, { false }
+        };
+    }
+
+    @Test(timeOut = 30000, dataProvider = "seekByFunction")
+    public void testSeekByTimestamp(boolean seekByFunction) throws Exception {
+        final var topic1 = TopicName.get(newTopicName()).toString();
+        final var topic2 = TopicName.get(newTopicName()).toString();
+        @Cleanup final var producer1 = pulsarClient.newProducer(Schema.STRING).topic(topic1).create();
+        @Cleanup final var producer2 = pulsarClient.newProducer(Schema.STRING).topic(topic2).create();
+        producer1.send("1-0");
+        producer2.send("2-0");
+        producer1.send("1-1");
+        producer2.send("2-1");
+        final var consumer1 = pulsarClient.newConsumer(Schema.STRING)
+                .topics(Arrays.asList(topic1, topic2)).subscriptionName("sub")
+                .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest).subscribe();
+        final var timestamps = new ArrayList<Long>();
+        for (int i = 0; i < 4; i++) {
+            timestamps.add(consumer1.receive().getPublishTime());
+        }
+        timestamps.sort(Comparator.naturalOrder());
+        final var timestamp = timestamps.get(2);
+        consumer1.close();
+
+        final Function<Consumer<String>, CompletableFuture<Void>> seekAsync = consumer ->
+                seekByFunction ? consumer.seekAsync(__ -> timestamp) : consumer.seekAsync(timestamp);
+
+        @Cleanup final var consumer2 = pulsarClient.newConsumer(Schema.STRING)
+                .topics(Arrays.asList(topic1, topic2)).subscriptionName("sub-2")
+                .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest).subscribe();
+        seekAsync.apply(consumer2).get();
+        final var values = new TreeSet<String>();
+        for (int i = 0; i < 2; i++) {
+            values.add(consumer2.receive().getValue());
+        }
+        assertEquals(values, new TreeSet<>(Arrays.asList("1-1", "2-1")));
+
+        final var valuesInListener = new CopyOnWriteArrayList<String>();
+        @Cleanup final var consumer3 = pulsarClient.newConsumer(Schema.STRING)
+                .topics(Arrays.asList(topic1, topic2)).subscriptionName("sub-3")
+                .messageListener((MessageListener<String>) (__, msg) -> valuesInListener.add(msg.getValue()))
+                .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest).subscribe();
+        seekAsync.apply(consumer3).get();
+        if (valuesInListener.isEmpty()) {
+            Awaitility.await().untilAsserted(() -> assertEquals(valuesInListener.size(), 2));
+            assertEquals(valuesInListener.stream().sorted().toList(), Arrays.asList("1-1", "2-1"));
+        } // else: consumer3 has passed messages to the listener before seek, in this case we cannot assume anything
+
+        @Cleanup final var consumer4 = pulsarClient.newConsumer(Schema.STRING)
+                .topics(Arrays.asList(topic1, topic2)).subscriptionName("sub-4")
+                .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest).subscribe();
+        seekAsync.apply(consumer4).get();
+        final var valuesInReceiveAsync = new ArrayList<String>();
+        valuesInReceiveAsync.add(consumer4.receiveAsync().get().getValue());
+        valuesInReceiveAsync.add(consumer4.receiveAsync().get().getValue());
+        assertEquals(valuesInReceiveAsync.stream().sorted().toList(), Arrays.asList("1-1", "2-1"));
+    }
 }

pulsar-client/src/main/java/org/apache/pulsar/client/impl/MultiTopicsConsumerImpl.java

@@ -49,6 +49,7 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import java.util.stream.IntStream;
+import javax.annotation.Nullable;
 import org.apache.commons.lang3.tuple.Pair;
 import org.apache.pulsar.client.api.BatchReceivePolicy;
 import org.apache.pulsar.client.api.Consumer;
@@ -101,7 +102,8 @@ public class MultiTopicsConsumerImpl<T> extends ConsumerBase<T> {
     private final MultiTopicConsumerStatsRecorderImpl stats;
     private final ConsumerConfigurationData<T> internalConfig;
 
-    private volatile MessageIdAdv startMessageId;
+    private final MessageIdAdv startMessageId;
+    private volatile boolean duringSeek = false;
     private final long startMessageRollbackDurationInSec;
     MultiTopicsConsumerImpl(PulsarClientImpl client, ConsumerConfigurationData<T> conf,
             ExecutorProvider executorProvider, CompletableFuture<Consumer<T>> subscribeFuture, Schema<T> schema,
@@ -235,6 +237,10 @@ private void startReceivingMessages(List<ConsumerImpl<T>> newConsumers) {
     }
 
     private void receiveMessageFromConsumer(ConsumerImpl<T> consumer, boolean batchReceive) {
+        if (duringSeek) {
+            log.info("[{}] Pause receiving messages for topic {} due to seek", subscription, consumer.getTopic());
+            return;
+        }
         CompletableFuture<List<Message<T>>> messagesFuture;
         if (batchReceive) {
             messagesFuture = consumer.batchReceiveAsync().thenApply(msgs -> ((MessagesImpl<T>) msgs).getMessageList());
@@ -252,7 +258,7 @@ private void receiveMessageFromConsumer(ConsumerImpl<T> consumer, boolean batchR
             }
             // Process the message, add to the queue and trigger listener or async callback
             messages.forEach(msg -> {
-                if (isValidConsumerEpoch((MessageImpl<T>) msg)) {
+                if (isValidConsumerEpoch((MessageImpl<T>) msg) && !duringSeek) {
                     messageReceived(consumer, msg);
                 }
             });
@@ -748,12 +754,7 @@ public void seek(Function<String, Object> function) throws PulsarClientException
 
     @Override
     public CompletableFuture<Void> seekAsync(Function<String, Object> function) {
-        List<CompletableFuture<Void>> futures = new ArrayList<>(consumers.size());
-        consumers.values().forEach(consumer -> futures.add(consumer.seekAsync(function)));
-        unAckedMessageTracker.clear();
-        incomingMessages.clear();
-        resetIncomingMessageSize();
-        return FutureUtil.waitForAll(futures);
+        return seekAllAsync(consumer -> consumer.seekAsync(function));
     }
 
     @Override
@@ -775,25 +776,44 @@ public CompletableFuture<Void> seekAsync(MessageId messageId) {
             );
         }
 
-        final CompletableFuture<Void> seekFuture;
         if (internalConsumer == null) {
-            List<CompletableFuture<Void>> futures = new ArrayList<>(consumers.size());
-            consumers.values().forEach(consumerImpl -> futures.add(consumerImpl.seekAsync(messageId)));
-            seekFuture = FutureUtil.waitForAll(futures);
+            return seekAllAsync(consumer -> consumer.seekAsync(messageId));
         } else {
-            seekFuture = internalConsumer.seekAsync(messageId);
+            beforeSeek();
+            final CompletableFuture<Void> future = new CompletableFuture<>();
+            internalConsumer.seekAsync(messageId).whenComplete((__, e) -> afterSeek(future, e));
+            return future;
         }
+    }
 
+    @Override
+    public CompletableFuture<Void> seekAsync(long timestamp) {
+        return seekAllAsync(consumer -> consumer.seekAsync(timestamp));
+    }
+
+    private CompletableFuture<Void> seekAllAsync(Function<ConsumerImpl<T>, CompletableFuture<Void>> seekFunc) {
+        beforeSeek();
+        final CompletableFuture<Void> future = new CompletableFuture<>();
+        FutureUtil.waitForAll(consumers.values().stream().map(seekFunc).collect(Collectors.toList()))
+                .whenComplete((__, e) -> afterSeek(future, e));
+        return future;
+    }
+
+    private void beforeSeek() {
+        duringSeek = true;
         unAckedMessageTracker.clear();
         clearIncomingMessages();
-        return seekFuture;
     }
 
-    @Override
-    public CompletableFuture<Void> seekAsync(long timestamp) {
-        List<CompletableFuture<Void>> futures = new ArrayList<>(consumers.size());
-        consumers.values().forEach(consumer -> futures.add(consumer.seekAsync(timestamp)));
-        return FutureUtil.waitForAll(futures);
+    private void afterSeek(CompletableFuture<Void> seekFuture, @Nullable Throwable throwable) {
+        duringSeek = false;
+        log.info("[{}] Resume receiving messages for {} since seek is done", subscription, consumers.keySet());
+        startReceivingMessages(new ArrayList<>(consumers.values()));
+        if (throwable == null) {
+            seekFuture.complete(null);
+        } else {
+            seekFuture.completeExceptionally(throwable);
+        }
     }
 
     @Override