[improve][broker][PIP-286] Make the TopicCompactionService to support find entry based on publishTime or index

pulsar-broker/src/main/java/org/apache/pulsar/compaction/PulsarTopicCompactionService.java

@@ -22,18 +22,23 @@
 import static org.apache.pulsar.compaction.CompactedTopicImpl.COMPACT_LEDGER_EMPTY;
 import static org.apache.pulsar.compaction.CompactedTopicImpl.NEWER_THAN_COMPACTED;
 import static org.apache.pulsar.compaction.CompactedTopicImpl.findStartPoint;
+import static org.apache.pulsar.compaction.CompactedTopicImpl.readEntries;
 import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.Objects;
 import java.util.concurrent.CompletableFuture;
+import java.util.function.Predicate;
 import java.util.function.Supplier;
 import javax.annotation.Nonnull;
 import org.apache.bookkeeper.client.BookKeeper;
+import org.apache.bookkeeper.client.LedgerHandle;
 import org.apache.bookkeeper.mledger.Entry;
 import org.apache.bookkeeper.mledger.Position;
 import org.apache.bookkeeper.mledger.impl.PositionImpl;
+import org.apache.pulsar.common.api.proto.BrokerEntryMetadata;
+import org.apache.pulsar.common.protocol.Commands;
 import org.apache.pulsar.common.util.FutureUtil;
 
 
@@ -106,6 +111,76 @@ public CompletableFuture<Position> getLastCompactedPosition() {
         return CompletableFuture.completedFuture(compactedTopic.getCompactionHorizon().orElse(null));
     }
 
+    @Override
+    public CompletableFuture<Entry> findEntryByPublishTime(long publishTime) {
+        final Predicate<Entry> predicate = entry -> {
+            return Commands.parseMessageMetadata(entry.getDataBuffer()).getPublishTime() >= publishTime;
+        };
+        return findFirstMatchEntry(predicate);
+    }
+
+    @Override
+    public CompletableFuture<Entry> findEntryByEntryIndex(long entryIndex) {
+        final Predicate<Entry> predicate = entry -> {
+            BrokerEntryMetadata brokerEntryMetadata = Commands.parseBrokerEntryMetadataIfExist(entry.getDataBuffer());
+            if (brokerEntryMetadata == null || !brokerEntryMetadata.hasIndex()) {
+                return false;
+            }
+            return brokerEntryMetadata.getIndex() >= entryIndex;
+        };
+        return findFirstMatchEntry(predicate);
+    }
+
+    private CompletableFuture<Entry> findFirstMatchEntry(final Predicate<Entry> predicate) {
+        var compactedTopicContextFuture = compactedTopic.getCompactedTopicContextFuture();
+
+        if (compactedTopicContextFuture == null) {
+            return CompletableFuture.completedFuture(null);
+        }
+        return compactedTopicContextFuture.thenCompose(compactedTopicContext -> {
+            LedgerHandle lh = compactedTopicContext.getLedger();
+            CompletableFuture<Long> promise = new CompletableFuture<>();
+            findFirstMatchIndexLoop(predicate, 0L, lh.getLastAddConfirmed(), promise, null, lh);
+            return promise.thenCompose(index -> {
+                if (index == null) {
+                    return CompletableFuture.completedFuture(null);
+                }
+                return readEntries(lh, index, index).thenApply(entries -> entries.get(0));
+            });
+        });
+    }
+
+    private static void findFirstMatchIndexLoop(final Predicate<Entry> predicate,
+                                           final long start, final long end,
+                                           final CompletableFuture<Long> promise,
+                                           final Long lastMatchIndex,
+                                           final LedgerHandle lh) {
+        if (start > end) {
+            promise.complete(lastMatchIndex);
+            return;
+        }
+
+        long mid = (start + end) / 2;
+        readEntries(lh, mid, mid).thenAccept(entries -> {
+            Entry entry = entries.get(0);
+            final boolean isMatch;
+            try {
+                isMatch = predicate.test(entry);
+            } finally {
+                entry.release();
+            }
+
+            if (isMatch) {
+                findFirstMatchIndexLoop(predicate, start, mid - 1, promise, mid, lh);
+            } else {
+                findFirstMatchIndexLoop(predicate, mid + 1, end, promise, lastMatchIndex, lh);
+            }
+        }).exceptionally(ex -> {
+            promise.completeExceptionally(ex);
+            return null;
+        });
+    }
+
     public CompactedTopicImpl getCompactedTopic() {
         return compactedTopic;
     }

