KeyShared stickyHashRange subscription: prevent message loss in case of consumer restart

(cherry picked from commit 202ad3d)

Author: eolivelli

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentStickyKeyDispatcherMultipleConsumers.java

@@ -181,6 +181,7 @@ protected void sendMessagesToConsumers(ReadType readType, List<Entry> entries) {
                 groupedEntries.computeIfAbsent(c, k -> new ArrayList<>()).add(entry);
                 consumerStickyKeyHashesMap.computeIfAbsent(c, k -> new HashSet<>()).add(stickyKeyHash);
             } else {
+                addMessageToReplay(entry.getLedgerId(), entry.getEntryId(), stickyKeyHash);
                 entry.release();
             }
         }

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

@@ -40,6 +40,8 @@
 import java.util.Set;
 import java.util.UUID;
 import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.ConcurrentSkipListSet;
+import java.util.concurrent.CountDownLatch;
 import java.util.concurrent.ExecutionException;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
@@ -1303,4 +1305,131 @@ public EncryptionKeyInfo getPrivateKey(String keyName, Map<String, String> keyMe
             return null;
         }
     }
+
+    @Test
+    public void testStickyKeyRangesRestartConsumers() throws PulsarClientException, InterruptedException {
+        final String topic = TopicName.get("persistent", "public", "default",
+                "testStickyKeyRangesRestartConsumers" + UUID.randomUUID()).toString();
+
+        final String subscriptionName = "my-sub";
+
+        final int numMessages = 100;
+        // start 2 consumers
+        Set<String> sentMessages = new ConcurrentSkipListSet<>();
+
+        CountDownLatch count1 = new CountDownLatch(2);
+        CountDownLatch count2 = new CountDownLatch(13); // consumer 2 usually receive the fix messages
+        CountDownLatch count3 = new CountDownLatch(numMessages);
+        Consumer<String> consumer1 = pulsarClient.newConsumer(
+                        Schema.STRING)
+                .topic(topic)
+                .subscriptionName(subscriptionName)
+                .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
+                .subscriptionType(SubscriptionType.Key_Shared)
+                .keySharedPolicy(KeySharedPolicy.stickyHashRange().ranges(Range.of(0, 65536 / 2)))
+                .messageListener((consumer, msg) -> {
+                    consumer.acknowledgeAsync(msg).whenComplete((m, e) -> {
+                        if (e != null) {
+                            log.error("error", e);
+                        } else {
+                            sentMessages.remove(msg.getKey());
+                            count1.countDown();
+                            count3.countDown();
+                        }
+                    });
+                })
+                .subscribe();
+
+        Consumer<String> consumer2 = pulsarClient.newConsumer(
+                        Schema.STRING)
+                .topic(topic)
+                .subscriptionName(subscriptionName)
+                .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
+                .subscriptionType(SubscriptionType.Key_Shared)
+                .keySharedPolicy(KeySharedPolicy.stickyHashRange().ranges(Range.of(65536 / 2 + 1, 65535)))
+                .messageListener((consumer, msg) -> {
+                    consumer.acknowledgeAsync(msg).whenComplete((m, e) -> {
+                        if (e != null) {
+                            log.error("error", e);
+                        } else {
+                            sentMessages.remove(msg.getKey());
+                            count2.countDown();
+                            count3.countDown();
+                        }
+                    });
+                })
+                .subscribe();
+
+        pulsar.getExecutor().submit(() -> {
+            try
+            {
+                try (Producer<String> producer = pulsarClient.newProducer(Schema.STRING)
+                        .topic(topic)
+                        .enableBatching(false)
+                        .create();) {
+                    for (int i = 0; i < numMessages; i++)
+                    {
+                        String key = "test" + i;
+                        sentMessages.add(key);
+                        producer.newMessage()
+                                .key(key)
+                                .value("test" + i).
+                                send();
+                        Thread.sleep(100);
+                    }
+                }
+            } catch (Throwable t) {
+                log.error("error", t);
+            }});
+
+        // wait for some messages to be received by both of the consumers
+        count1.await();
+        count2.await();
+        consumer1.close();
+        consumer2.close();
+
+        // this sleep is to trigger a race condition that happens
+        // when there are some messages that cannot be dispatched while consuming
+        Thread.sleep(3000);
+
+        // start consuming again...
+
+        pulsarClient.newConsumer(Schema.STRING)
+                .topic(topic)
+                .subscriptionName(subscriptionName)
+                .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
+                .subscriptionType(SubscriptionType.Key_Shared)
+                .keySharedPolicy(KeySharedPolicy.stickyHashRange().ranges(Range.of(0, 65536 / 2)))
+                .messageListener((consumer, msg) -> {
+                    consumer.acknowledgeAsync(msg).whenComplete((m, e) -> {
+                        if (e != null) {
+                            log.error("error", e);
+                        } else {
+                            sentMessages.remove(msg.getKey());
+                            count3.countDown();
+                        }
+                    });
+                })
+                .subscribe();
+        pulsarClient.newConsumer(Schema.STRING)
+                .topic(topic)
+                .subscriptionName(subscriptionName)
+                .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest)
+                .subscriptionType(SubscriptionType.Key_Shared)
+                .keySharedPolicy(KeySharedPolicy.stickyHashRange().ranges(Range.of(65536 / 2 + 1, 65535)))
+                .messageListener((consumer, msg) -> {
+                    consumer.acknowledgeAsync(msg).whenComplete((m, e) -> {
+                        if (e != null) {
+                            log.error("error", e);
+                        } else {
+                            sentMessages.remove(msg.getKey());
+                            count3.countDown();
+                        }
+                    });
+                })
+                .subscribe();
+        // wait for all the messages to be delivered
+        count3.await();
+        assertTrue(sentMessages.isEmpty(), "didn't receive " + sentMessages);
+    }
 }