[Pipe] Distribute parser admission across processor workers (#19217) (#19231)

(cherry picked from commit 7666083)

Author: Caideyipi

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/event/common/tsfile/PipeTsFileInsertionEvent.java

@@ -1076,7 +1076,8 @@ private void releaseTsFileParserMemoryIfReserved() {
     synchronized (isTsFileParserMemoryReserved) {
       if (isTsFileParserMemoryReserved.compareAndSet(true, false)) {
         PipeDataNodeResourceManager.memory()
-            .releaseTsFileParserMemory(pipeName, creationTime, dataRegionId);
+            .releaseTsFileParserMemory(
+                pipeName, creationTime, dataRegionId, tsFileParserMemoryReservationKey);
       }
     }
   }
@@ -1339,7 +1340,8 @@ protected void finalizeResource() {
         synchronized (isTsFileParserMemoryReserved) {
           if (isTsFileParserMemoryReserved.compareAndSet(true, false)) {
             PipeDataNodeResourceManager.memory()
-                .releaseTsFileParserMemory(pipeName, creationTime, dataRegionId);
+                .releaseTsFileParserMemory(
+                    pipeName, creationTime, dataRegionId, tsFileParserMemoryReservationKey);
           }
         }
       } catch (final Exception e) {

iotdb-core/datanode/src/main/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManager.java

@@ -36,6 +36,7 @@
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
+import java.util.LinkedHashMap;
 import java.util.LinkedHashSet;
 import java.util.List;
 import java.util.Map;
@@ -74,6 +75,9 @@ public class PipeMemoryManager {
   private final Map<PipeIdentity, Integer> reservedTsFileParserCountByPipe = new HashMap<>();
   private final Map<PipeRegionIdentity, Integer> reservedTsFileParserCountByPipeRegion =
       new HashMap<>();
+  private final Map<PipeRegionIdentity, Map<TsFileParserMemoryReservation, Runnable>>
+      reservedTsFileParserWorkersByPipeRegion = new HashMap<>();
+  private final Map<Runnable, Integer> reservedTsFileParserCountByWorker = new HashMap<>();
   private final Map<PipeRegionIdentity, LinkedHashSet<TsFileParserMemoryReservation>>
       waitingTsFileParserRequestsByPipeRegion = new HashMap<>();
   private final Map<PipeIdentity, ArrayDeque<PipeRegionIdentity>>
@@ -242,8 +246,10 @@ public synchronized boolean tryReserveTsFileParserMemory(
             1, Math.min(globalLimit, PIPE_CONFIG.getPipeTsFileParserInFlightMaxNumPerPipeRegion()));
     final int reservedCountOfPipeRegion =
         reservedTsFileParserCountByPipeRegion.getOrDefault(pipeRegionIdentity, 0);
+    final Runnable worker = reservationKey.wakeUpCallback;
     if (reservedTsFileParserCount >= globalLimit
-        || reservedCountOfPipeRegion >= perPipeRegionLimit) {
+        || reservedCountOfPipeRegion >= perPipeRegionLimit
+        || (worker != null && reservedTsFileParserCountByWorker.containsKey(worker))) {
       PipeResourceMetrics.getInstance()
           .recordTsFileParserAdmissionReject(TsFileParserAdmissionRejectReason.CAP);
       notifyNextTsFileParserMemoryReservationInternal();
@@ -291,6 +297,12 @@ public synchronized boolean tryReserveTsFileParserMemory(
     reservedTsFileParserCount++;
     reservedTsFileParserCountByPipe.merge(pipeIdentity, 1, Integer::sum);
     reservedTsFileParserCountByPipeRegion.put(pipeRegionIdentity, reservedCountOfPipeRegion + 1);
+    reservedTsFileParserWorkersByPipeRegion
+        .computeIfAbsent(pipeRegionIdentity, key -> new LinkedHashMap<>())
+        .put(reservationKey, worker);
+    if (worker != null) {
+      reservedTsFileParserCountByWorker.merge(worker, 1, Integer::sum);
+    }
     notifyNextTsFileParserMemoryReservationInternal();
     return true;
   }
@@ -312,12 +324,23 @@ public synchronized void cancelTsFileParserMemoryReservation(
 
   public synchronized void releaseTsFileParserMemory(
       final String pipeName, final long creationTime, final String dataRegionId) {
+    releaseTsFileParserMemory(pipeName, creationTime, dataRegionId, null);
+  }
+
+  public synchronized void releaseTsFileParserMemory(
+      final String pipeName,
+      final long creationTime,
+      final String dataRegionId,
+      final TsFileParserMemoryReservation reservationKey) {
     final PipeIdentity pipeIdentity = new PipeIdentity(pipeName, creationTime);
     final PipeRegionIdentity pipeRegionIdentity =
         new PipeRegionIdentity(pipeIdentity, dataRegionId);
     final int reservedCountOfPipeRegion =
         reservedTsFileParserCountByPipeRegion.getOrDefault(pipeRegionIdentity, 0);
-    if (reservedCountOfPipeRegion <= 0) {
+    final Map<TsFileParserMemoryReservation, Runnable> workersOfPipeRegion =
+        reservedTsFileParserWorkersByPipeRegion.get(pipeRegionIdentity);
+    if (reservedCountOfPipeRegion <= 0
+        || (reservationKey != null && !workersOfPipeRegion.containsKey(reservationKey))) {
       LOGGER.warn(
           "Failed to release TsFile parser memory for pipe {} (creation time {}) in DataRegion {} because no reservation exists.",
           pipeName,
@@ -326,6 +349,19 @@ public synchronized void releaseTsFileParserMemory(
       return;
     }
 
+    final Runnable worker =
+        workersOfPipeRegion.remove(
+            reservationKey != null
+                ? reservationKey
+                : workersOfPipeRegion.keySet().iterator().next());
+    if (workersOfPipeRegion.isEmpty()) {
+      reservedTsFileParserWorkersByPipeRegion.remove(pipeRegionIdentity);
+    }
+    if (worker != null) {
+      reservedTsFileParserCountByWorker.computeIfPresent(
+          worker, (key, count) -> count == 1 ? null : count - 1);
+    }
+
     if (reservedCountOfPipeRegion == 1) {
       reservedTsFileParserCountByPipeRegion.remove(pipeRegionIdentity);
     } else {
@@ -406,6 +442,7 @@ private List<TsFileParserMemoryReservation> getEligibleTsFileParserMemoryReserva
     final List<TsFileParserMemoryReservation> eligibleRequests = new ArrayList<>();
     final Map<PipeRegionIdentity, Integer> selectedRequestsByPipeRegion = new HashMap<>();
     final Set<PipeIdentity> selectedPipes = new HashSet<>();
+    final Set<Runnable> selectedWorkers = new HashSet<>();
     PipeIdentity lastSelectedPipe = lastAdmittedWaitingTsFileParserPipe;
     // Expose one fair turn per available slot, instead of serializing all free slots behind a
     // single worker. Account for every earlier turn's memory even if its worker has not run yet.
@@ -429,6 +466,7 @@ private List<TsFileParserMemoryReservation> getEligibleTsFileParserMemoryReserva
               lastSelectedPipe,
               selectedRequestsByPipeRegion,
               selectedPipes,
+              selectedWorkers,
               true);
       final PipeRegionIdentity nextRegionWhenNoOtherPipeIsEligible =
           nextPipeRegion == null
@@ -438,6 +476,7 @@ private List<TsFileParserMemoryReservation> getEligibleTsFileParserMemoryReserva
                   lastSelectedPipe,
                   selectedRequestsByPipeRegion,
                   selectedPipes,
+                  selectedWorkers,
                   false)
               : nextPipeRegion;
       if (nextRegionWhenNoOtherPipeIsEligible == null) {
@@ -453,7 +492,11 @@ private List<TsFileParserMemoryReservation> getEligibleTsFileParserMemoryReserva
           i--) {
         requests.next();
       }
-      eligibleRequests.add(requests.next());
+      final TsFileParserMemoryReservation request = requests.next();
+      eligibleRequests.add(request);
+      if (request.wakeUpCallback != null) {
+        selectedWorkers.add(request.wakeUpCallback);
+      }
       selectedRequestsByPipeRegion.merge(nextRegionWhenNoOtherPipeIsEligible, 1, Integer::sum);
       selectedPipes.add(nextRegionWhenNoOtherPipeIsEligible.pipeIdentity);
       lastSelectedPipe = nextRegionWhenNoOtherPipeIsEligible.pipeIdentity;
@@ -467,6 +510,7 @@ private PipeRegionIdentity getNextEligibleTsFileParserPipeRegion(
       final PipeIdentity lastSelectedPipe,
       final Map<PipeRegionIdentity, Integer> selectedRequestsByPipeRegion,
       final Set<PipeIdentity> selectedPipes,
+      final Set<Runnable> selectedWorkers,
       final boolean requireUnselectedPipe) {
     PipeRegionIdentity firstEligiblePipeRegion = null;
     boolean hasVisitedLastAdmittedPipe = lastSelectedPipe == null;
@@ -499,6 +543,20 @@ private PipeRegionIdentity getNextEligibleTsFileParserPipeRegion(
             && reservedTsFileParserCountByPipeRegion.getOrDefault(pipeRegionIdentity, 0)
                     + selectedCount
                 < perPipeRegionLimit) {
+          final java.util.Iterator<TsFileParserMemoryReservation> requests =
+              waitingTsFileParserRequestsByPipeRegion.get(pipeRegionIdentity).iterator();
+          for (int i = 0; i < selectedCount; i++) {
+            requests.next();
+          }
+          final Runnable worker = requests.next().wakeUpCallback;
+          // A processor worker executes its subtasks serially. Leaving multiple fair turns for
+          // that worker strands free parser slots while other workers have runnable requests.
+          // Requests without a processor worker (for example subscriptions) remain independent.
+          if (worker != null
+              && (reservedTsFileParserCountByWorker.containsKey(worker)
+                  || selectedWorkers.contains(worker))) {
+            continue;
+          }
           if (eligiblePipeRegion == null) {
             eligiblePipeRegion = pipeRegionIdentity;
           }
@@ -579,7 +637,10 @@ public static final class TsFileParserMemoryReservation {
     // Accessed only under the PipeMemoryManager monitor while this request is queued.
     private long waitingStartTimeInNanos;
 
-    /** The callback must only wake the worker; it runs while the memory manager monitor is held. */
+    /**
+     * The callback must only wake the worker; it runs while the memory manager monitor is held.
+     * Requests on the same serial processor worker must share the same callback instance.
+     */
     public void setWakeUpCallback(final Runnable wakeUpCallback) {
       this.wakeUpCallback = wakeUpCallback;
     }

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/agent/task/subtask/processor/PipeProcessorSubtaskWorkerTest.java

@@ -66,6 +66,104 @@ public static void initializeSubtaskMock() {
     mock(PipeProcessorSubtask.class);
   }
 
+  @Test(timeout = 30000)
+  public void testRequestsOnOneWorkerDoNotStrandOtherWorkersParserSlots() throws Exception {
+    final CommonConfig config = CommonDescriptor.getInstance().getConfig();
+    final PipeMemoryManager manager = PipeDataNodeResourceManager.memory();
+    final long originalParserMemory = config.getPipeTsFileParserMemory();
+    final int originalGlobalLimit = config.getPipeTsFileParserInFlightMaxNum();
+    final int originalRegionLimit = config.getPipeTsFileParserInFlightMaxNumPerPipeRegion();
+    final List<PipeProcessorSubtaskWorker> workers = new ArrayList<>();
+    final List<TsFileParserMemoryReservation> requests = new ArrayList<>();
+    final boolean[] reservedBlockers = new boolean[4];
+    final CountDownLatch parsing = new CountDownLatch(4);
+    final CountDownLatch finishParsing = new CountDownLatch(1);
+    final ExecutorService executor = Executors.newFixedThreadPool(4);
+    try {
+      config.setPipeTsFileParserMemory(1);
+      config.setPipeTsFileParserInFlightMaxNum(4);
+      config.setPipeTsFileParserInFlightMaxNumPerPipeRegion(1);
+      for (int i = 0; i < 4; i++) {
+        reservedBlockers[i] =
+            manager.tryReserveTsFileParserMemory(
+                "blocker" + i, 0, "0", new TsFileParserMemoryReservation());
+        Assert.assertTrue(reservedBlockers[i]);
+        workers.add(new PipeProcessorSubtaskWorker(new LinkedHashSet<>()));
+      }
+      // Queue four pipes on the same serial worker before the other workers' requests. A fair
+      // window based only on pipe order leaves all four free slots waiting for that one worker.
+      for (int i = 0; i < 7; i++) {
+        final String pipeName = "skewed" + i;
+        final TsFileParserMemoryReservation request = new TsFileParserMemoryReservation();
+        requests.add(request);
+        final PipeProcessorSubtask subtask = createRunnableSubtask(pipeName);
+        final AtomicBoolean processed = new AtomicBoolean();
+        when(subtask.isWaitingForTsFileParserMemory())
+            .thenAnswer(ignored -> request.isWaitingForSignal());
+        doAnswer(
+                invocation -> {
+                  request.setWakeUpCallback(invocation.getArgument(0));
+                  return null;
+                })
+            .when(subtask)
+            .setWorkerWakeUpCallback(any());
+        when(subtask.call())
+            .thenAnswer(
+                ignored -> {
+                  if (processed.get()) {
+                    return false;
+                  }
+                  if (!manager.tryReserveTsFileParserMemory(pipeName, 0, "0", request)) {
+                    throw PipeProcessorSubtaskYieldException.parserNotAdmitted();
+                  }
+                  try {
+                    parsing.countDown();
+                    Assert.assertTrue(finishParsing.await(10, TimeUnit.SECONDS));
+                    processed.set(true);
+                    return true;
+                  } finally {
+                    manager.releaseTsFileParserMemory(pipeName, 0, "0");
+                  }
+                });
+        workers.get(i < 4 ? 0 : i - 3).schedule(subtask);
+        Assert.assertFalse(manager.tryReserveTsFileParserMemory(pipeName, 0, "0", request));
+      }
+      for (int i = 0; i < 4; i++) {
+        manager.releaseTsFileParserMemory("blocker" + i, 0, "0");
+        reservedBlockers[i] = false;
+      }
+      final List<Future<?>> rounds = new ArrayList<>();
+      for (final PipeProcessorSubtaskWorker worker : workers) {
+        rounds.add(executor.submit(worker::runSubtasks));
+      }
+      Assert.assertTrue(
+          "Free parser slots must run on distinct available workers",
+          parsing.await(5, TimeUnit.SECONDS));
+      Assert.assertEquals(4, manager.getReservedTsFileParserCount());
+      finishParsing.countDown();
+      for (final Future<?> round : rounds) {
+        round.get(5, TimeUnit.SECONDS);
+      }
+      Assert.assertEquals(0, manager.getReservedTsFileParserCount());
+      Assert.assertEquals(0, manager.getWaitingTsFileParserCount());
+    } finally {
+      finishParsing.countDown();
+      executor.shutdownNow();
+      Assert.assertTrue(executor.awaitTermination(5, TimeUnit.SECONDS));
+      for (int i = 0; i < requests.size(); i++) {
+        manager.cancelTsFileParserMemoryReservation("skewed" + i, 0, "0", requests.get(i));
+      }
+      for (int i = 0; i < reservedBlockers.length; i++) {
+        if (reservedBlockers[i]) {
+          manager.releaseTsFileParserMemory("blocker" + i, 0, "0");
+        }
+      }
+      config.setPipeTsFileParserMemory(originalParserMemory);
+      config.setPipeTsFileParserInFlightMaxNum(originalGlobalLimit);
+      config.setPipeTsFileParserInFlightMaxNumPerPipeRegion(originalRegionLimit);
+    }
+  }
+
   @Test(timeout = 30000)
   public void testOneHundredWaitingPipesDoNotWaitForWorkerPolling() throws Exception {
     final CommonConfig config = CommonDescriptor.getInstance().getConfig();

iotdb-core/datanode/src/test/java/org/apache/iotdb/db/pipe/resource/memory/PipeMemoryManagerTest.java

@@ -143,6 +143,70 @@ public void testFreeSlotsAreAvailableWithoutWaitingForTheFirstWorker() {
     Assert.assertEquals(96, memoryManager.getWaitingTsFileParserCount());
   }
 
+  @Test
+  public void testReleasingParserOnlyFreesItsOwnWorker() {
+    commonConfig.setPipeTsFileParserInFlightMaxNum(4);
+    commonConfig.setPipeTsFileParserInFlightMaxNumPerPipeRegion(2);
+    final PipeMemoryManager manager = createTestMemoryManager();
+    final Reservation firstWorker = new Reservation("pipe", 0);
+    final Reservation secondWorker = new Reservation("pipe", 0);
+    final Reservation waitingOnFirstWorker = new Reservation("otherPipe", 1);
+    final Reservation waitingOnSecondWorker = new Reservation("otherPipe", 2);
+    final Runnable firstWakeUp = new AtomicInteger()::incrementAndGet;
+    final Runnable secondWakeUp = new AtomicInteger()::incrementAndGet;
+    firstWorker.key.setWakeUpCallback(firstWakeUp);
+    waitingOnFirstWorker.key.setWakeUpCallback(firstWakeUp);
+    secondWorker.key.setWakeUpCallback(secondWakeUp);
+    waitingOnSecondWorker.key.setWakeUpCallback(secondWakeUp);
+    final List<Reservation> testReservations =
+        List.of(firstWorker, secondWorker, waitingOnFirstWorker, waitingOnSecondWorker);
+    try {
+      Assert.assertTrue(tryAcquire(manager, firstWorker));
+      Assert.assertTrue(tryAcquire(manager, secondWorker));
+      Assert.assertFalse(tryAcquire(manager, waitingOnFirstWorker));
+      Assert.assertFalse(tryAcquire(manager, waitingOnSecondWorker));
+
+      // The two parsers share a pipe/region but finish in reverse order. Releasing the second
+      // reservation must wake only its worker, rather than freeing the first worker's slot.
+      manager.releaseTsFileParserMemory(
+          secondWorker.pipeName,
+          secondWorker.creationTime,
+          secondWorker.dataRegionId,
+          secondWorker.key);
+      secondWorker.acquired = false;
+      Assert.assertTrue(waitingOnFirstWorker.key.isWaitingForSignal());
+      Assert.assertFalse(waitingOnSecondWorker.key.isWaitingForSignal());
+      Assert.assertTrue(tryAcquire(manager, waitingOnSecondWorker));
+      Assert.assertFalse(tryAcquire(manager, waitingOnFirstWorker));
+
+      manager.releaseTsFileParserMemory(
+          firstWorker.pipeName,
+          firstWorker.creationTime,
+          firstWorker.dataRegionId,
+          firstWorker.key);
+      firstWorker.acquired = false;
+      Assert.assertTrue(tryAcquire(manager, waitingOnFirstWorker));
+      Assert.assertEquals(2, manager.getReservedTsFileParserCount());
+    } finally {
+      for (final Reservation reservation : testReservations) {
+        manager.cancelTsFileParserMemoryReservation(
+            reservation.pipeName,
+            reservation.creationTime,
+            reservation.dataRegionId,
+            reservation.key);
+        if (reservation.acquired) {
+          manager.releaseTsFileParserMemory(
+              reservation.pipeName,
+              reservation.creationTime,
+              reservation.dataRegionId,
+              reservation.key);
+        }
+      }
+    }
+    Assert.assertEquals(0, manager.getReservedTsFileParserCount());
+    Assert.assertEquals(0, manager.getWaitingTsFileParserCount());
+  }
+
   @Test
   public void testAdmissionWindowRespectsPipeFairnessAcrossRegions() {
     commonConfig.setPipeTsFileParserInFlightMaxNum(2);