Address comment

Author: nodece

pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/InMemoryDelayedDeliveryTracker.java

@@ -66,6 +66,9 @@ public class InMemoryDelayedDeliveryTracker extends AbstractDelayedDeliveryTrack
     // Count of delayed messages in the tracker.
     private final AtomicLong delayedMessagesCount = new AtomicLong(0);
 
+    // Cached memory usage of the delayed message bitmaps, maintained via delta on each mutation.
+    private final AtomicLong memoryUsage = new AtomicLong(0);
+
     InMemoryDelayedDeliveryTracker(AbstractPersistentDispatcherMultipleConsumers dispatcher, Timer timer,
                                    long tickTimeMillis,
                                    boolean isDelayedDeliveryDeliverAtTimeStrict,
@@ -142,7 +145,9 @@ public boolean addMessage(long ledgerId, long entryId, long deliverAt) {
         boolean isNew = !bitmap.contains(entryId);
 
         if (isNew) {
+            long oldSize = bitmap.getLongSizeInBytes();
             bitmap.add(entryId);
+            memoryUsage.addAndGet(bitmap.getLongSizeInBytes() - oldSize);
             delayedMessagesCount.incrementAndGet();
         }
 
@@ -205,13 +210,16 @@ public NavigableSet<Position> getScheduledMessages(int maxMessages) {
                     });
                     n -= cardinality;
                     delayedMessagesCount.addAndGet(-cardinality);
+                    memoryUsage.addAndGet(-entryIds.getLongSizeInBytes());
                     ledgerIdToDelete.add(ledgerId);
                 } else {
+                    long oldSize = entryIds.getLongSizeInBytes();
                     long[] entryIdsArray = entryIds.toArray();
                     for (int i = 0; i < n; i++) {
                         positions.add(PositionFactory.create(ledgerId, entryIdsArray[i]));
                         entryIds.removeLong(entryIdsArray[i]);
                     }
+                    memoryUsage.addAndGet(entryIds.getLongSizeInBytes() - oldSize);
                     delayedMessagesCount.addAndGet(-n);
                     n = 0;
                 }
