[improve][broker] Reduce PendingAcksMap memory by replacing IntIntPair values with packed longs (#26030)
diff --git a/microbench/build.gradle.kts b/microbench/build.gradle.kts
index 9cea5eb..2164862 100644
--- a/microbench/build.gradle.kts
+++ b/microbench/build.gradle.kts
@@ -29,6 +29,7 @@
implementation(project(":pulsar-common"))
api(project(":pulsar-broker"))
implementation(libs.bookkeeper.server)
+ implementation(libs.fastutil)
api(libs.guava)
api("org.openjdk.jmh:jmh-core:1.37")
annotationProcessor("org.openjdk.jmh:jmh-generator-annprocess:1.37")
diff --git a/microbench/src/main/java/org/apache/pulsar/broker/service/PendingAcksMapPackingBenchmark.java b/microbench/src/main/java/org/apache/pulsar/broker/service/PendingAcksMapPackingBenchmark.java
new file mode 100644
index 0000000..70e711c
--- /dev/null
+++ b/microbench/src/main/java/org/apache/pulsar/broker/service/PendingAcksMapPackingBenchmark.java
@@ -0,0 +1,536 @@
+/*
+ * 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.service;
+
+import it.unimi.dsi.fastutil.ints.IntIntPair;
+import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
+import it.unimi.dsi.fastutil.longs.Long2ObjectRBTreeMap;
+import it.unimi.dsi.fastutil.longs.Long2ObjectSortedMap;
+import it.unimi.dsi.fastutil.objects.ObjectBidirectionalIterator;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.Lock;
+import java.util.concurrent.locks.ReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.function.Supplier;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@BenchmarkMode(Mode.AverageTime)
+@Fork(1)
+@Warmup(iterations = 2, time = 1, timeUnit = TimeUnit.SECONDS)
+@Measurement(iterations = 3, time = 1, timeUnit = TimeUnit.SECONDS)
+public class PendingAcksMapPackingBenchmark {
+ private static final int PENDING_ACK_NOT_FOUND = -1;
+
+ @Benchmark
+ public boolean addOrReplace(MapState state, CursorState cursor) {
+ int index = cursor.next(state.entries);
+ return state.store.addOrReplace(state.ledgerIds[index], state.entryIds[index],
+ remainingUnacked(index), stickyKeyHash(index));
+ }
+
+ @Benchmark
+ public boolean containsHit(MapState state, CursorState cursor) {
+ int index = cursor.next(state.entries);
+ return state.store.contains(state.ledgerIds[index], state.entryIds[index]);
+ }
+
+ @Benchmark
+ public long forEachScan(MapState state) {
+ return state.store.forEachScan();
+ }
+
+ @Benchmark
+ public int getRemainingUnackedHit(MapState state, CursorState cursor) {
+ int index = cursor.next(state.entries);
+ return state.store.getRemainingUnacked(state.ledgerIds[index], state.entryIds[index]);
+ }
+
+ @Benchmark
+ public long removeAllUpToBeforeFirstEntry(MapState state) {
+ return state.store.removeAllUpTo(0, -1);
+ }
+
+ @Benchmark
+ public long removeAllUpToSmallPrefixAndRefill(PrefixRemoveState state) {
+ long removed = state.store.removeAllUpTo(0, state.prefixEntries - 1L);
+ for (int i = 0; i < state.prefixEntries; i++) {
+ state.store.addOrReplace(0, i, remainingUnacked(i), stickyKeyHash(i));
+ }
+ return removed;
+ }
+
+ @Benchmark
+ public int removeAndGetRemainingAndAdd(MapState state, CursorState cursor) {
+ int index = cursor.next(state.entries);
+ long ledgerId = state.ledgerIds[index];
+ long entryId = state.entryIds[index];
+ int remainingUnacked = state.store.removeAndGetRemainingUnacked(ledgerId, entryId);
+ state.store.addOrReplace(ledgerId, entryId,
+ remainingUnacked == PENDING_ACK_NOT_FOUND ? remainingUnacked(index) : remainingUnacked,
+ stickyKeyHash(index));
+ return remainingUnacked;
+ }
+
+ @Benchmark
+ public boolean removeWithValueAndAdd(MapState state, CursorState cursor) {
+ int index = cursor.next(state.entries);
+ long ledgerId = state.ledgerIds[index];
+ long entryId = state.entryIds[index];
+ int remainingUnacked = remainingUnacked(index);
+ int stickyKeyHash = stickyKeyHash(index);
+ boolean removed = state.store.remove(ledgerId, entryId, remainingUnacked, stickyKeyHash);
+ state.store.addOrReplace(ledgerId, entryId, remainingUnacked, stickyKeyHash);
+ return removed;
+ }
+
+ @Benchmark
+ public boolean updateRemainingUnacked(MapState state, CursorState cursor) {
+ int index = cursor.next(state.entries);
+ return state.store.updateRemainingUnacked(state.ledgerIds[index], state.entryIds[index], 1);
+ }
+
+ @State(Scope.Benchmark)
+ public static class MapState {
+ @Param({"FORMER_PENDING_ACKS_MAP", "CURRENT_PENDING_ACKS_MAP"})
+ private Implementation implementation;
+
+ @Param({"50000"})
+ private int entries;
+
+ @Param({"1"})
+ private int ledgers;
+
+ private PendingAcksStore store;
+ private long[] ledgerIds;
+ private long[] entryIds;
+
+ @Setup(Level.Trial)
+ public void setup() {
+ store = implementation.createStore();
+ ledgerIds = new long[entries];
+ entryIds = new long[entries];
+ populate(store, entries, ledgers, ledgerIds, entryIds);
+ }
+ }
+
+ @State(Scope.Thread)
+ public static class PrefixRemoveState {
+ @Param({"FORMER_PENDING_ACKS_MAP", "CURRENT_PENDING_ACKS_MAP"})
+ private Implementation implementation;
+
+ @Param({"50000"})
+ private int entries;
+
+ private PendingAcksStore store;
+ private int prefixEntries;
+
+ @Setup(Level.Trial)
+ public void setup() {
+ store = implementation.createStore();
+ prefixEntries = Math.max(1, entries / 50);
+ populate(store, entries, 1, null, null);
+ }
+ }
+
+ @State(Scope.Thread)
+ public static class CursorState {
+ private int cursor;
+
+ int next(int bound) {
+ int next = cursor++;
+ if (cursor == bound) {
+ cursor = 0;
+ }
+ return next;
+ }
+ }
+
+ public enum Implementation {
+ FORMER_PENDING_ACKS_MAP {
+ @Override
+ PendingAcksStore createStore() {
+ return new FormerPendingAcksMapStore();
+ }
+ },
+ CURRENT_PENDING_ACKS_MAP {
+ @Override
+ PendingAcksStore createStore() {
+ return new CurrentPendingAcksMapStore();
+ }
+ };
+
+ abstract PendingAcksStore createStore();
+ }
+
+ private interface PendingAcksStore {
+ boolean addOrReplace(long ledgerId, long entryId, int remainingUnacked, int stickyKeyHash);
+
+ boolean contains(long ledgerId, long entryId);
+
+ int getRemainingUnacked(long ledgerId, long entryId);
+
+ boolean remove(long ledgerId, long entryId, int remainingUnacked, int stickyKeyHash);
+
+ int removeAndGetRemainingUnacked(long ledgerId, long entryId);
+
+ boolean updateRemainingUnacked(long ledgerId, long entryId, int ackedDelta);
+
+ long forEachScan();
+
+ long removeAllUpTo(long markDeleteLedgerId, long markDeleteEntryId);
+ }
+
+ private static final class FormerPendingAcksMapStore implements PendingAcksStore {
+ private final Consumer consumer = null;
+ private final Long2ObjectSortedMap<Long2ObjectSortedMap<IntIntPair>> pendingAcks =
+ new Long2ObjectRBTreeMap<>();
+ private final Supplier<PendingAcksMap.PendingAcksAddHandler> pendingAcksAddHandlerSupplier = () -> null;
+ private final Supplier<PendingAcksMap.PendingAcksRemoveHandler> pendingAcksRemoveHandlerSupplier = () -> null;
+ private final Lock readLock;
+ private final Lock writeLock;
+ private boolean closed = false;
+ private volatile long size;
+ private long forEachSum;
+ private long removedCount;
+ private final PendingAcksMap.PendingAcksConsumer pendingAckScanner =
+ (ledgerId, entryId, remainingUnacked, stickyKeyHash) ->
+ forEachSum += ledgerId + entryId + remainingUnacked + stickyKeyHash;
+ private final PendingAcksMap.PendingAcksConsumer removedCounter =
+ (ledgerId, entryId, remainingUnacked, stickyKeyHash) -> removedCount++;
+
+ private FormerPendingAcksMapStore() {
+ ReadWriteLock readWriteLock = new ReentrantReadWriteLock();
+ this.writeLock = readWriteLock.writeLock();
+ this.readLock = readWriteLock.readLock();
+ }
+
+ @Override
+ public boolean addOrReplace(long ledgerId, long entryId, int remainingUnacked, int stickyKeyHash) {
+ try {
+ writeLock.lock();
+ if (closed) {
+ return false;
+ }
+ PendingAcksMap.PendingAcksAddHandler pendingAcksAddHandler = pendingAcksAddHandlerSupplier.get();
+ if (pendingAcksAddHandler != null
+ && !pendingAcksAddHandler.handleAdding(consumer, ledgerId, entryId, stickyKeyHash)) {
+ return false;
+ }
+ Long2ObjectSortedMap<IntIntPair> ledgerPendingAcks =
+ pendingAcks.computeIfAbsent(ledgerId, k -> new Long2ObjectRBTreeMap<>());
+ IntIntPair previous = ledgerPendingAcks.put(entryId, IntIntPair.of(remainingUnacked, stickyKeyHash));
+ if (previous == null) {
+ size++;
+ }
+ return true;
+ } finally {
+ writeLock.unlock();
+ }
+ }
+
+ @Override
+ public boolean contains(long ledgerId, long entryId) {
+ try {
+ readLock.lock();
+ Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ return ledgerMap != null && ledgerMap.containsKey(entryId);
+ } finally {
+ readLock.unlock();
+ }
+ }
+
+ @Override
+ public int getRemainingUnacked(long ledgerId, long entryId) {
+ try {
+ readLock.lock();
+ Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ IntIntPair value = ledgerMap == null ? null : ledgerMap.get(entryId);
+ return value == null ? PENDING_ACK_NOT_FOUND : value.leftInt();
+ } finally {
+ readLock.unlock();
+ }
+ }
+
+ @Override
+ public boolean remove(long ledgerId, long entryId, int remainingUnacked, int stickyKeyHash) {
+ try {
+ writeLock.lock();
+ Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ if (ledgerMap == null) {
+ return false;
+ }
+ boolean removed = ledgerMap.remove(entryId, IntIntPair.of(remainingUnacked, stickyKeyHash));
+ if (removed) {
+ size--;
+ handleRemovePendingAck(ledgerId, entryId, stickyKeyHash);
+ }
+ if (removed && ledgerMap.isEmpty()) {
+ pendingAcks.remove(ledgerId);
+ }
+ return removed;
+ } finally {
+ writeLock.unlock();
+ }
+ }
+
+ @Override
+ public int removeAndGetRemainingUnacked(long ledgerId, long entryId) {
+ try {
+ writeLock.lock();
+ Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ if (ledgerMap == null) {
+ return PENDING_ACK_NOT_FOUND;
+ }
+ IntIntPair removed = ledgerMap.remove(entryId);
+ if (removed == null) {
+ return PENDING_ACK_NOT_FOUND;
+ }
+ size--;
+ handleRemovePendingAck(ledgerId, entryId, removed.rightInt());
+ if (ledgerMap.isEmpty()) {
+ pendingAcks.remove(ledgerId);
+ }
+ return removed.leftInt();
+ } finally {
+ writeLock.unlock();
+ }
+ }
+
+ @Override
+ public boolean updateRemainingUnacked(long ledgerId, long entryId, int ackedDelta) {
+ try {
+ writeLock.lock();
+ Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ IntIntPair value = ledgerMap == null ? null : ledgerMap.get(entryId);
+ if (value == null) {
+ return false;
+ }
+ ledgerMap.put(entryId, IntIntPair.of(value.leftInt() - ackedDelta, value.rightInt()));
+ return true;
+ } finally {
+ writeLock.unlock();
+ }
+ }
+
+ @Override
+ public long forEachScan() {
+ forEachSum = 0;
+ forEach(pendingAckScanner);
+ return forEachSum;
+ }
+
+ @Override
+ public long removeAllUpTo(long markDeleteLedgerId, long markDeleteEntryId) {
+ removedCount = 0;
+ internalRemoveAllUpTo(markDeleteLedgerId, markDeleteEntryId, false, removedCounter);
+ return removedCount;
+ }
+
+ private void forEach(PendingAcksMap.PendingAcksConsumer processor) {
+ try {
+ readLock.lock();
+ processPendingAcks(processor);
+ } finally {
+ readLock.unlock();
+ }
+ }
+
+ private void processPendingAcks(PendingAcksMap.PendingAcksConsumer processor) {
+ for (Long2ObjectMap.Entry<Long2ObjectSortedMap<IntIntPair>> entry : pendingAcks.long2ObjectEntrySet()) {
+ long ledgerId = entry.getLongKey();
+ Long2ObjectSortedMap<IntIntPair> ledgerPendingAcks = entry.getValue();
+ for (Long2ObjectMap.Entry<IntIntPair> e : ledgerPendingAcks.long2ObjectEntrySet()) {
+ long entryId = e.getLongKey();
+ IntIntPair batchSizeAndStickyKeyHash = e.getValue();
+ processor.accept(ledgerId, entryId, batchSizeAndStickyKeyHash.leftInt(),
+ batchSizeAndStickyKeyHash.rightInt());
+ }
+ }
+ }
+
+ private void internalRemoveAllUpTo(long markDeleteLedgerId, long markDeleteEntryId, boolean useWriteLock,
+ PendingAcksMap.PendingAcksConsumer removedEntryCallback) {
+ PendingAcksMap.PendingAcksRemoveHandler pendingAcksRemoveHandler =
+ pendingAcksRemoveHandlerSupplier.get();
+ boolean acquiredWriteLock = false;
+ boolean batchStarted = false;
+ boolean retryWithWriteLock = false;
+ try {
+ if (useWriteLock) {
+ writeLock.lock();
+ acquiredWriteLock = true;
+ } else {
+ readLock.lock();
+ }
+ ObjectBidirectionalIterator<Long2ObjectMap.Entry<Long2ObjectSortedMap<IntIntPair>>> ledgerIterator =
+ pendingAcks.headMap(markDeleteLedgerId + 1).long2ObjectEntrySet().iterator();
+ while (ledgerIterator.hasNext()) {
+ Long2ObjectMap.Entry<Long2ObjectSortedMap<IntIntPair>> entry = ledgerIterator.next();
+ long ledgerId = entry.getLongKey();
+ Long2ObjectSortedMap<IntIntPair> ledgerMap = entry.getValue();
+ Long2ObjectSortedMap<IntIntPair> ledgerMapHead =
+ ledgerId == markDeleteLedgerId ? ledgerMap.headMap(markDeleteEntryId + 1) : ledgerMap;
+ ObjectBidirectionalIterator<Long2ObjectMap.Entry<IntIntPair>> entryIterator =
+ ledgerMapHead.long2ObjectEntrySet().iterator();
+ while (entryIterator.hasNext()) {
+ Long2ObjectMap.Entry<IntIntPair> pendingAckEntry = entryIterator.next();
+ long entryId = pendingAckEntry.getLongKey();
+ if (!acquiredWriteLock) {
+ retryWithWriteLock = true;
+ return;
+ }
+ IntIntPair value = pendingAckEntry.getValue();
+ int remainingUnacked = value.leftInt();
+ int stickyKeyHash = value.rightInt();
+ if (pendingAcksRemoveHandler != null) {
+ if (!batchStarted) {
+ pendingAcksRemoveHandler.startBatch();
+ batchStarted = true;
+ }
+ pendingAcksRemoveHandler.handleRemoving(consumer, ledgerId, entryId, stickyKeyHash,
+ closed);
+ }
+ if (removedEntryCallback != null) {
+ removedEntryCallback.accept(ledgerId, entryId, remainingUnacked, stickyKeyHash);
+ }
+ entryIterator.remove();
+ size--;
+ }
+ if (ledgerMap.isEmpty()) {
+ if (!acquiredWriteLock) {
+ retryWithWriteLock = true;
+ return;
+ }
+ ledgerIterator.remove();
+ }
+ }
+ } finally {
+ if (batchStarted) {
+ pendingAcksRemoveHandler.endBatch();
+ }
+ if (acquiredWriteLock) {
+ writeLock.unlock();
+ } else {
+ readLock.unlock();
+ if (retryWithWriteLock) {
+ internalRemoveAllUpTo(markDeleteLedgerId, markDeleteEntryId, true, removedEntryCallback);
+ }
+ }
+ }
+ }
+
+ private void handleRemovePendingAck(long ledgerId, long entryId, int stickyKeyHash) {
+ PendingAcksMap.PendingAcksRemoveHandler pendingAcksRemoveHandler =
+ pendingAcksRemoveHandlerSupplier.get();
+ if (pendingAcksRemoveHandler != null) {
+ pendingAcksRemoveHandler.handleRemoving(consumer, ledgerId, entryId, stickyKeyHash, closed);
+ }
+ }
+ }
+
+ private static final class CurrentPendingAcksMapStore implements PendingAcksStore {
+ private final PendingAcksMap pendingAcks = new PendingAcksMap(null, () -> null, () -> null);
+ private long forEachSum;
+ private long removedCount;
+ private final PendingAcksMap.PendingAcksConsumer pendingAckScanner =
+ (ledgerId, entryId, remainingUnacked, stickyKeyHash) ->
+ forEachSum += ledgerId + entryId + remainingUnacked + stickyKeyHash;
+ private final PendingAcksMap.PendingAcksConsumer removedCounter =
+ (ledgerId, entryId, remainingUnacked, stickyKeyHash) -> removedCount++;
+
+ @Override
+ public boolean addOrReplace(long ledgerId, long entryId, int remainingUnacked, int stickyKeyHash) {
+ return pendingAcks.addPendingAckIfAllowed(ledgerId, entryId, remainingUnacked, stickyKeyHash);
+ }
+
+ @Override
+ public boolean contains(long ledgerId, long entryId) {
+ return pendingAcks.contains(ledgerId, entryId);
+ }
+
+ @Override
+ public int getRemainingUnacked(long ledgerId, long entryId) {
+ return pendingAcks.getRemainingUnacked(ledgerId, entryId);
+ }
+
+ @Override
+ public boolean remove(long ledgerId, long entryId, int remainingUnacked, int stickyKeyHash) {
+ return pendingAcks.remove(ledgerId, entryId, remainingUnacked, stickyKeyHash);
+ }
+
+ @Override
+ public int removeAndGetRemainingUnacked(long ledgerId, long entryId) {
+ return pendingAcks.removeAndGetRemainingUnacked(ledgerId, entryId);
+ }
+
+ @Override
+ public boolean updateRemainingUnacked(long ledgerId, long entryId, int ackedDelta) {
+ return pendingAcks.updateRemainingUnacked(ledgerId, entryId, ackedDelta);
+ }
+
+ @Override
+ public long forEachScan() {
+ forEachSum = 0;
+ pendingAcks.forEach(pendingAckScanner);
+ return forEachSum;
+ }
+
+ @Override
+ public long removeAllUpTo(long markDeleteLedgerId, long markDeleteEntryId) {
+ removedCount = 0;
+ pendingAcks.removeAllUpTo(markDeleteLedgerId, markDeleteEntryId, removedCounter);
+ return removedCount;
+ }
+ }
+
+ private static void populate(PendingAcksStore store, int entries, int ledgers, long[] ledgerIds, long[] entryIds) {
+ int index = 0;
+ for (int ledger = 0; ledger < ledgers; ledger++) {
+ int entriesInLedger = entries / ledgers + (ledger < entries % ledgers ? 1 : 0);
+ for (int entry = 0; entry < entriesInLedger; entry++) {
+ if (ledgerIds != null) {
+ ledgerIds[index] = ledger;
+ entryIds[index] = entry;
+ }
+ store.addOrReplace(ledger, entry, remainingUnacked(index), stickyKeyHash(index));
+ index++;
+ }
+ }
+ }
+
+ private static int remainingUnacked(int index) {
+ return 1_000_000_000 - (index & 1023);
+ }
+
+ private static int stickyKeyHash(int index) {
+ return index * 0x9E3779B9;
+ }
+}
diff --git a/microbench/src/main/java/org/apache/pulsar/broker/service/package-info.java b/microbench/src/main/java/org/apache/pulsar/broker/service/package-info.java
new file mode 100644
index 0000000..ee8034b
--- /dev/null
+++ b/microbench/src/main/java/org/apache/pulsar/broker/service/package-info.java
@@ -0,0 +1,22 @@
+/*
+ * 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.
+ */
+/**
+ * Benchmarks for broker service internals.
+ */
+package org.apache.pulsar.broker.service;
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java
index fdad320..a7f2965 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java
@@ -28,7 +28,6 @@
import io.netty.util.concurrent.Future;
import io.netty.util.concurrent.Promise;
import io.opentelemetry.api.common.Attributes;
-import it.unimi.dsi.fastutil.ints.IntIntPair;
import it.unimi.dsi.fastutil.objects.ObjectIntPair;
import java.time.Instant;
import java.util.ArrayList;
@@ -83,6 +82,7 @@
public class Consumer {
private static final Logger LOG = Logger.get(Consumer.class);
+ private static final int PENDING_ACK_NOT_FOUND = PendingAcksMap.PENDING_ACK_NOT_FOUND;
private final Logger log;
private final Subscription subscription;
@@ -649,10 +649,10 @@
addAndGetUnAckedMsgs(ackOwnerConsumer, -(int) ackedCount);
}
} else if (!hasAckSet) {
- IntIntPair removed = ackOwnerConsumer.removePendingAckAndGet(
+ int removed = ackOwnerConsumer.removePendingAckAndGetRemainingUnacked(
position.getLedgerId(), position.getEntryId());
- if (removed != null) {
- addAndGetUnAckedMsgs(ackOwnerConsumer, -removed.leftInt());
+ if (removed != PENDING_ACK_NOT_FOUND) {
+ addAndGetUnAckedMsgs(ackOwnerConsumer, -removed);
updateBlockedConsumerOnUnackedMsgs(ackOwnerConsumer);
}
}
@@ -744,10 +744,10 @@
}
}
} else {
- IntIntPair removed = ackOwnerConsumer.removePendingAckAndGet(
+ int removed = ackOwnerConsumer.removePendingAckAndGetRemainingUnacked(
position.getLedgerId(), position.getEntryId());
- if (removed != null) {
- addAndGetUnAckedMsgs(ackOwnerConsumer, -removed.leftInt());
+ if (removed != PENDING_ACK_NOT_FOUND) {
+ addAndGetUnAckedMsgs(ackOwnerConsumer, -removed);
}
}
updateBlockedConsumerOnUnackedMsgs(ackOwnerConsumer);
@@ -841,16 +841,16 @@
*/
private ObjectIntPair<Consumer> getAckOwnerConsumerAndBatchSize(long ledgerId, long entryId) {
if (Subscription.isIndividualAckMode(subType)) {
- IntIntPair pendingAck = getPendingAcks().get(ledgerId, entryId);
- if (pendingAck != null) {
- return ObjectIntPair.of(this, pendingAck.leftInt());
+ int remainingUnacked = getPendingAcks().getRemainingUnacked(ledgerId, entryId);
+ if (remainingUnacked != PENDING_ACK_NOT_FOUND) {
+ return ObjectIntPair.of(this, remainingUnacked);
} else {
// If there are more consumers, this step will consume more CPU, and it should be optimized later.
for (Consumer consumer : subscription.getConsumers()) {
if (consumer != this) {
- pendingAck = consumer.getPendingAcks().get(ledgerId, entryId);
- if (pendingAck != null) {
- return ObjectIntPair.of(consumer, pendingAck.leftInt());
+ remainingUnacked = consumer.getPendingAcks().getRemainingUnacked(ledgerId, entryId);
+ if (remainingUnacked != PENDING_ACK_NOT_FOUND) {
+ return ObjectIntPair.of(consumer, remainingUnacked);
}
}
}
@@ -1193,14 +1193,13 @@
*
* <p>No-op if {@code pendingAcks} is not initialized.
*
- * @return the removed {@link IntIntPair#leftInt() remainingUnacked} and
- * {@link IntIntPair#rightInt() stickyKeyHash}, or {@code null} if not found
+ * @return the remaining unacked count, or {@link PendingAcksMap#PENDING_ACK_NOT_FOUND} if not found
*/
- public IntIntPair removePendingAckAndGet(long ledgerId, long entryId) {
+ public int removePendingAckAndGetRemainingUnacked(long ledgerId, long entryId) {
if (pendingAcks != null) {
- return pendingAcks.removeAndGet(ledgerId, entryId);
+ return pendingAcks.removeAndGetRemainingUnacked(ledgerId, entryId);
}
- return null;
+ return PENDING_ACK_NOT_FOUND;
}
/**
@@ -1269,9 +1268,10 @@
List<Position> pendingPositions = new ArrayList<>();
for (MessageIdData msg : messageIds) {
Position position = PositionFactory.create(msg.getLedgerId(), msg.getEntryId());
- IntIntPair pendingAck = pendingAcks.removeAndGet(position.getLedgerId(), position.getEntryId());
- if (pendingAck != null) {
- totalRedeliveryMessages += pendingAck.leftInt();
+ int remainingUnacked = pendingAcks.removeAndGetRemainingUnacked(
+ position.getLedgerId(), position.getEntryId());
+ if (remainingUnacked != PENDING_ACK_NOT_FOUND) {
+ totalRedeliveryMessages += remainingUnacked;
pendingPositions.add(position);
}
}
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PendingAckValues.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PendingAckValues.java
new file mode 100644
index 0000000..a27e937
--- /dev/null
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PendingAckValues.java
@@ -0,0 +1,77 @@
+/*
+ * 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.service;
+
+/**
+ * Utility methods for storing pending ack values in primitive maps.
+ *
+ * <p>The high 32 bits contain the remaining unacked count and the low 32 bits contain the sticky key hash.
+ * A value with {@link #NOT_FOUND} as the remaining unacked count and 0 as the sticky key hash is reserved
+ * as the missing-entry sentinel.
+ */
+final class PendingAckValues {
+ /**
+ * Remaining unacked count reserved for missing entries.
+ */
+ static final int NOT_FOUND = -1;
+
+ /**
+ * Packed value reserved for missing entries in primitive maps.
+ */
+ static final long PACKED_NOT_FOUND = (long) NOT_FOUND << Integer.SIZE;
+
+ private static final long STICKY_KEY_HASH_MASK = 0xFFFF_FFFFL;
+
+ private PendingAckValues() {
+ }
+
+ /**
+ * Packs a pending ack's remaining unacked count and sticky key hash into one long value.
+ *
+ * @throws IllegalArgumentException if the remaining unacked count is negative
+ */
+ static long pack(int remainingUnacked, int stickyKeyHash) {
+ if (remainingUnacked < 0) {
+ throw new IllegalArgumentException("remainingUnacked must be non-negative");
+ }
+ return ((long) remainingUnacked << Integer.SIZE) | (stickyKeyHash & STICKY_KEY_HASH_MASK);
+ }
+
+ /**
+ * Returns true if the packed value is the missing-entry sentinel.
+ */
+ static boolean isNotFound(long packedValue) {
+ return packedValue == PACKED_NOT_FOUND;
+ }
+
+ /**
+ * Returns the remaining unacked count stored in the high 32 bits.
+ */
+ static int remainingUnacked(long packedValue) {
+ return (int) (packedValue >> Integer.SIZE);
+ }
+
+ /**
+ * Returns the sticky key hash stored in the low 32 bits.
+ */
+ static int stickyKeyHash(long packedValue) {
+ return (int) packedValue;
+ }
+
+}
diff --git a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PendingAcksMap.java b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PendingAcksMap.java
index 9eb43ef..7932106 100644
--- a/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PendingAcksMap.java
+++ b/pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PendingAcksMap.java
@@ -18,7 +18,9 @@
*/
package org.apache.pulsar.broker.service;
-import it.unimi.dsi.fastutil.ints.IntIntPair;
+import it.unimi.dsi.fastutil.longs.Long2LongMap;
+import it.unimi.dsi.fastutil.longs.Long2LongRBTreeMap;
+import it.unimi.dsi.fastutil.longs.Long2LongSortedMap;
import it.unimi.dsi.fastutil.longs.Long2ObjectMap;
import it.unimi.dsi.fastutil.longs.Long2ObjectRBTreeMap;
import it.unimi.dsi.fastutil.longs.Long2ObjectSortedMap;
@@ -40,6 +42,8 @@
* running.
*/
public class PendingAcksMap {
+ static final int PENDING_ACK_NOT_FOUND = PendingAckValues.NOT_FOUND;
+
/**
* Callback interface for handling the addition of pending acknowledgments.
*/
@@ -98,7 +102,12 @@
}
private final Consumer consumer;
- private final Long2ObjectSortedMap<Long2ObjectSortedMap<IntIntPair>> pendingAcks;
+
+ /**
+ * Pending acknowledgments grouped by ledger id. The outer map key is ledgerId and each inner map key is entryId.
+ * The inner map value is packed by {@link PendingAckValues} with the remaining unacked count and sticky key hash.
+ */
+ private final Long2ObjectSortedMap<Long2LongSortedMap> pendingAcks;
private final Supplier<PendingAcksAddHandler> pendingAcksAddHandlerSupplier;
private final Supplier<PendingAcksRemoveHandler> pendingAcksRemoveHandlerSupplier;
private final Lock readLock;
@@ -145,10 +154,11 @@
&& !pendingAcksAddHandler.handleAdding(consumer, ledgerId, entryId, stickyKeyHash)) {
return false;
}
- Long2ObjectSortedMap<IntIntPair> ledgerPendingAcks =
- pendingAcks.computeIfAbsent(ledgerId, k -> new Long2ObjectRBTreeMap<>());
- IntIntPair previous = ledgerPendingAcks.put(entryId, IntIntPair.of(remainingUnacked, stickyKeyHash));
- if (previous == null) {
+ Long2LongSortedMap ledgerPendingAcks =
+ pendingAcks.computeIfAbsent(ledgerId, k -> newLedgerPendingAcks());
+ long packedValue = PendingAckValues.pack(remainingUnacked, stickyKeyHash);
+ long previous = ledgerPendingAcks.put(entryId, packedValue);
+ if (PendingAckValues.isNotFound(previous)) {
size++;
}
return true;
@@ -180,19 +190,14 @@
}
}
- // iterate all pending acks and process them
+ // this code uses for loops intentionally, don't refactor to use forEach
private void processPendingAcks(PendingAcksConsumer processor) {
- // this code uses for loops intentionally, don't refactor to use forEach
- // iterate the outer map
- for (Long2ObjectMap.Entry<Long2ObjectSortedMap<IntIntPair>> entry : pendingAcks.long2ObjectEntrySet()) {
+ for (Long2ObjectMap.Entry<Long2LongSortedMap> entry : pendingAcks.long2ObjectEntrySet()) {
long ledgerId = entry.getLongKey();
- Long2ObjectSortedMap<IntIntPair> ledgerPendingAcks = entry.getValue();
- // iterate the inner map
- for (Long2ObjectMap.Entry<IntIntPair> e : ledgerPendingAcks.long2ObjectEntrySet()) {
- long entryId = e.getLongKey();
- IntIntPair batchSizeAndStickyKeyHash = e.getValue();
- processor.accept(ledgerId, entryId, batchSizeAndStickyKeyHash.leftInt(),
- batchSizeAndStickyKeyHash.rightInt());
+ for (Long2LongMap.Entry pendingAckEntry : entry.getValue().long2LongEntrySet()) {
+ long packedValue = pendingAckEntry.getLongValue();
+ processor.accept(ledgerId, pendingAckEntry.getLongKey(),
+ PendingAckValues.remainingUnacked(packedValue), PendingAckValues.stickyKeyHash(packedValue));
}
}
}
@@ -255,7 +260,7 @@
public boolean contains(long ledgerId, long entryId) {
try {
readLock.lock();
- Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ Long2LongSortedMap ledgerMap = pendingAcks.get(ledgerId);
if (ledgerMap == null) {
return false;
}
@@ -266,20 +271,20 @@
}
/**
- * Get the pending ack for the given ledger ID and entry ID.
+ * Get the remaining unacked count for the given ledger ID and entry ID.
*
- * @param ledgerId the ledger ID
- * @param entryId the entry ID
- * @return the pending ack, or null if not found
+ * @return the remaining unacked count, or {@link #PENDING_ACK_NOT_FOUND} if not found
*/
- public IntIntPair get(long ledgerId, long entryId) {
+ int getRemainingUnacked(long ledgerId, long entryId) {
try {
readLock.lock();
- Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ Long2LongSortedMap ledgerMap = pendingAcks.get(ledgerId);
if (ledgerMap == null) {
- return null;
+ return PENDING_ACK_NOT_FOUND;
}
- return ledgerMap.get(entryId);
+ long packedValue = ledgerMap.get(entryId);
+ return PendingAckValues.isNotFound(packedValue)
+ ? PENDING_ACK_NOT_FOUND : PendingAckValues.remainingUnacked(packedValue);
} finally {
readLock.unlock();
}
@@ -297,19 +302,21 @@
public boolean remove(long ledgerId, long entryId, int batchSize, int stickyKeyHash) {
try {
writeLock.lock();
- Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
- if (ledgerMap == null) {
+ Long2LongSortedMap ledgerMap = pendingAcks.get(ledgerId);
+ if (batchSize < 0) {
return false;
}
- boolean removed = ledgerMap.remove(entryId, IntIntPair.of(batchSize, stickyKeyHash));
- if (removed) {
- size--;
- handleRemovePendingAck(ledgerId, entryId, stickyKeyHash);
+ long expectedValue = PendingAckValues.pack(batchSize, stickyKeyHash);
+ if (ledgerMap == null || ledgerMap.get(entryId) != expectedValue) {
+ return false;
}
- if (removed && ledgerMap.isEmpty()) {
+ ledgerMap.remove(entryId);
+ size--;
+ handleRemovePendingAck(ledgerId, entryId, stickyKeyHash);
+ if (ledgerMap.isEmpty()) {
pendingAcks.remove(ledgerId);
}
- return removed;
+ return true;
} finally {
writeLock.unlock();
}
@@ -327,16 +334,19 @@
public boolean updateRemainingUnacked(long ledgerId, long entryId, int ackedDelta) {
try {
writeLock.lock();
- Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ Long2LongSortedMap ledgerMap = pendingAcks.get(ledgerId);
if (ledgerMap == null) {
return false;
}
- IntIntPair current = ledgerMap.get(entryId);
- if (current == null) {
+ long packedValue = ledgerMap.get(entryId);
+ if (PendingAckValues.isNotFound(packedValue)) {
return false;
}
- int newRemaining = current.leftInt() - ackedDelta;
- ledgerMap.put(entryId, IntIntPair.of(newRemaining, current.rightInt()));
+ int newRemaining = PendingAckValues.remainingUnacked(packedValue) - ackedDelta;
+ if (newRemaining < 0) {
+ return false;
+ }
+ ledgerMap.put(entryId, PendingAckValues.pack(newRemaining, PendingAckValues.stickyKeyHash(packedValue)));
return true;
} finally {
writeLock.unlock();
@@ -353,51 +363,47 @@
public boolean remove(long ledgerId, long entryId) {
try {
writeLock.lock();
- Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ Long2LongSortedMap ledgerMap = pendingAcks.get(ledgerId);
if (ledgerMap == null) {
return false;
}
- IntIntPair removedEntry = ledgerMap.remove(entryId);
- boolean removed = removedEntry != null;
- if (removed) {
- size--;
- int stickyKeyHash = removedEntry.rightInt();
- handleRemovePendingAck(ledgerId, entryId, stickyKeyHash);
+ long removedEntry = ledgerMap.remove(entryId);
+ if (PendingAckValues.isNotFound(removedEntry)) {
+ return false;
}
- if (removed && ledgerMap.isEmpty()) {
+ size--;
+ handleRemovePendingAck(ledgerId, entryId, PendingAckValues.stickyKeyHash(removedEntry));
+ if (ledgerMap.isEmpty()) {
pendingAcks.remove(ledgerId);
}
- return removed;
+ return true;
} finally {
writeLock.unlock();
}
}
/**
- * Atomically remove and return the pending ack for the given ledger ID and entry ID.
- * Unlike {@link #remove(long, long)}, this method returns the removed entry so the caller
- * can access the batch size and sticky key hash without a separate get operation.
+ * Atomically remove and return the remaining unacked count for the given ledger ID and entry ID.
*
- * @param ledgerId the ledger ID
- * @param entryId the entry ID
- * @return the removed entry as an IntIntPair (batchSize, stickyKeyHash), or null if not found
+ * @return the remaining unacked count, or {@link #PENDING_ACK_NOT_FOUND} if not found
*/
- public IntIntPair removeAndGet(long ledgerId, long entryId) {
+ int removeAndGetRemainingUnacked(long ledgerId, long entryId) {
try {
writeLock.lock();
- Long2ObjectSortedMap<IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+ Long2LongSortedMap ledgerMap = pendingAcks.get(ledgerId);
if (ledgerMap == null) {
- return null;
+ return PENDING_ACK_NOT_FOUND;
}
- IntIntPair removedEntry = ledgerMap.remove(entryId);
- if (removedEntry != null) {
- size--;
- handleRemovePendingAck(ledgerId, entryId, removedEntry.rightInt());
+ long removedEntry = ledgerMap.remove(entryId);
+ if (PendingAckValues.isNotFound(removedEntry)) {
+ return PENDING_ACK_NOT_FOUND;
}
- if (removedEntry != null && ledgerMap.isEmpty()) {
+ size--;
+ handleRemovePendingAck(ledgerId, entryId, PendingAckValues.stickyKeyHash(removedEntry));
+ if (ledgerMap.isEmpty()) {
pendingAcks.remove(ledgerId);
}
- return removedEntry;
+ return PendingAckValues.remainingUnacked(removedEntry);
} finally {
writeLock.unlock();
}
@@ -444,30 +450,30 @@
} else {
readLock.lock();
}
- ObjectBidirectionalIterator<Long2ObjectMap.Entry<Long2ObjectSortedMap<IntIntPair>>> ledgerMapIterator =
+ ObjectBidirectionalIterator<Long2ObjectMap.Entry<Long2LongSortedMap>> ledgerMapIterator =
pendingAcks.headMap(markDeleteLedgerId + 1).long2ObjectEntrySet().iterator();
while (ledgerMapIterator.hasNext()) {
- Long2ObjectMap.Entry<Long2ObjectSortedMap<IntIntPair>> entry = ledgerMapIterator.next();
+ Long2ObjectMap.Entry<Long2LongSortedMap> entry = ledgerMapIterator.next();
long ledgerId = entry.getLongKey();
- Long2ObjectSortedMap<IntIntPair> ledgerMap = entry.getValue();
- Long2ObjectSortedMap<IntIntPair> ledgerMapHead;
+ Long2LongSortedMap ledgerMap = entry.getValue();
+ Long2LongSortedMap ledgerMapHead;
if (ledgerId == markDeleteLedgerId) {
ledgerMapHead = ledgerMap.headMap(markDeleteEntryId + 1);
} else {
ledgerMapHead = ledgerMap;
}
- ObjectBidirectionalIterator<Long2ObjectMap.Entry<IntIntPair>> entryMapIterator =
- ledgerMapHead.long2ObjectEntrySet().iterator();
+ ObjectBidirectionalIterator<Long2LongMap.Entry> entryMapIterator =
+ ledgerMapHead.long2LongEntrySet().iterator();
while (entryMapIterator.hasNext()) {
- Long2ObjectMap.Entry<IntIntPair> intIntPairEntry = entryMapIterator.next();
- long entryId = intIntPairEntry.getLongKey();
+ Long2LongMap.Entry pendingAckEntry = entryMapIterator.next();
+ long entryId = pendingAckEntry.getLongKey();
if (!acquiredWriteLock) {
retryWithWriteLock = true;
return;
}
- IntIntPair value = intIntPairEntry.getValue();
- int batchSize = value.leftInt();
- int stickyKeyHash = value.rightInt();
+ long packedValue = pendingAckEntry.getLongValue();
+ int batchSize = PendingAckValues.remainingUnacked(packedValue);
+ int stickyKeyHash = PendingAckValues.stickyKeyHash(packedValue);
if (pendingAcksRemoveHandler != null) {
if (!batchStarted) {
pendingAcksRemoveHandler.startBatch();
@@ -504,6 +510,12 @@
}
}
+ private static Long2LongSortedMap newLedgerPendingAcks() {
+ Long2LongRBTreeMap ledgerPendingAcks = new Long2LongRBTreeMap();
+ ledgerPendingAcks.defaultReturnValue(PendingAckValues.PACKED_NOT_FOUND);
+ return ledgerPendingAcks;
+ }
+
private void handleRemovePendingAck(long ledgerId, long entryId, int stickyKeyHash) {
PendingAcksRemoveHandler pendingAcksRemoveHandler = pendingAcksRemoveHandlerSupplier.get();
if (pendingAcksRemoveHandler != null) {
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PendingAckValuesTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PendingAckValuesTest.java
new file mode 100644
index 0000000..1931654
--- /dev/null
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PendingAckValuesTest.java
@@ -0,0 +1,60 @@
+/*
+ * 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.service;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+import org.testng.annotations.Test;
+
+public class PendingAckValuesTest {
+ @Test
+ public void pack_RoundTripsRemainingUnackedAndStickyKeyHash() {
+ int[][] values = new int[][] {
+ {0, 0},
+ {1, -1},
+ {Integer.MAX_VALUE, Integer.MIN_VALUE},
+ {42, Integer.MAX_VALUE},
+ {1_000_000, -123456789}
+ };
+
+ for (int[] value : values) {
+ long packedValue = PendingAckValues.pack(value[0], value[1]);
+
+ assertFalse(PendingAckValues.isNotFound(packedValue));
+ assertEquals(PendingAckValues.remainingUnacked(packedValue), value[0]);
+ assertEquals(PendingAckValues.stickyKeyHash(packedValue), value[1]);
+ }
+ }
+
+ @Test
+ public void packedNotFound_UsesReservedNegativeRemainingUnackedValue() {
+ assertTrue(PendingAckValues.isNotFound(PendingAckValues.PACKED_NOT_FOUND));
+ assertEquals(PendingAckValues.remainingUnacked(PendingAckValues.PACKED_NOT_FOUND),
+ PendingAckValues.NOT_FOUND);
+ assertEquals(PendingAckValues.stickyKeyHash(PendingAckValues.PACKED_NOT_FOUND), 0);
+ assertFalse(PendingAckValues.isNotFound(PendingAckValues.pack(0, 0)));
+ }
+
+ @Test
+ public void pack_RejectsNegativeRemainingUnacked() {
+ assertThrows(IllegalArgumentException.class, () -> PendingAckValues.pack(-1, 0));
+ }
+}
diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PendingAcksMapTest.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PendingAcksMapTest.java
index c13babb..59aa302 100644
--- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PendingAcksMapTest.java
+++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/service/PendingAcksMapTest.java
@@ -25,11 +25,11 @@
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.never;
import static org.mockito.Mockito.verify;
+import static org.mockito.Mockito.verifyNoMoreInteractions;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertTrue;
-import it.unimi.dsi.fastutil.ints.IntIntPair;
import java.util.ArrayList;
import java.util.List;
import org.testng.annotations.Test;
@@ -295,41 +295,178 @@
}
@Test
- public void removeAndGet_RemovesAndReturnsEntry() {
+ public void removeAndGetRemainingUnacked_RemovesAndReturnsEntry() {
Consumer consumer = createMockConsumer("consumer1");
PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> null);
pendingAcksMap.addPendingAckIfAllowed(1L, 1L, 5, 123);
- IntIntPair result = pendingAcksMap.removeAndGet(1L, 1L);
+ int result = pendingAcksMap.removeAndGetRemainingUnacked(1L, 1L);
- assertTrue(result != null);
- assertEquals(result.leftInt(), 5);
- assertEquals(result.rightInt(), 123);
+ assertEquals(result, 5);
assertFalse(pendingAcksMap.contains(1L, 1L));
assertEquals(pendingAcksMap.size(), 0);
}
@Test
- public void removeAndGet_ReturnsNullWhenNotFound() {
+ public void removeAndGetRemainingUnacked_ReturnsNotFoundWhenNotFound() {
Consumer consumer = createMockConsumer("consumer1");
PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> null);
- IntIntPair result = pendingAcksMap.removeAndGet(1L, 1L);
+ int result = pendingAcksMap.removeAndGetRemainingUnacked(1L, 1L);
- assertTrue(result == null);
+ assertEquals(result, PendingAcksMap.PENDING_ACK_NOT_FOUND);
assertEquals(pendingAcksMap.size(), 0);
}
@Test
- public void removeAndGet_InvokesRemoveHandler() {
+ public void removeAndGetRemainingUnacked_InvokesRemoveHandler() {
Consumer consumer = createMockConsumer("consumer1");
PendingAcksMap.PendingAcksRemoveHandler removeHandler = mock(PendingAcksMap.PendingAcksRemoveHandler.class);
PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> removeHandler);
pendingAcksMap.addPendingAckIfAllowed(1L, 1L, 5, 123);
- pendingAcksMap.removeAndGet(1L, 1L);
+ pendingAcksMap.removeAndGetRemainingUnacked(1L, 1L);
verify(removeHandler).handleRemoving(consumer, 1L, 1L, 123, false);
}
+ @Test
+ public void getRemainingUnacked_ReturnsStoredCount() {
+ Consumer consumer = createMockConsumer("consumer1");
+ PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> null);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 1L, 0, -1);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 2L, Integer.MAX_VALUE, Integer.MIN_VALUE);
+
+ assertEquals(pendingAcksMap.getRemainingUnacked(1L, 1L), 0);
+ assertEquals(pendingAcksMap.getRemainingUnacked(1L, 2L), Integer.MAX_VALUE);
+ assertEquals(pendingAcksMap.getRemainingUnacked(1L, 3L), PendingAcksMap.PENDING_ACK_NOT_FOUND);
+ }
+
+ @Test
+ public void removeAndGetRemainingUnacked_RemovesAndInvokesRemoveHandler() {
+ Consumer consumer = createMockConsumer("consumer1");
+ PendingAcksMap.PendingAcksRemoveHandler removeHandler = mock(PendingAcksMap.PendingAcksRemoveHandler.class);
+ PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> removeHandler);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 1L, 0, Integer.MIN_VALUE);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 2L, Integer.MAX_VALUE, -1);
+
+ assertEquals(pendingAcksMap.removeAndGetRemainingUnacked(1L, 1L), 0);
+ assertEquals(pendingAcksMap.removeAndGetRemainingUnacked(1L, 2L), Integer.MAX_VALUE);
+ assertEquals(pendingAcksMap.removeAndGetRemainingUnacked(1L, 3L), PendingAcksMap.PENDING_ACK_NOT_FOUND);
+
+ verify(removeHandler).handleRemoving(consumer, 1L, 1L, Integer.MIN_VALUE, false);
+ verify(removeHandler).handleRemoving(consumer, 1L, 2L, -1, false);
+ verifyNoMoreInteractions(removeHandler);
+ assertEquals(pendingAcksMap.size(), 0);
+ }
+
+ @Test
+ public void packedPendingAckFields_RoundTripThroughPublicAccessors() {
+ Consumer consumer = createMockConsumer("consumer1");
+ PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> null);
+ int[][] values = new int[][] {
+ {1, 123},
+ {0, 0},
+ {Integer.MAX_VALUE, Integer.MIN_VALUE},
+ {2, Integer.MAX_VALUE},
+ {42, -123456789}
+ };
+
+ for (int i = 0; i < values.length; i++) {
+ pendingAcksMap.addPendingAckIfAllowed(10L + i, 100L + i, values[i][0], values[i][1]);
+ }
+
+ List<int[]> forEachValues = new ArrayList<>();
+ pendingAcksMap.forEach((ledgerId, entryId, remainingUnacked, stickyKeyHash) ->
+ forEachValues.add(new int[] {
+ (int) ledgerId, (int) entryId, remainingUnacked, stickyKeyHash
+ }));
+ assertEquals(forEachValues.size(), values.length);
+ for (int i = 0; i < values.length; i++) {
+ assertEquals(pendingAcksMap.getRemainingUnacked(10L + i, 100L + i), values[i][0]);
+ assertEquals(forEachValues.get(i), new int[] {10 + i, 100 + i, values[i][0], values[i][1]});
+ }
+
+ for (int i = 0; i < values.length; i++) {
+ assertEquals(pendingAcksMap.removeAndGetRemainingUnacked(10L + i, 100L + i), values[i][0]);
+ }
+ assertEquals(pendingAcksMap.size(), 0);
+ }
+
+ @Test
+ public void updateRemainingUnacked_PreservesStickyKeyHash() {
+ Consumer consumer = createMockConsumer("consumer1");
+ PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> null);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 1L, 10, Integer.MIN_VALUE);
+
+ assertTrue(pendingAcksMap.updateRemainingUnacked(1L, 1L, 4));
+
+ assertOnlyPendingAck(pendingAcksMap, 1L, 1L, 6, Integer.MIN_VALUE);
+ }
+
+ @Test
+ public void updateRemainingUnacked_AllowsZeroRemainingUnacked() {
+ Consumer consumer = createMockConsumer("consumer1");
+ PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> null);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 1L, 1, -1);
+
+ assertTrue(pendingAcksMap.updateRemainingUnacked(1L, 1L, 1));
+
+ assertEquals(pendingAcksMap.getRemainingUnacked(1L, 1L), 0);
+ assertOnlyPendingAck(pendingAcksMap, 1L, 1L, 0, -1);
+ assertEquals(pendingAcksMap.size(), 1);
+ }
+
+ @Test
+ public void updateRemainingUnacked_DoesNotStoreNegativeRemainingUnacked() {
+ Consumer consumer = createMockConsumer("consumer1");
+ PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> null);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 1L, 1, 123);
+
+ assertFalse(pendingAcksMap.updateRemainingUnacked(1L, 1L, 2));
+
+ assertOnlyPendingAck(pendingAcksMap, 1L, 1L, 1, 123);
+ }
+
+ @Test
+ public void removeAllUpTo_PreservesPackedFieldsInCallback() {
+ Consumer consumer = createMockConsumer("consumer1");
+ PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> null);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 1L, Integer.MAX_VALUE, Integer.MIN_VALUE);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 2L, 1, -1);
+ pendingAcksMap.addPendingAckIfAllowed(2L, 1L, 2, Integer.MAX_VALUE);
+
+ List<int[]> callbackInvocations = new ArrayList<>();
+ pendingAcksMap.removeAllUpTo(1L, 2L, (ledgerId, entryId, remainingUnacked, stickyKeyHash) ->
+ callbackInvocations.add(new int[] {
+ (int) ledgerId, (int) entryId, remainingUnacked, stickyKeyHash
+ }));
+
+ assertEquals(callbackInvocations.size(), 2);
+ assertEquals(callbackInvocations.get(0), new int[] {1, 1, Integer.MAX_VALUE, Integer.MIN_VALUE});
+ assertEquals(callbackInvocations.get(1), new int[] {1, 2, 1, -1});
+ assertTrue(pendingAcksMap.contains(2L, 1L));
+ }
+
+ @Test
+ public void removeWithValue_MatchesSignedPackedFields() {
+ Consumer consumer = createMockConsumer("consumer1");
+ PendingAcksMap pendingAcksMap = new PendingAcksMap(consumer, () -> null, () -> null);
+ pendingAcksMap.addPendingAckIfAllowed(1L, 1L, Integer.MAX_VALUE, -1);
+
+ assertFalse(pendingAcksMap.remove(1L, 1L, Integer.MAX_VALUE, 1));
+ assertTrue(pendingAcksMap.remove(1L, 1L, Integer.MAX_VALUE, -1));
+ assertFalse(pendingAcksMap.contains(1L, 1L));
+ }
+
+ private static void assertOnlyPendingAck(PendingAcksMap pendingAcksMap, long ledgerId, long entryId,
+ int remainingUnacked, int stickyKeyHash) {
+ List<long[]> pendingAcks = new ArrayList<>();
+ pendingAcksMap.forEach((actualLedgerId, actualEntryId, actualRemainingUnacked, actualStickyKeyHash) ->
+ pendingAcks.add(new long[] {
+ actualLedgerId, actualEntryId, actualRemainingUnacked, actualStickyKeyHash
+ }));
+ assertEquals(pendingAcks.size(), 1);
+ assertEquals(pendingAcks.get(0), new long[] {ledgerId, entryId, remainingUnacked, stickyKeyHash});
+ }
}
