diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/InMemoryDelayedDeliveryTracker.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/InMemoryDelayedDeliveryTracker.java index b1e50700c72f8..cf674c3b0520f 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/InMemoryDelayedDeliveryTracker.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/InMemoryDelayedDeliveryTracker.java @@ -72,6 +72,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, @@ -144,7 +147,11 @@ public boolean addMessage(long ledgerId, long entryId, long deliverAt) { LongBitmap bitmap = delayedMessageMap.computeIfAbsent(timestamp, k -> new Long2ObjectRBTreeMap<>()) .computeIfAbsent(ledgerId, k -> LongBitmaps.create()); + + long oldSize = bitmap.serializedSize(); if (bitmap.checkedAdd(entryId)) { + long newSize = bitmap.serializedSize(); + memoryUsage.addAndGet(newSize - oldSize); delayedMessagesCount.incrementAndGet(); } @@ -222,9 +229,12 @@ public NavigableSet getScheduledMessages(int maxMessages) { long ledgerId = ledgerEntry.getLongKey(); LongBitmap entryIds = ledgerEntry.getValue(); long cardinality = entryIds.cardinality(); + long oldSize = entryIds.serializedSize(); long drained = entryIds.drainTo(n, entryId -> { positions.add(PositionFactory.create(ledgerId, entryId)); }); + long newSize = entryIds.serializedSize(); + memoryUsage.addAndGet(newSize - oldSize); delayedMessagesCount.addAndGet(-drained); n -= drained; if (drained == cardinality) { @@ -264,6 +274,7 @@ public NavigableSet getScheduledMessages(int maxMessages) { public CompletableFuture clear() { this.delayedMessageMap.clear(); this.delayedMessagesCount.set(0); + this.memoryUsage.set(0); return CompletableFuture.completedFuture(null); } @@ -279,8 +290,7 @@ public long getNumberOfDelayedMessages() { */ @Override public long getBufferMemoryUsage() { - return delayedMessageMap.values().stream().mapToLong( - ledgerMap -> ledgerMap.values().stream().mapToLong(LongBitmap::serializedSize).sum()).sum(); + return memoryUsage.get(); } @Override diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedDeliveryTracker.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedDeliveryTracker.java index b8a8eb7f0f034..964e84e16b40c 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedDeliveryTracker.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedDeliveryTracker.java @@ -32,7 +32,6 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; -import java.util.Iterator; import java.util.List; import java.util.Map; import java.util.NavigableSet; @@ -111,6 +110,14 @@ public static record SnapshotKey(long ledgerId, long entryId) {} @VisibleForTesting private final RangeMap immutableBuckets; + @Getter + @VisibleForTesting + private final AtomicLong bucketsCount = new AtomicLong(0); + + @Getter + @VisibleForTesting + private final AtomicLong totalSnapshotLengthBytes = new AtomicLong(0); + private final ConcurrentHashMap snapshotSegmentLastIndexMap; private final BucketDelayedMessageIndexStats stats; @@ -245,15 +252,19 @@ private synchronized long recoverBucketSnapshot() throws RecoverDelayedDeliveryT for (Map.Entry, ImmutableBucket> mapEntry : toBeDeletedBucketMap.entrySet()) { Range key = mapEntry.getKey(); ImmutableBucket immutableBucket = mapEntry.getValue(); - immutableBucketMap.remove(key); + removeBucket(key); // delete asynchronously without waiting for completion immutableBucket.asyncDeleteBucketSnapshot(stats); } MutableLong numberDelayedMessages = new MutableLong(0); - immutableBucketMap.values().forEach(bucket -> { + long totalLength = 0; + for (ImmutableBucket bucket : immutableBucketMap.values()) { numberDelayedMessages.add(bucket.numberBucketDelayedMessages); - }); + totalLength += bucket.getSnapshotLength(); + } + totalSnapshotLengthBytes.set(totalLength); + bucketsCount.set(immutableBuckets.asMapOfRanges().size()); log.info() .attr("buckets", immutableBucketMap.size()) @@ -292,12 +303,16 @@ private CompletableFuture> handleRecoverBucketSnapshotEntry(I private synchronized void putAndCleanOverlapRange(Range range, ImmutableBucket immutableBucket, Map, ImmutableBucket> toBeDeletedBucketMap) { - RangeMap subRangeMap = immutableBuckets.subRangeMap(range); + Map, ImmutableBucket> subRangeMap = immutableBuckets.subRangeMap(range).asMapOfRanges(); boolean canPut = false; - if (!subRangeMap.asMapOfRanges().isEmpty()) { - for (Map.Entry, ImmutableBucket> rangeEntry : subRangeMap.asMapOfRanges().entrySet()) { - if (range.encloses(rangeEntry.getKey())) { - toBeDeletedBucketMap.put(rangeEntry.getKey(), rangeEntry.getValue()); + if (!subRangeMap.isEmpty()) { + for (Map.Entry, ImmutableBucket> rangeEntry : subRangeMap.entrySet()) { + // Use original key instead of truncated key for encloses check + ImmutableBucket bucket = rangeEntry.getValue(); + Range originalKey = Range.closed(bucket.startLedgerId, bucket.endLedgerId); + + if (range.encloses(originalKey)) { + toBeDeletedBucketMap.put(originalKey, bucket); canPut = true; } } @@ -306,7 +321,7 @@ private synchronized void putAndCleanOverlapRange(Range range, ImmutableBu } if (canPut) { - immutableBuckets.put(range, immutableBucket); + putBucket(range, immutableBucket); } } @@ -333,7 +348,7 @@ private void afterCreateImmutableBucket(Pair immu long startTime) { if (immutableBucketDelayedIndexPair != null) { ImmutableBucket immutableBucket = immutableBucketDelayedIndexPair.getLeft(); - immutableBuckets.put(Range.closed(immutableBucket.startLedgerId, immutableBucket.endLedgerId), + putBucket(Range.closed(immutableBucket.startLedgerId, immutableBucket.endLedgerId), immutableBucket); DelayedIndex lastDelayedIndex = immutableBucketDelayedIndexPair.getRight(); @@ -345,7 +360,12 @@ private void afterCreateImmutableBucket(Pair immu CompletableFuture future = createFuture.handle((bucketId, ex) -> { if (ex == null) { immutableBucket.setSnapshotSegments(null); - immutableBucket.asyncUpdateSnapshotLength(); + immutableBucket.asyncUpdateSnapshotLength() + .thenAccept(newLength -> { + synchronized (BucketDelayedDeliveryTracker.this) { + updateBucketSnapshotLength(immutableBucket, newLength); + } + }); log.info() .attr("bucketKey", immutableBucket.bucketKey()) .log("Create bucket snapshot finish, bucketKey"); @@ -375,7 +395,7 @@ private void afterCreateImmutableBucket(Pair immu }); immutableBucket.setCurrentSegmentEntryId(immutableBucket.lastSegmentEntryId); - immutableBuckets.asMapOfRanges().remove( + removeBucket( Range.closed(immutableBucket.startLedgerId, immutableBucket.endLedgerId)); snapshotSegmentLastIndexMap.remove( new SnapshotKey(lastDelayedIndex.getLedgerId(), lastDelayedIndex.getEntryId())); @@ -413,7 +433,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 && bucketsCount.get() > maxNumBuckets && (trimFuture == null || trimFuture.isDone())) { trimFuture = asyncTrimImmutableBuckets() .thenCompose(ignore -> asyncMergeBucketSnapshot()) @@ -483,11 +503,10 @@ private synchronized List selectMergedBuckets(final List asyncMergeBucketSnapshot() { - List immutableBucketList = immutableBuckets.asMapOfRanges().values().stream().toList(); - if (maxNumBuckets <= 0 || immutableBucketList.size() <= maxNumBuckets) { + if (maxNumBuckets <= 0 || bucketsCount.get() <= maxNumBuckets) { return CompletableFuture.completedFuture(null); } - + List immutableBucketList = immutableBuckets.asMapOfRanges().values().stream().toList(); List toBeMergeImmutableBuckets = selectMergedBuckets(immutableBucketList, MAX_MERGE_NUM); if (toBeMergeImmutableBuckets.isEmpty()) { @@ -522,7 +541,7 @@ private synchronized CompletableFuture asyncMergeBucketSnapshot() { } else { log.info() .attr("bucketKeys", bucketsStr) - .attr("bucketNum", immutableBuckets.asMapOfRanges().size()) + .attr("bucketNum", bucketsCount.get()) .log("Merge bucket snapshot finish"); stats.recordSuccessEvent(BucketDelayedMessageIndexStats.Type.merge, @@ -590,8 +609,7 @@ private synchronized CompletableFuture asyncMergeBucketSnapshot(List getScheduledMessages(int maxMessages) synchronized (BucketDelayedDeliveryTracker.this) { this.snapshotSegmentLastIndexMap.remove(snapshotKey); if (CollectionUtils.isEmpty(indexList)) { - immutableBuckets.asMapOfRanges() - .remove(Range.closed(bucket.startLedgerId, bucket.endLedgerId)); + removeBucket(Range.closed(bucket.startLedgerId, bucket.endLedgerId)); bucket.asyncDeleteBucketSnapshot(stats); return; } @@ -820,14 +837,16 @@ public CompletableFuture closeAsync() { } private CompletableFuture cleanImmutableBuckets() { + Map, ImmutableBucket> bucketsToDelete = + new HashMap<>(immutableBuckets.asMapOfRanges()); + List> futures = new ArrayList<>(); - Iterator iterator = immutableBuckets.asMapOfRanges().values().iterator(); - while (iterator.hasNext()) { - ImmutableBucket bucket = iterator.next(); - futures.add(bucket.clear(stats)); + bucketsToDelete.forEach((range, bucket) -> { + removeBucket(range); numberDelayedMessages.addAndGet(-bucket.getNumberBucketDelayedMessages()); - iterator.remove(); - } + futures.add(bucket.clear(stats)); + }); + return FutureUtil.waitForAll(futures); } @@ -850,13 +869,9 @@ public synchronized boolean containsMessage(long ledgerId, long entryId) { } public Map genTopicMetricMap() { - stats.recordNumOfBuckets(immutableBuckets.asMapOfRanges().size() + 1); + stats.recordNumOfBuckets((int) (bucketsCount.get() + 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(totalSnapshotLengthBytes.get()); return stats.genTopicMetricMap(); } @@ -900,7 +915,7 @@ private CompletableFuture deleteBucketSnapshot(String ledgerName, } synchronized (this) { snapshotSegmentLastIndexMap.entrySet().removeIf(entry -> entry.getValue() == bucket); - immutableBuckets.remove(range); + removeBucket(range); numberDelayedMessages.addAndGet(-bucket.getNumberBucketDelayedMessages()); } return null; @@ -912,4 +927,35 @@ private Long firstActiveLedgerId() { Position mdp = cursor.getMarkDeletedPosition(); return mdp == null ? null : mdp.getLedgerId(); } + + private void putBucket(Range range, ImmutableBucket bucket) { + long removedLength = immutableBuckets.subRangeMap(range).asMapOfRanges().values().stream() + .mapToLong(ImmutableBucket::getSnapshotLength) + .sum(); + + immutableBuckets.put(range, bucket); + bucketsCount.set(immutableBuckets.asMapOfRanges().size()); + totalSnapshotLengthBytes.addAndGet(bucket.getSnapshotLength() - removedLength); + } + + private void removeBucket(Range range) { + // Use exact key matching - all callers should provide exact keys + ImmutableBucket bucket = immutableBuckets.asMapOfRanges().get(range); + + if (bucket != null) { + // Remove even if snapshot length is 0 (for newly created buckets) + immutableBuckets.asMapOfRanges().remove(range); + bucketsCount.set(immutableBuckets.asMapOfRanges().size()); + totalSnapshotLengthBytes.addAndGet(-bucket.getSnapshotLength()); + } + } + + private void updateBucketSnapshotLength(ImmutableBucket bucket, long newLength) { + if (!immutableBuckets.asMapOfRanges().containsValue(bucket)) { + return; + } + long oldLength = bucket.getSnapshotLength(); + bucket.setSnapshotLength(newLength); + totalSnapshotLengthBytes.addAndGet(newLength - oldLength); + } } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedMessageIndexStats.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedMessageIndexStats.java index b9e6c7dc64c7c..502ca9f3d4edd 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedMessageIndexStats.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedMessageIndexStats.java @@ -59,7 +59,7 @@ enum Type { public BucketDelayedMessageIndexStats() { } - public Map genTopicMetricMap() { + public synchronized Map genTopicMetricMap() { Map metrics = new HashMap<>(); metrics.put(BUCKET_TOTAL_NAME, diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/ImmutableBucket.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/ImmutableBucket.java index 2124c6e9540d5..544cedce497ae 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/ImmutableBucket.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/delayed/bucket/ImmutableBucket.java @@ -131,9 +131,9 @@ private CompletableFuture> asyncLoadNextBucketSnapshotEntry(b .log("Failed to get bucket snapshot segment"); } }), BucketSnapshotPersistenceException.class, MaxRetryTimes) - .thenApply(bucketSnapshotSegments -> { + .thenCompose(bucketSnapshotSegments -> { if (CollectionUtils.isEmpty(bucketSnapshotSegments)) { - return Collections.emptyList(); + return CompletableFuture.completedFuture(Collections.emptyList()); } SnapshotSegment snapshotSegment = @@ -141,9 +141,11 @@ private CompletableFuture> asyncLoadNextBucketSnapshotEntry(b List indexList = snapshotSegment.getIndexesList(); this.setCurrentSegmentEntryId(nextSegmentEntryId); if (isRecover) { - this.asyncUpdateSnapshotLength(); + return this.asyncUpdateSnapshotLength() + .thenAccept(this::setSnapshotLength) + .thenApply(__ -> indexList); } - return indexList; + return CompletableFuture.completedFuture(indexList); }); }); } @@ -245,8 +247,6 @@ protected CompletableFuture asyncUpdateSnapshotLength() { .attr("bucketKey", bucketKey()) .exception(ex) .log("Failed to get snapshot length"); - } else { - setSnapshotLength(length); } }); } diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java index 388536907b9aa..b659f6e2200d8 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumers.java @@ -94,7 +94,7 @@ public class PersistentDispatcherMultipleConsumers extends AbstractPersistentDis protected final MessageRedeliveryController redeliveryMessages; protected final RedeliveryTracker redeliveryTracker; - private Optional delayedDeliveryTracker = Optional.empty(); + private volatile Optional delayedDeliveryTracker = Optional.empty(); protected volatile boolean havePendingRead = false; protected volatile boolean havePendingReplayRead = false; @@ -1374,13 +1374,12 @@ protected boolean isNormalReadAllowed() { } - - protected synchronized boolean shouldPauseDeliveryForDelayTracker() { - return delayedDeliveryTracker.isPresent() && delayedDeliveryTracker.get().shouldPauseAllDeliveries(); + protected boolean shouldPauseDeliveryForDelayTracker() { + return delayedDeliveryTracker.map(DelayedDeliveryTracker::shouldPauseAllDeliveries).orElse(false); } @Override - public synchronized long getNumberOfDelayedMessages() { + public long getNumberOfDelayedMessages() { return delayedDeliveryTracker.map(DelayedDeliveryTracker::getNumberOfDelayedMessages).orElse(0L); } @@ -1466,20 +1465,15 @@ public PersistentTopic getTopic() { } - public synchronized long getDelayedTrackerMemoryUsage() { + public long getDelayedTrackerMemoryUsage() { return delayedDeliveryTracker.map(DelayedDeliveryTracker::getBufferMemoryUsage).orElse(0L); } - public synchronized Map getBucketDelayedIndexStats() { - if (delayedDeliveryTracker.isEmpty()) { - return Collections.emptyMap(); - } - - if (delayedDeliveryTracker.get() instanceof BucketDelayedDeliveryTracker) { - return ((BucketDelayedDeliveryTracker) delayedDeliveryTracker.get()).genTopicMetricMap(); - } - - return Collections.emptyMap(); + public Map getBucketDelayedIndexStats() { + return delayedDeliveryTracker + .filter(BucketDelayedDeliveryTracker.class::isInstance) + .map(tracker -> ((BucketDelayedDeliveryTracker) tracker).genTopicMetricMap()) + .orElse(Collections.emptyMap()); } @Override diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java index 02353bc0dc818..3de50042b592d 100644 --- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java +++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/persistent/PersistentDispatcherMultipleConsumersClassic.java @@ -96,7 +96,7 @@ public class PersistentDispatcherMultipleConsumersClassic extends AbstractPersis protected final MessageRedeliveryController redeliveryMessages; protected final RedeliveryTracker redeliveryTracker; - private Optional delayedDeliveryTracker = Optional.empty(); + private volatile Optional delayedDeliveryTracker = Optional.empty(); protected volatile boolean havePendingRead = false; protected volatile boolean havePendingReplayRead = false; @@ -1209,12 +1209,12 @@ protected boolean hasConsumersNeededNormalRead() { return true; } - protected synchronized boolean shouldPauseDeliveryForDelayTracker() { - return delayedDeliveryTracker.isPresent() && delayedDeliveryTracker.get().shouldPauseAllDeliveries(); + protected boolean shouldPauseDeliveryForDelayTracker() { + return delayedDeliveryTracker.map(DelayedDeliveryTracker::shouldPauseAllDeliveries).orElse(false); } @Override - public synchronized long getNumberOfDelayedMessages() { + public long getNumberOfDelayedMessages() { return delayedDeliveryTracker.map(DelayedDeliveryTracker::getNumberOfDelayedMessages).orElse(0L); } @@ -1292,21 +1292,15 @@ public PersistentTopic getTopic() { return topic; } - - public synchronized long getDelayedTrackerMemoryUsage() { + public long getDelayedTrackerMemoryUsage() { return delayedDeliveryTracker.map(DelayedDeliveryTracker::getBufferMemoryUsage).orElse(0L); } - public synchronized Map getBucketDelayedIndexStats() { - if (delayedDeliveryTracker.isEmpty()) { - return Collections.emptyMap(); - } - - if (delayedDeliveryTracker.get() instanceof BucketDelayedDeliveryTracker) { - return ((BucketDelayedDeliveryTracker) delayedDeliveryTracker.get()).genTopicMetricMap(); - } - - return Collections.emptyMap(); + public Map getBucketDelayedIndexStats() { + return delayedDeliveryTracker + .filter(BucketDelayedDeliveryTracker.class::isInstance) + .map(tracker -> ((BucketDelayedDeliveryTracker) tracker).genTopicMetricMap()) + .orElse(Collections.emptyMap()); } @Override diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/MockBucketSnapshotStorage.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/MockBucketSnapshotStorage.java index fa31c27d72b9d..28d25ca00366c 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/MockBucketSnapshotStorage.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/MockBucketSnapshotStorage.java @@ -189,6 +189,15 @@ public void clean() { } } bucketSnapshots.clear(); - executorService.shutdownNow(); + // Gracefully shutdown: allow pending tasks to complete + executorService.shutdown(); + try { + if (!executorService.awaitTermination(10, TimeUnit.SECONDS)) { + executorService.shutdownNow(); + } + } catch (InterruptedException e) { + executorService.shutdownNow(); + Thread.currentThread().interrupt(); + } } } diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedDeliveryTrackerTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedDeliveryTrackerTest.java index 7205d8e6815b2..5fd2fec277720 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedDeliveryTrackerTest.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/delayed/bucket/BucketDelayedDeliveryTrackerTest.java @@ -22,6 +22,7 @@ import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotSame; @@ -728,4 +729,296 @@ public void testGetScheduledMessagesWhenAllOrphaned() throws Exception { ts.close(); } + + /** + * Test that overlapping buckets are correctly cleaned up during recovery. + * This verifies the fix for the subRangeMap clipped key issue where + * putAndCleanOverlapRange would store clipped keys that couldn't be + * removed by exact key matching in removeBucket(). + */ + @Test + public void testOverlappingBucketsCleanupDuringRecovery() throws Exception { + // Setup mocks + AbstractPersistentDispatcherMultipleConsumers testDispatcher = + mock(AbstractPersistentDispatcherMultipleConsumers.class); + Clock testClock = mock(Clock.class); + AtomicLong testClockTime = new AtomicLong(); + when(testClock.millis()).then(x -> testClockTime.get()); + + MockBucketSnapshotStorage storage = new MockBucketSnapshotStorage(); + storage.start(); + + ManagedCursor cursor = new MockManagedCursor("test_overlap_cursor"); + doReturn(cursor).when(testDispatcher).getCursor(); + doReturn("persistent://public/default/testOverlap / " + cursor.getName()) + .when(testDispatcher).getName(); + + try { + // Create first tracker with small minIndexCountPerBucket + BucketDelayedDeliveryTracker tracker1 = new BucketDelayedDeliveryTracker( + testDispatcher, timer, 100000, testClock, true, storage, + 3, TimeUnit.MILLISECONDS.toMillis(10), -1, 50); + + // Add messages to create multiple immutable buckets + for (int i = 1; i <= 12; i++) { + tracker1.addMessage(i, i, i * 10); + } + + // Wait for all bucket operations to complete + Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> { + assertTrue(tracker1.getImmutableBuckets().asMapOfRanges().values().stream() + .noneMatch(x -> x.merging), + "All buckets should finish merging"); + assertTrue(tracker1.getImmutableBuckets().asMapOfRanges().size() >= 2, + "Should have created multiple buckets"); + }); + + int bucketCountBeforeClose = tracker1.getImmutableBuckets().asMapOfRanges().size(); + + tracker1.close(); + + // Create second tracker - triggers recovery with putAndCleanOverlapRange + BucketDelayedDeliveryTracker tracker2 = new BucketDelayedDeliveryTracker( + testDispatcher, timer, 100000, testClock, true, storage, + 3, TimeUnit.MILLISECONDS.toMillis(10), -1, 50); + + // Verify buckets were recovered + int bucketCountAfterRecovery = tracker2.getImmutableBuckets().asMapOfRanges().size(); + assertTrue(bucketCountAfterRecovery > 0, "Should have recovered buckets"); + + // Key assertion: verify no orphaned buckets remain + // If clipped keys weren't fixed, removeBucket() would fail and buckets would accumulate + assertTrue(bucketCountAfterRecovery <= bucketCountBeforeClose, + String.format("Orphaned buckets detected: %d after recovery > %d before close", + bucketCountAfterRecovery, bucketCountBeforeClose)); + + // Verify messages were recovered + assertTrue(tracker2.getNumberOfDelayedMessages() > 0, + "Should have recovered messages"); + + // Verify snapshot length tracking is correct + long totalSnapshotLength = tracker2.getImmutableBuckets().asMapOfRanges().values().stream() + .mapToLong(ImmutableBucket::getSnapshotLength) + .sum(); + assertTrue(totalSnapshotLength >= 0, + "Snapshot length tracking broken - likely due to failed removeBucket()"); + + tracker2.close(); + } finally { + storage.clean(); + } + } + + /** + * Test that putAndCleanOverlapRange correctly uses original keys instead of truncated keys + * when checking if a new range encloses existing buckets. + * + * This prevents the bug where a truncated key from subRangeMap() would incorrectly pass + * the encloses() check, causing a bucket to be replaced when it shouldn't be. + */ + @Test + public void testPutAndCleanOverlapRangeWithTruncatedKeys() throws Exception { + // Setup mocks + AbstractPersistentDispatcherMultipleConsumers testDispatcher = + mock(AbstractPersistentDispatcherMultipleConsumers.class); + Clock testClock = mock(Clock.class); + AtomicLong testClockTime = new AtomicLong(); + when(testClock.millis()).then(x -> testClockTime.get()); + + MockBucketSnapshotStorage storage = new MockBucketSnapshotStorage(); + storage.start(); + + ManagedCursor cursor = new MockManagedCursor("test_truncated_cursor"); + doReturn(cursor).when(testDispatcher).getCursor(); + doReturn("persistent://public/default/testTruncated / " + cursor.getName()) + .when(testDispatcher).getName(); + + try { + // Create tracker + BucketDelayedDeliveryTracker tracker = new BucketDelayedDeliveryTracker( + testDispatcher, timer, 100000, testClock, true, storage, + 3, TimeUnit.MILLISECONDS.toMillis(10), -1, 50); + + // Add messages to create a bucket [1-6] + for (int i = 1; i <= 6; i++) { + tracker.addMessage(i, i, i * 10); + } + + // Wait for bucket to be created + Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> { + assertTrue(tracker.getImmutableBuckets().asMapOfRanges().size() >= 1, + "Should have created at least one bucket"); + }); + + int initialBucketCount = tracker.getImmutableBuckets().asMapOfRanges().size(); + long initialMessageCount = tracker.getNumberOfDelayedMessages(); + + // Now add messages that would create a bucket [7-9] + // This should NOT replace the existing bucket [1-6] + for (int i = 7; i <= 9; i++) { + tracker.addMessage(i, i, i * 10); + } + + // Wait for new bucket operations + Awaitility.await().atMost(5, TimeUnit.SECONDS).untilAsserted(() -> { + assertTrue(tracker.getImmutableBuckets().asMapOfRanges().values().stream() + .noneMatch(x -> x.merging), + "All buckets should finish processing"); + }); + + // Verify bucket count increased (or stayed same if they got merged) + int finalBucketCount = tracker.getImmutableBuckets().asMapOfRanges().size(); + assertTrue(finalBucketCount >= initialBucketCount, + "Bucket count should not decrease when adding non-overlapping ranges"); + + // Verify all messages are tracked + long finalMessageCount = tracker.getNumberOfDelayedMessages(); + assertTrue(finalMessageCount >= initialMessageCount, + String.format("Message count should not decrease: initial=%d, final=%d", + initialMessageCount, finalMessageCount)); + + // Verify no bucket was incorrectly replaced + // If putAndCleanOverlapRange used truncated keys, it might have incorrectly + // removed a bucket that shouldn't have been removed + tracker.getImmutableBuckets().asMapOfRanges().forEach((range, bucket) -> { + assertTrue(bucket.getNumberBucketDelayedMessages() > 0, + "All buckets should have messages - bucket " + range + " is empty"); + }); + + tracker.close(); + } finally { + storage.clean(); + } + } + + @Test + public void testLateSnapshotLengthUpdateAfterClearDoesNotInflateCounter() throws Exception { + AbstractPersistentDispatcherMultipleConsumers testDispatcher = + mock(AbstractPersistentDispatcherMultipleConsumers.class); + Clock testClock = mock(Clock.class); + AtomicLong testClockTime = new AtomicLong(); + when(testClock.millis()).then(x -> testClockTime.get()); + + MockBucketSnapshotStorage storage = new MockBucketSnapshotStorage(); + storage.start(); + MockBucketSnapshotStorage spyStorage = spy(storage); + + CompletableFuture blockedLength = new CompletableFuture<>(); + when(spyStorage.getBucketSnapshotLength(anyLong())).thenReturn(blockedLength); + + ManagedCursor cursor = new MockManagedCursor("test_late_update_cursor"); + doReturn(cursor).when(testDispatcher).getCursor(); + doReturn("persistent://public/default/testLateUpdate / " + cursor.getName()) + .when(testDispatcher).getName(); + + try { + BucketDelayedDeliveryTracker tracker = new BucketDelayedDeliveryTracker( + testDispatcher, timer, 100000, testClock, true, spyStorage, + 3, TimeUnit.MILLISECONDS.toMillis(10), -1, 50); + + for (int i = 1; i <= 6; i++) { + tracker.addMessage(i, i, i * 10); + } + + Awaitility.await().untilAsserted(() -> + assertTrue(tracker.getBucketsCount().get() >= 1, + "Should have created at least one immutable bucket")); + assertCountersConsistent(tracker); + + tracker.clear(); + + assertEquals(tracker.getBucketsCount().get(), 0, "All buckets should be removed"); + assertCountersConsistent(tracker); + + blockedLength.complete(999_999L); + + Awaitility.await().untilAsserted(() -> { + assertEquals(tracker.getTotalSnapshotLengthBytes().get(), 0, + "Late length update inflated totalSnapshotLengthBytes after clear"); + }); + + tracker.close(); + } finally { + storage.clean(); + } + } + + @Test + public void testLateSnapshotLengthUpdateAfterTrimDoesNotInflateCounter() throws Exception { + AbstractPersistentDispatcherMultipleConsumers testDispatcher = + mock(AbstractPersistentDispatcherMultipleConsumers.class); + Clock testClock = mock(Clock.class); + AtomicLong testClockTime = new AtomicLong(); + when(testClock.millis()).then(x -> testClockTime.get()); + + MockBucketSnapshotStorage storage = new MockBucketSnapshotStorage(); + storage.start(); + MockBucketSnapshotStorage spyStorage = spy(storage); + + CompletableFuture blockedLength = new CompletableFuture<>(); + when(spyStorage.getBucketSnapshotLength(anyLong())).thenReturn(blockedLength); + + ManagedCursor spyCursor = spy(new MockManagedCursor("test_late_trim_cursor")); + AtomicLong markDeletedLedger = new AtomicLong(0); + when(spyCursor.getMarkDeletedPosition()).thenAnswer(inv -> + PositionFactory.create(markDeletedLedger.get(), 0)); + ManagedLedger mockLedger = mock(ManagedLedger.class); + when(mockLedger.getName()).thenReturn("test_ledger"); + when(spyCursor.getManagedLedger()).thenReturn(mockLedger); + + doReturn(spyCursor).when(testDispatcher).getCursor(); + doReturn("persistent://public/default/testLateTrim / " + spyCursor.getName()) + .when(testDispatcher).getName(); + + try { + BucketDelayedDeliveryTracker tracker = new BucketDelayedDeliveryTracker( + testDispatcher, timer, 100000, testClock, true, spyStorage, + 3, TimeUnit.MILLISECONDS.toMillis(10), -1, 3); + + for (int i = 1; i <= 12; i++) { + tracker.addMessage(i, i, i * 10); + } + + Awaitility.await().untilAsserted(() -> + assertTrue(tracker.getBucketsCount().get() >= 3, + "Should have created at least 3 immutable buckets")); + assertCountersConsistent(tracker); + + markDeletedLedger.set(5); + + for (int i = 13; i <= 15; i++) { + tracker.addMessage(i, i, i * 10); + } + + Awaitility.await().untilAsserted(() -> { + boolean hasOldBucket = tracker.getImmutableBuckets().asMapOfRanges().keySet().stream() + .anyMatch(r -> r.upperEndpoint() < 5); + Assert.assertFalse(hasOldBucket, "Buckets with endLedgerId < 5 should be trimmed"); + }); + assertCountersConsistent(tracker); + + blockedLength.complete(999_999L); + + Awaitility.await().untilAsserted(() -> + assertCountersConsistent(tracker)); + + tracker.close(); + } finally { + storage.clean(); + } + } + + private static void assertCountersConsistent(BucketDelayedDeliveryTracker tracker) { + int liveBucketCount = tracker.getImmutableBuckets().asMapOfRanges().size(); + long liveSnapshotLength = tracker.getImmutableBuckets().asMapOfRanges().values().stream() + .mapToLong(ImmutableBucket::getSnapshotLength) + .sum(); + + assertEquals(tracker.getBucketsCount().get(), liveBucketCount, + String.format("bucketsCount drift: cached=%d live=%d", + tracker.getBucketsCount().get(), liveBucketCount)); + assertEquals(tracker.getTotalSnapshotLengthBytes().get(), liveSnapshotLength, + String.format("totalSnapshotLengthBytes drift: cached=%d live=%d", + tracker.getTotalSnapshotLengthBytes().get(), liveSnapshotLength)); + } }