@@ -248,6 +256,7 @@ public NavigableSet<Position> getScheduledMessages(int maxMessages) {
     public CompletableFuture<Void> clear() {
         this.delayedMessageMap.clear();
         this.delayedMessagesCount.set(0);
+        this.memoryUsage.set(0);
         return CompletableFuture.completedFuture(null);
     }
 
@@ -264,9 +273,7 @@ public long getNumberOfDelayedMessages() {
      */
     @Override
     public long getBufferMemoryUsage() {
-        return delayedMessageMap.values().stream().mapToLong(
-                ledgerMap -> ledgerMap.values().stream().mapToLong(
-                        Roaring64Bitmap::getLongSizeInBytes).sum()).sum();
+        return memoryUsage.get();
     }
 
     @Override

pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedDeliveryTracker.java

@@ -24,7 +24,6 @@
 import com.google.common.annotations.VisibleForTesting;
 import com.google.common.collect.Range;
 import com.google.common.collect.RangeMap;
-import com.google.common.collect.TreeRangeMap;
 import io.github.merlimat.slog.Logger;
 import io.netty.util.Timeout;
 import io.netty.util.Timer;
@@ -107,7 +106,7 @@ public static record SnapshotKey(long ledgerId, long entryId) {}
 
     @Getter
     @VisibleForTesting
-    private final RangeMap<Long, ImmutableBucket> immutableBuckets;
+    private final ImmutableBucketIndex immutableBuckets;
 
     private final ConcurrentHashMap<SnapshotKey, ImmutableBucket> snapshotSegmentLastIndexMap;
 
@@ -154,7 +153,7 @@ public BucketDelayedDeliveryTracker(DelayedDeliveryContext context,
         this.maxIndexesPerBucketSnapshotSegment = maxIndexesPerBucketSnapshotSegment;
         this.maxNumBuckets = maxNumBuckets;
         this.sharedBucketPriorityQueue = new TripleLongPriorityQueue();
-        this.immutableBuckets = TreeRangeMap.create();
+        this.immutableBuckets = new ImmutableBucketIndex();
         this.snapshotSegmentLastIndexMap = new ConcurrentHashMap<>();
         this.lastMutableBucket =
                 new MutableBucket(context.getName(), context.getCursor(), FutureUtil.Sequencer.create(),
@@ -241,7 +240,7 @@ private synchronized long recoverBucketSnapshot() throws RecoverDelayedDeliveryT
         for (Map.Entry<Range<Long>, ImmutableBucket> mapEntry : toBeDeletedBucketMap.entrySet()) {
             Range<Long> key = mapEntry.getKey();
             ImmutableBucket immutableBucket = mapEntry.getValue();
-            immutableBucketMap.remove(key);
+            immutableBuckets.remove(key);
             // delete asynchronously without waiting for completion
             immutableBucket.asyncDeleteBucketSnapshot(stats);
         }
@@ -251,6 +250,8 @@ private synchronized long recoverBucketSnapshot() throws RecoverDelayedDeliveryT
             numberDelayedMessages.add(bucket.numberBucketDelayedMessages);
         });
 
+        immutableBuckets.recomputeCounters();
+
         log.info()
                 .attr("buckets", immutableBucketMap.size())
                 .attr("numberDelayedMessages", numberDelayedMessages.longValue())
@@ -371,7 +372,7 @@ private void afterCreateImmutableBucket(Pair<ImmutableBucket, DelayedIndex> immu
                         });
 
                         immutableBucket.setCurrentSegmentEntryId(immutableBucket.lastSegmentEntryId);
-                        immutableBuckets.asMapOfRanges().remove(
+                        immutableBuckets.remove(
                                 Range.closed(immutableBucket.startLedgerId, immutableBucket.endLedgerId));
                         snapshotSegmentLastIndexMap.remove(
                                 new SnapshotKey(lastDelayedIndex.getLedgerId(), lastDelayedIndex.getEntryId()));
@@ -409,7 +410,7 @@ public synchronized boolean addMessage(long ledgerId, long entryId, long deliver
             afterCreateImmutableBucket(immutableBucketDelayedIndexPair, createStartTime);
             lastMutableBucket.resetLastMutableBucketRange();
 
-            if (maxNumBuckets > 0 && immutableBuckets.asMapOfRanges().size() > maxNumBuckets) {
+            if (maxNumBuckets > 0 && immutableBuckets.count() > maxNumBuckets) {
                 asyncMergeBucketSnapshot();
             }
         }
@@ -507,7 +508,7 @@ private synchronized CompletableFuture<Void> asyncMergeBucketSnapshot() {
             } else {
                 log.info()
                         .attr("bucketKeys", bucketsStr)
-                        .attr("bucketNum", immutableBuckets.asMapOfRanges().size())
+                        .attr("bucketNum", immutableBuckets.count())
                         .log("Merge bucket snapshot finish");
 
                 stats.recordSuccessEvent(BucketDelayedMessageIndexStats.Type.merge,
@@ -577,8 +578,7 @@ private synchronized CompletableFuture<Void> asyncMergeBucketSnapshot(List<Immut
                                     });
 
                             for (ImmutableBucket bucket : buckets) {
-                                immutableBuckets.asMapOfRanges()
-                                        .remove(Range.closed(bucket.startLedgerId, bucket.endLedgerId));
+                                immutableBuckets.remove(Range.closed(bucket.startLedgerId, bucket.endLedgerId));
                             }
                         }
                     });
@@ -644,7 +644,7 @@ public synchronized NavigableSet<Position> getScheduledMessages(int maxMessages)
             SnapshotKey snapshotKey = new SnapshotKey(ledgerId, entryId);
 
             ImmutableBucket bucket = snapshotSegmentLastIndexMap.get(snapshotKey);
-            if (bucket != null && immutableBuckets.asMapOfRanges().containsValue(bucket)) {
+            if (bucket != null && immutableBuckets.containsValue(bucket)) {
                 // All message of current snapshot segment are scheduled, try load next snapshot segment
                 if (bucket.merging) {
                     log.info()
@@ -674,8 +674,7 @@ public synchronized NavigableSet<Position> getScheduledMessages(int maxMessages)
                     synchronized (BucketDelayedDeliveryTracker.this) {
                         this.snapshotSegmentLastIndexMap.remove(snapshotKey);
                         if (CollectionUtils.isEmpty(indexList)) {
-                            immutableBuckets.asMapOfRanges()
-                                    .remove(Range.closed(bucket.startLedgerId, bucket.endLedgerId));
+                            immutableBuckets.remove(Range.closed(bucket.startLedgerId, bucket.endLedgerId));
                             bucket.asyncDeleteBucketSnapshot(stats);
                             return;
                         }
@@ -807,13 +806,9 @@ public synchronized boolean containsMessage(long ledgerId, long entryId) {
     }
 
     public Map<String, TopicMetricBean> genTopicMetricMap() {
-        stats.recordNumOfBuckets(immutableBuckets.asMapOfRanges().size() + 1);
+        stats.recordNumOfBuckets((int) (immutableBuckets.count() + 1));
         stats.recordDelayedMessageIndexLoaded(this.sharedBucketPriorityQueue.size() + this.lastMutableBucket.size());
-        MutableLong totalSnapshotLength = new MutableLong();
-        immutableBuckets.asMapOfRanges().values().forEach(immutableBucket -> {
-            totalSnapshotLength.add(immutableBucket.getSnapshotLength());
-        });
-        stats.recordBucketSnapshotSizeBytes(totalSnapshotLength.longValue());
+        stats.recordBucketSnapshotSizeBytes(immutableBuckets.totalSnapshotLength());
         return stats.genTopicMetricMap();
     }
 }

pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/ImmutableBucketIndex.java

@@ -0,0 +1,118 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.pulsar.broker.delayed.bucket;
+
+import com.google.common.collect.Range;
+import com.google.common.collect.RangeMap;
+import com.google.common.collect.TreeRangeMap;
+import java.util.Map;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * Index of {@link ImmutableBucket}s keyed by ledger-ID {@link Range}.
+ *
+ * <p>The underlying {@link TreeRangeMap} is not thread-safe. All mutations must be
+ * performed under an external lock (for example {@code synchronized(tracker)}).
+ *
+ * <p>To support lock-free stats reads, this class maintains cached counters for:
+ * <ul>
+ *   <li>number of buckets ({@link #count()})</li>
+ *   <li>total snapshot length ({@link #totalSnapshotLength()})</li>
+ * </ul>
+ *
+ * <p>Note that {@link TreeRangeMap#put} automatically removes any existing entries
+ * whose ranges overlap with the inserted range. Because these removals are implicit,
+ * counter updates cannot be derived solely from the caller's operation. Therefore
+ * {@link #put(Range, ImmutableBucket)} and {@link #recomputeCounters()} rebuild the
+ * counters from the actual map state.
+ */
+class ImmutableBucketIndex {
+
+    private final TreeRangeMap<Long, ImmutableBucket> map = TreeRangeMap.create();
+
+    private final AtomicLong count = new AtomicLong();
+    private final AtomicLong totalSnapshotLength = new AtomicLong();
+
+    /**
+     * Adds a bucket to the index.
+     *
+     * <p>If the supplied range overlaps existing entries, {@link TreeRangeMap}
+     * removes them automatically. Counters are therefore recomputed from the
+     * resulting map state.
+     */
+    void put(Range<Long> range, ImmutableBucket bucket) {
+        map.put(range, bucket);
+        recomputeCounters();
+    }
+
+    /**
+     * Removes the bucket associated with the given range.
+     *
+     * <p>The supplied range is expected to exactly match an existing entry.
+     */
+    void remove(Range<Long> range) {
+        ImmutableBucket removed = map.asMapOfRanges().remove(range);
+        if (removed != null) {
+            count.decrementAndGet();
+            totalSnapshotLength.addAndGet(-removed.getSnapshotLength());
+        }
+    }
+
+    long count() {
+        return count.get();
+    }
+
+    long totalSnapshotLength() {
+        return totalSnapshotLength.get();
+    }
+
+    ImmutableBucket get(long key) {
+        return map.get(key);
+    }
+
+    RangeMap<Long, ImmutableBucket> subRangeMap(Range<Long> range) {
+        return map.subRangeMap(range);
+    }
+
+    Map<Range<Long>, ImmutableBucket> asMapOfRanges() {
+        return map.asMapOfRanges();
+    }
+
+    boolean containsValue(ImmutableBucket bucket) {
+        return map.asMapOfRanges().containsValue(bucket);
+    }
+
+    /**
+     * Rebuilds cached counters from the current map state.
+     *
+     * <p>This method is useful after bulk updates or whenever the map state is
+     * considered the source of truth.
+     */
+    void recomputeCounters() {
+        Map<Range<Long>, ImmutableBucket> ranges = map.asMapOfRanges();
+
+        count.set(ranges.size());
+
+        long snapshotLength = 0;
+        for (ImmutableBucket bucket : ranges.values()) {
+            snapshotLength += bucket.getSnapshotLength();
+        }
+        totalSnapshotLength.set(snapshotLength);
+    }
+}

pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/bucket/DelayedIndexQueueTest.java

@@ -152,4 +152,29 @@ public void tripleLongPriorityDelayedIndexQueueTest() {
 
         Assert.assertEquals(10, count);
     }
+
+    @Test
+    public void testTripleLongPriorityDelayedIndexQueueDrainTo() {
+        @Cleanup
+        TripleLongPriorityQueue queue = new TripleLongPriorityQueue();
+        queue.add(30, 3, 3);
+        queue.add(10, 1, 1);
+        queue.add(20, 2, 2);
+        queue.add(10, 1, 2);
+
+        TripleLongPriorityDelayedIndexQueue delayedIndexQueue =
+                TripleLongPriorityDelayedIndexQueue.wrap(queue);
+
+        List<long[]> result = new ArrayList<>();
+        delayedIndexQueue.drainTo((ts, lid, eid) -> result.add(new long[]{ts, lid, eid}));
+
+        Assert.assertTrue(delayedIndexQueue.isEmpty());
+        Assert.assertEquals(result.size(), 4);
+
+        // Verify sorted order: (10,1,1) < (10,1,2) < (20,2,2) < (30,3,3)
+        Assert.assertEquals(result.get(0), new long[]{10, 1, 1});
+        Assert.assertEquals(result.get(1), new long[]{10, 1, 2});
+        Assert.assertEquals(result.get(2), new long[]{20, 2, 2});
+        Assert.assertEquals(result.get(3), new long[]{30, 3, 3});
+    }
 }