[fix][broker] Move pending acks cleanup to selected mark-delete callbacks

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Consumer.java

@@ -385,9 +385,18 @@ public Future<Void> sendMessages(final List<? extends Entry> entries,
                     } else {
                         stickyKeyHash = stickyKeyHashes.get(i);
                     }
-                    boolean sendingAllowed =
-                            pendingAcks.addPendingAckIfAllowed(entry.getLedgerId(), entry.getEntryId(), batchSize,
-                                    stickyKeyHash);
+                    boolean sendingAllowed;
+                    long[] ackSet = batchIndexesAcks == null ? null : batchIndexesAcks.getAckSet(i);
+                    int remainingUnacked;
+                    if (ackSet != null) {
+                        remainingUnacked = BitSet.valueOf(ackSet).cardinality();
+                        unackedMessages -= (batchSize - remainingUnacked);
+                    } else {
+                        remainingUnacked = batchSize;
+                    }
+                    sendingAllowed =
+                            pendingAcks.addPendingAckIfAllowed(entry.getLedgerId(), entry.getEntryId(),
+                                    remainingUnacked, stickyKeyHash);
                     if (!sendingAllowed) {
                         // sending isn't allowed when pending acks doesn't accept adding the entry
                         // this happens when Key_Shared draining hashes contains the stickyKeyHash
@@ -401,10 +410,6 @@ public Future<Void> sendMessages(final List<? extends Entry> entries,
                                 .attr("batchSize", batchSize)
                                 .log("Skipping sending of entry since adding to pending acks failed");
                     } else {
-                        long[] ackSet = batchIndexesAcks == null ? null : batchIndexesAcks.getAckSet(i);
-                        if (ackSet != null) {
-                            unackedMessages -= (batchSize - BitSet.valueOf(ackSet).cardinality());
-                        }
                         log.debug()
                                 .attr("ledgerId", entry.getLedgerId())
                                 .attr("entryId", entry.getEntryId())
@@ -596,6 +601,10 @@ private CompletableFuture<Long> individualAckNormal(CommandAck ack, Map<String,
                                 .syncBatchPositionBitSetForPendingAck(position);
                     }
                 }
+                if (ackedCount > 0) {
+                    ackOwnerConsumer.getPendingAcks().updateRemainingUnacked(
+                            position.getLedgerId(), position.getEntryId(), (int) ackedCount);
+                }
                 addAndGetUnAckedMsgs(ackOwnerConsumer, -(int) ackedCount);
             } else {
                 position = PositionFactory.create(msgId.getLedgerId(), msgId.getEntryId());
@@ -679,6 +688,12 @@ private CompletableFuture<Long> individualAckWithTransaction(CommandAck ack) {
                 }
                 AckSetStateUtil.getAckSetState(position).setAckSet(ackSets);
                 ackedCount = getAckedCountForTransactionAck(batchSize, ackSets);
+                // Update the stored remaining unacked count for this entry so that
+                // removal paths (expire, skip, redeliver) can use it directly.
+                if (ackedCount > 0) {
+                    ackOwnerConsumer.getPendingAcks().updateRemainingUnacked(
+                            position.getLedgerId(), position.getEntryId(), (int) ackedCount);
+                }
             }
 
             addAndGetUnAckedMsgs(ackOwnerConsumer, -(int) ackedCount);
@@ -747,19 +762,6 @@ private long getAckedCountForTransactionAck(int batchSize, long[] ackSets) {
         return ackedCount;
     }
 