pulsar-broker/src/main/java/org/apache/pulsar/compaction/TopicCompactionService.java

@@ -60,4 +60,20 @@ public interface TopicCompactionService extends AutoCloseable {
      * @return a future that will be completed with the last compacted position, this position can be null.
      */
     CompletableFuture<Position> getLastCompactedPosition();
+
+    /**
+    * Find the first entry that greater or equal to target publishTime.
+    *
+    * @param publishTime  the publish time of entry.
+    * @return the first entry metadata that greater or equal to target publishTime, this entry can be null.
+    */
+    CompletableFuture<Entry> findEntryByPublishTime(long publishTime);
+
+    /**
+    * Find the first entry that greater or equal to target entryIndex.
+    *
+    * @param entryIndex  the index of entry.
+    * @return the first entry that greater or equal to target entryIndex, this entry can be null.
+    */
+    CompletableFuture<Entry> findEntryByEntryIndex(long entryIndex);
 }

pulsar-broker/src/test/java/org/apache/pulsar/compaction/TopicCompactionServiceTest.java

@@ -21,6 +21,7 @@
 import static org.apache.pulsar.compaction.Compactor.COMPACTED_TOPIC_LEDGER_PROPERTY;
 import static org.apache.pulsar.compaction.Compactor.COMPACTION_SUBSCRIPTION;
 import static org.testng.Assert.assertEquals;
+import static org.testng.AssertJUnit.assertNotNull;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import java.io.IOException;
 import java.util.List;
@@ -40,8 +41,11 @@
 import org.apache.pulsar.client.api.Producer;
 import org.apache.pulsar.client.api.PulsarClientException;
 import org.apache.pulsar.client.impl.MessageImpl;
+import org.apache.pulsar.common.api.proto.BrokerEntryMetadata;
+import org.apache.pulsar.common.api.proto.MessageMetadata;
 import org.apache.pulsar.common.policies.data.ClusterData;
 import org.apache.pulsar.common.policies.data.TenantInfoImpl;
+import org.apache.pulsar.common.protocol.Commands;
 import org.testng.annotations.AfterMethod;
 import org.testng.annotations.BeforeMethod;
 import org.testng.annotations.Test;
@@ -55,6 +59,11 @@ public class TopicCompactionServiceTest extends MockedPulsarServiceBaseTest {
     @BeforeMethod
     @Override
     public void setup() throws Exception {
+        conf.setBrokerEntryMetadataInterceptors(org.assertj.core.util.Sets.newTreeSet(
+                "org.apache.pulsar.common.intercept.AppendIndexMetadataInterceptor"
+        ));
+        conf.setExposingBrokerEntryMetadataToClientEnabled(true);
+
         super.internalSetup();
 
         admin.clusters().createCluster("test", ClusterData.builder().serviceUrl(pulsar.getWebServiceAddress()).build());
@@ -93,6 +102,8 @@ public void test() throws PulsarClientException, PulsarAdminException {
                 .messageRoutingMode(MessageRoutingMode.SinglePartition)
                 .create();
 
+        long startTime = System.currentTimeMillis();
+
         producer.newMessage()
                 .key("a")
                 .value("A_1".getBytes())
@@ -151,5 +162,21 @@ public void test() throws PulsarClientException, PulsarAdminException {
 
         List<Entry> entries2 = service.readCompactedEntries(PositionImpl.EARLIEST, 1).join();
         assertEquals(entries2.size(), 1);
+
+        Entry entry = service.findEntryByEntryIndex(3).join();
+        BrokerEntryMetadata brokerEntryMetadata = Commands.peekBrokerEntryMetadataIfExist(entry.getDataBuffer());
+        assertNotNull(brokerEntryMetadata);
+        assertEquals(brokerEntryMetadata.getIndex(), 4);
+        MessageMetadata metadata = Commands.parseMessageMetadata(entry.getDataBuffer());
+        assertEquals(metadata.getPartitionKey(), "b");
+        entry.release();
+
+        entry = service.findEntryByPublishTime(startTime).join();
+        brokerEntryMetadata = Commands.peekBrokerEntryMetadataIfExist(entry.getDataBuffer());
+        assertNotNull(brokerEntryMetadata);
+        assertEquals(brokerEntryMetadata.getIndex(), 2);
+        metadata = Commands.parseMessageMetadata(entry.getDataBuffer());
+        assertEquals(metadata.getPartitionKey(), "a");
+        entry.release();
     }
 }