-    private long getUnAckedCountForBatchIndexLevelEnabled(Position position, int batchSize) {
-        long unAckedCount = batchSize;
-        if (isAcknowledgmentAtBatchIndexLevelEnabled) {
-            long[] cursorAckSet = getCursorAckSet(position);
-            if (cursorAckSet != null) {
-                BitSetRecyclable cursorBitSet = BitSetRecyclable.create().resetWords(cursorAckSet);
-                unAckedCount = cursorBitSet.cardinality();
-                cursorBitSet.recycle();
-            }
-        }
-        return unAckedCount;
-    }
-
     private void checkAckValidationError(CommandAck ack, Position position) {
         if (ack.hasValidationError()) {
             log.warn()
@@ -1139,9 +1141,8 @@ public void removePendingAcksUpToPositionAndDecrementUnacked(long markDeleteLedg
 
         MutableInt mutableTotalUnacked = new MutableInt(0);
         pendingAcks.removeAllUpTo(markDeleteLedgerId, markDeleteEntryId,
-                (ledgerId, entryId, batchSize, stickyKeyHash) -> {
-                    mutableTotalUnacked.add((int) getUnAckedCountForBatchIndexLevelEnabled(
-                            PositionFactory.create(ledgerId, entryId), batchSize));
+                (ledgerId, entryId, remainingUnacked, stickyKeyHash) -> {
+                    mutableTotalUnacked.add(remainingUnacked);
                 });
         int totalUnacked = mutableTotalUnacked.intValue();
         if (totalUnacked > 0) {
@@ -1160,11 +1161,8 @@ public void redeliverUnacknowledgedMessages(long consumerEpoch) {
         if (pendingAcks != null) {
             List<Position> pendingPositions = new ArrayList<>((int) pendingAcks.size());
             MutableInt totalRedeliveryMessages = new MutableInt(0);
-            pendingAcks.forEachAndClear((ledgerId, entryId, batchSize, stickyKeyHash) -> {
-                int unAckedCount =
-                        (int) getUnAckedCountForBatchIndexLevelEnabled(PositionFactory.create(ledgerId, entryId),
-                                batchSize);
-                totalRedeliveryMessages.add(unAckedCount);
+            pendingAcks.forEachAndClear((ledgerId, entryId, remainingUnacked, stickyKeyHash) -> {
+                totalRedeliveryMessages.add(remainingUnacked);
                 pendingPositions.add(PositionFactory.create(ledgerId, entryId));
             });
 
@@ -1193,8 +1191,7 @@ public void redeliverUnacknowledgedMessages(List<MessageIdData> messageIds) {
             Position position = PositionFactory.create(msg.getLedgerId(), msg.getEntryId());
             IntIntPair pendingAck = pendingAcks.removeAndGet(position.getLedgerId(), position.getEntryId());
             if (pendingAck != null) {
-                int unAckedCount = (int) getUnAckedCountForBatchIndexLevelEnabled(position, pendingAck.leftInt());
-                totalRedeliveryMessages += unAckedCount;
+                totalRedeliveryMessages += pendingAck.leftInt();
                 pendingPositions.add(position);
             }
         }

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/Dispatcher.java

@@ -142,6 +142,23 @@ default boolean checkAndUnblockIfStuck() {
         return false;
     }
 
+    /**
+     * Prune pending-ack entries up to the specified position and updates the
+     * consumer unacked-message counters accordingly.
+     *
+     * <p>This hook is invoked after the cursor mark-delete operation completes
+     * (for example, during message expiry, skip, or clear backlog). Since the
+     * cursor ack set may no longer be available after mark-delete, the counter
+     * adjustment relies on the remaining unacked count stored in the
+     * {@code PendingAcksMap} entries.
+     *
+     * @param ledgerId the ledger ID of the inclusive upper bound position
+     * @param entryId the entry ID of the inclusive upper bound position
+     */
+    default void prunePendingAcksUpToPosition(long ledgerId, long entryId) {
+        // No-op by default
+    }
+
     /**
      * A callback hook after acknowledge messages.
      * @param exOfDeletion the ex of {@link org.apache.bookkeeper.mledger.ManagedCursor#asyncDelete},

pulsar-broker/src/main/java/org/apache/pulsar/broker/service/PendingAcksMap.java

@@ -87,12 +87,13 @@ public interface PendingAcksConsumer {
         /**
          * Accept a pending acknowledgment.
          *
-         * @param ledgerId      the ledger ID
-         * @param entryId       the entry ID
-         * @param batchSize     the batch size
-         * @param stickyKeyHash the sticky key hash
+         * @param ledgerId          the ledger ID
+         * @param entryId           the entry ID
+         * @param remainingUnacked  the number of remaining unacked messages in this entry
+         *                          (accounts for batch index level acknowledgments)
+         * @param stickyKeyHash     the sticky key hash
          */
-        void accept(long ledgerId, long entryId, int batchSize, int stickyKeyHash);
+        void accept(long ledgerId, long entryId, int remainingUnacked, int stickyKeyHash);
     }
 
     private final Consumer consumer;
@@ -122,11 +123,12 @@ public interface PendingAcksConsumer {
      *
      * @param ledgerId the ledger ID
      * @param entryId the entry ID
-     * @param batchSize the batch size
+     * @param remainingUnacked the number of remaining unacked messages in this entry
+     *                         (for batch entries with some indexes already acked, this may be less than batchSize)
      * @param stickyKeyHash the sticky key hash
      * @return true if the pending ack was added, and it's allowed to send a message, false otherwise
      */
-    public boolean addPendingAckIfAllowed(long ledgerId, long entryId, int batchSize, int stickyKeyHash) {
+    public boolean addPendingAckIfAllowed(long ledgerId, long entryId, int remainingUnacked, int stickyKeyHash) {
         try {
             writeLock.lock();
             // prevent adding sticky hash to pending acks if the PendingAcksMap has already been closed
@@ -143,7 +145,7 @@ public boolean addPendingAckIfAllowed(long ledgerId, long entryId, int batchSize
             }
             TreeMap<Long, IntIntPair> ledgerPendingAcks =
                     pendingAcks.computeIfAbsent(ledgerId, k -> new TreeMap<>());
-            ledgerPendingAcks.put(entryId, IntIntPair.of(batchSize, stickyKeyHash));
+            ledgerPendingAcks.put(entryId, IntIntPair.of(remainingUnacked, stickyKeyHash));
             return true;
         } finally {
             writeLock.unlock();
@@ -311,6 +313,34 @@ public boolean remove(long ledgerId, long entryId, int batchSize, int stickyKeyH
         }
     }
 
+    /**
+     * Atomically update the remaining unacked count for a pending ack entry by subtracting the given delta.
+     * Called from the ack handler after computing the number of batch indexes acknowledged in a partial ack.
+     *
+     * @param ledgerId the ledger ID
+     * @param entryId the entry ID
+     * @param ackedDelta the number of batch indexes that were just acknowledged
+     * @return true if the entry was found and updated, false otherwise
+     */
+    public boolean updateRemainingUnacked(long ledgerId, long entryId, int ackedDelta) {
+        try {
+            writeLock.lock();
+            TreeMap<Long, IntIntPair> ledgerMap = pendingAcks.get(ledgerId);
+            if (ledgerMap == null) {
+                return false;
+            }
+            IntIntPair current = ledgerMap.get(entryId);
+            if (current == null) {
+                return false;
+            }
+            int newRemaining = current.leftInt() - ackedDelta;
+            ledgerMap.put(entryId, IntIntPair.of(newRemaining, current.rightInt()));
+            return true;
+        } finally {
+            writeLock.unlock();
+        }
+    }
+
     /**
      * Remove the pending ack for the given ledger ID and entry ID.
      *

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

@@ -145,7 +145,6 @@ public class PersistentDispatcherMultipleConsumers extends AbstractPersistentDis
     protected enum ReadType {
         Normal, Replay
     }
-    private Position lastMarkDeletePositionBeforeReadMoreEntries;
     private volatile long readMoreEntriesCallCount;
 
     public PersistentDispatcherMultipleConsumers(PersistentTopic topic, ManagedCursor cursor,
@@ -362,17 +361,6 @@ public synchronized void readMoreEntries() {
         // increment the counter for readMoreEntries calls, to track the number of times readMoreEntries is called
         readMoreEntriesCallCount++;
 
-        // remove possible expired messages from redelivery tracker and pending acks
-        Position markDeletePosition = cursor.getMarkDeletedPosition();
-        if (lastMarkDeletePositionBeforeReadMoreEntries != markDeletePosition) {
-            redeliveryMessages.removeAllUpTo(markDeletePosition.getLedgerId(), markDeletePosition.getEntryId());
-            for (Consumer consumer : consumerList) {
-                consumer.removePendingAcksUpToPositionAndDecrementUnacked(
-                        markDeletePosition.getLedgerId(), markDeletePosition.getEntryId());
-            }
-            lastMarkDeletePositionBeforeReadMoreEntries = markDeletePosition;
-        }
-
         // totalAvailablePermits may be updated by other threads
         int firstAvailableConsumerPermits = getFirstAvailableConsumerPermits();
         int currentTotalAvailablePermits = Math.max(totalAvailablePermits, firstAvailableConsumerPermits);
@@ -600,6 +588,14 @@ public CopyOnWriteArrayList<Consumer> getConsumers() {
         return consumerList;
     }
 
+    @Override
+    public void prunePendingAcksUpToPosition(long ledgerId, long entryId) {
+        redeliveryMessages.removeAllUpTo(ledgerId, entryId);
+        for (Consumer consumer : consumerList) {
+            consumer.removePendingAcksUpToPositionAndDecrementUnacked(ledgerId, entryId);
+        }
+    }
+
     @Override
     public synchronized boolean canUnsubscribe(Consumer consumer) {
         return consumerList.size() == 1 && consumerSet.contains(consumer);

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

@@ -38,6 +38,7 @@
 import org.apache.bookkeeper.mledger.PositionFactory;
 import org.apache.bookkeeper.mledger.impl.ManagedLedgerImpl;
 import org.apache.bookkeeper.mledger.proto.ManagedLedgerInfo;
+import org.apache.pulsar.broker.service.Dispatcher;
 import org.apache.pulsar.broker.service.MessageExpirer;
 import org.apache.pulsar.client.impl.MessageImpl;
 import org.apache.pulsar.common.api.proto.CommandSubscribe.SubType;
@@ -211,9 +212,18 @@ public void markDeleteComplete(Object ctx) {
             long numMessagesExpired = (long) ctx - cursor.getNumberOfEntriesInBacklog(false);
             msgExpired.recordMultipleEvents(numMessagesExpired, 0 /* no value stats */);
             totalMsgExpired.add(numMessagesExpired);
-            // If the subscription is a Key_Shared subscription, we should to trigger message dispatch.
-            if (subscription != null && subscription.getType() == SubType.Key_Shared) {
-                subscription.getDispatcher().markDeletePositionMoveForward();
+            if (subscription != null) {
+                Dispatcher dispatcher = subscription.getDispatcher();
+                if (dispatcher != null) {
+                    Position mdPos = cursor.getMarkDeletedPosition();
+                    if (mdPos != null) {
+                        dispatcher.prunePendingAcksUpToPosition(mdPos.getLedgerId(), mdPos.getEntryId());
+                    }
+                    // If the subscription is a Key_Shared subscription, we should to trigger message dispatch.
+                    if (subscription.getType() == SubType.Key_Shared) {
+                        dispatcher.markDeletePositionMoveForward();
+                    }
+                }
             }
             expirationCheckInProgress = FALSE;
             log.debug()

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

@@ -798,6 +798,10 @@ public void clearBacklogComplete(Object ctx) {
                             future.complete(null);
                         }
                     });
+                    Position mdPos = cursor.getMarkDeletedPosition();
+                    if (mdPos != null) {
+                        dispatcher.prunePendingAcksUpToPosition(mdPos.getLedgerId(), mdPos.getEntryId());
+                    }
                     dispatcher.afterAckMessages(null, ctx);
                 } else {
                     future.complete(null);
@@ -837,6 +841,10 @@ public void skipEntriesComplete(Object ctx) {
                                 .log("Skipped messages");
                         future.complete(null);
                         if (dispatcher != null) {
+                            Position mdPos = cursor.getMarkDeletedPosition();
+                            if (mdPos != null) {
+                                dispatcher.prunePendingAcksUpToPosition(mdPos.getLedgerId(), mdPos.getEntryId());
+                            }
                             dispatcher.afterAckMessages(null, ctx);
                         }
                     }