Add poller balancer

Author: Quinn-With-Two-Ns

temporal-sdk/src/main/java/io/temporal/internal/worker/AsyncPoller.java

@@ -116,7 +116,7 @@ public boolean start() {
               });
       PollQueueTask pollQueue =
           new PollQueueTask(asyncTaskPoller, pollerSemaphore, pollScaleReportHandle);
-      balancer.addPoller(pollQueue);
+      balancer.addPoller(asyncTaskPoller.getLabel());
       exec.execute(pollQueue);
       exec.scheduleAtFixedRate(pollScaleReportHandle, 0, 100, TimeUnit.MILLISECONDS);
     }
@@ -129,6 +129,7 @@ public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean
         .thenApply(
             (f) -> {
               for (PollTaskAsync<T> asyncTaskPoller : asyncTaskPollers) {
+                log.info("Shutting down async poller: {} !!!!", asyncTaskPoller);
                 try {
                   asyncTaskPoller.cancel(new RuntimeException("Shutting down poller"));
                 } catch (Throwable e) {
@@ -152,6 +153,10 @@ public interface PollTaskAsync<TT> {
     default void cancel(Throwable cause) {
       // no-op
     }
+
+    default String getLabel() {
+      return "AsyncPollTask";
+    }
   }
 
   class PollQueueTask implements Runnable {
@@ -179,171 +184,211 @@ class PollQueueTask implements Runnable {
 
     @Override
     public void run() {
-      // Permit to reserve a slot for the poll request.
-      SlotPermit permit = null;
-      // Flag to check if pollerSemaphore was acquired, if so, we need to release it if
-      // an exception occurs.
-      boolean pollerSemaphoreAcquired = false;
-      // Flag to check if poll request was made, if not, we need to release the slot
-      // permit and pollerSemaphore in this method.
-      boolean pollRequestMade = false;
-      try {
-        long throttleMs = pollBackoffThrottler.getSleepTime();
-        if (throttleMs > 0) {
-          Thread.sleep(throttleMs);
-        }
-        if (pollRateThrottler != null) {
-          pollRateThrottler.throttle();
-        }
+      int slotsReserved = 0;
+      while (!abort) {
+        // Permit to reserve a slot for the poll request.
+        SlotPermit permit = null;
+        // Flag to check if pollerSemaphore was acquired, if so, we need to release it if
+        // an exception occurs.
+        boolean pollerSemaphoreAcquired = false;
+        // Flag to check if poll request was made, if not, we need to release the slot
+        // permit and pollerSemaphore in this method.
+        boolean pollRequestMade = false;
+        try {
+          long throttleMs = pollBackoffThrottler.getSleepTime();
+          if (throttleMs > 0) {
+            Thread.sleep(throttleMs);
+          }
+          if (pollRateThrottler != null) {
+            pollRateThrottler.throttle();
+          }
 
-        CountDownLatch suspender = suspendLatch.get();
-        if (suspender != null) {
-          if (log.isDebugEnabled()) {
-            log.debug("poll task suspending latchCount=" + suspender.getCount());
+          CountDownLatch suspender = suspendLatch.get();
+          if (suspender != null) {
+            if (log.isDebugEnabled()) {
+              log.debug("poll task suspending latchCount=" + suspender.getCount());
+            }
+            suspender.await();
           }
-          suspender.await();
-        }
 
-        if (shouldTerminate()) {
-          return;
-        }
+          if (shouldTerminate()) {
+            return;
+          }
 
-        balancer.balance(this);
-        if (shouldTerminate()) {
-          return;
-        }
-        // Reserve a slot for the poll request
-        SlotSupplierFuture future;
-        try {
-          future = slotSupplier.reserveSlot(slotReservationData);
-        } catch (Exception e) {
-          log.warn("Error while trying to reserve a slot", e.getCause());
-          return;
-        }
-        permit = BasePoller.getSlotPermitAndHandleInterrupts(future, slotSupplier);
-        if (permit == null || shouldTerminate()) {
-          return;
-        }
+          balancer.balance(asyncTaskPoller.getLabel());
+          if (shouldTerminate()) {
+            return;
+          }
+          // Reserve a slot for the poll request
+          SlotSupplierFuture future;
+          try {
+            future = slotSupplier.reserveSlot(slotReservationData);
+          } catch (Exception e) {
+            log.warn("Error while trying to reserve a slot", e.getCause());
+            return;
+          }
+          permit = BasePoller.getSlotPermitAndHandleInterrupts(future, slotSupplier);
+          if (permit == null || shouldTerminate()) {
+            if (permit != null) {
+              slotsReserved++;
+            }
+            return;
+          }
+          slotsReserved++;
 
-        pollerSemaphore.acquire();
-        pollerSemaphoreAcquired = true;
-        if (shouldTerminate()) {
-          return;
-        }
-        workerMetricsScope.counter(MetricsType.POLLER_START_COUNTER).inc(1);
+          pollerSemaphore.acquire();
+          pollerSemaphoreAcquired = true;
+          if (shouldTerminate()) {
+            return;
+          }
+          workerMetricsScope.counter(MetricsType.POLLER_START_COUNTER).inc(1);
 
-        SlotPermit finalPermit = permit;
-        CompletableFuture<T> pollRequest = asyncTaskPoller.poll(permit);
-        // Mark that we have made a poll request
-        pollRequestMade = true;
-        balancer.startPoll(this);
+          SlotPermit finalPermit = permit;
+          CompletableFuture<T> pollRequest = asyncTaskPoller.poll(permit);
+          // Mark that we have made a poll request
+          pollRequestMade = true;
+          balancer.startPoll(asyncTaskPoller.getLabel());
 
-        pollRequest.handle(
-            (task, e) -> {
-              balancer.endPoll(this);
-              if (e instanceof CompletionException) {
-                e = e.getCause();
-              }
+          pollRequest
+              .handle(
+                  (task, e) -> {
+                    balancer.endPoll(asyncTaskPoller.getLabel());
+                    if (e instanceof CompletionException) {
+                      e = e.getCause();
+                    }
+                    pollerSemaphore.release();
+                    pollScaleReportHandle.report(task, e);
+                    if (e != null) {
+                      // uncaughtExceptionHandler.uncaughtException(Thread.currentThread(), e);
+                      pollBackoffThrottler.failure(
+                          (e instanceof StatusRuntimeException)
+                              ? ((StatusRuntimeException) e).getStatus().getCode()
+                              : Status.Code.UNKNOWN);
+                      slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), finalPermit);
+                      return null;
+                    }
+                    log.trace("$$$$$ Picked up task: {}", task);
+                    if (task != null) {
+                      taskExecutor.process(task);
+                    } else {
+                      slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), finalPermit);
+                    }
+                    pollBackoffThrottler.success();
+                    return null;
+                  })
+              .exceptionally(
+                  throwable -> {
+                    log.error("Error while trying to poll task", throwable);
+                    return null;
+                  });
+        } catch (PollTaskAsyncAbort ab) {
+          abort = true;
+        } catch (Throwable e) {
+          if (e instanceof InterruptedException) {
+            // we restore the flag here, so it can be checked and processed (with exit) in finally.
+            Thread.currentThread().interrupt();
+          }
+          uncaughtExceptionHandler.uncaughtException(Thread.currentThread(), e);
+        } finally {
+          // release the slot if it was acquired, but a poll request was not made
+          if (!pollRequestMade) {
+            if (permit != null) {
+              slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), permit);
+            }
+            if (pollerSemaphoreAcquired) {
               pollerSemaphore.release();
-              pollScaleReportHandle.report(task, e);
-              if (e != null) {
-                uncaughtExceptionHandler.uncaughtException(Thread.currentThread(), e);
-                pollBackoffThrottler.failure(
-                    (e instanceof StatusRuntimeException)
-                        ? ((StatusRuntimeException) e).getStatus().getCode()
-                        : Status.Code.UNKNOWN);
-                slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), finalPermit);
-                return null;
-              }
-              log.trace("$$$$$ Picked up task: {}", task);
-              if (task != null) {
-                taskExecutor.process(task);
-              } else {
-                slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), finalPermit);
-              }
-              pollBackoffThrottler.success();
-              return null;
-            });
-      } catch (PollTaskAsyncAbort ab) {
-        abort = true;
-      } catch (Throwable e) {
-        if (e instanceof InterruptedException) {
-          // we restore the flag here, so it can be checked and processed (with exit) in finally.
-          Thread.currentThread().interrupt();
-        }
-        uncaughtExceptionHandler.uncaughtException(Thread.currentThread(), e);
-      } finally {
-        // release the slot if it was acquired, but a poll request was not made
-        if (!pollRequestMade && permit != null) {
-          slotSupplier.releaseSlot(SlotReleaseReason.taskComplete(), permit);
-        }
-        if (!pollRequestMade && pollerSemaphoreAcquired) {
-          pollerSemaphore.release();
-        }
+            }
+          }
 
-        if (!shouldTerminate() && !abort) {
-          // Resubmit itself back to pollExecutor
-          pollExecutor.execute(this);
-        } else {
-          balancer.removePoller(this);
-          log.info("poll loop is terminated: {}", asyncTaskPoller.getClass().getSimpleName());
+          if (shouldTerminate()) {
+            log.info(
+                "ABORT {} permit {}, pollRequestMade {}",
+                asyncTaskPoller.getLabel(),
+                permit,
+                pollRequestMade);
+            balancer.removePoller(asyncTaskPoller.getLabel());
+            abort = true;
+          }
         }
       }
+      log.info(
+          "poll loop is terminated: {}, slotsReserved {}",
+          asyncTaskPoller.getLabel(),
+          slotsReserved);
     }
   }
 
   class PollQueueBalancer {
-    Map<PollQueueTask, AtomicInteger> taskCounts = new HashMap<>();
+    Map<String, AtomicInteger> taskCounts = new HashMap<>();
     private final Lock balancerLock = new ReentrantLock();
     private final Condition balancerCondition = balancerLock.newCondition();
 
-    void startPoll(PollQueueTask task) {
+    void startPoll(String pollerName) {
       balancerLock.lock();
-      int currentPolls = taskCounts.get(task).addAndGet(1);
+      int currentPolls = taskCounts.get(pollerName).addAndGet(1);
       if (currentPolls == 1) {
         balancerCondition.signalAll();
       }
       balancerLock.unlock();
     }
 
-    void endPoll(PollQueueTask task) {
+    void endPoll(String pollerName) {
       balancerLock.lock();
-      int currentPolls = taskCounts.get(task).addAndGet(-1);
+      if (!taskCounts.containsKey(pollerName)) {
+        log.info("endPoll has null key for tasks {}", pollerName);
+        balancerLock.unlock();
+        return;
+      }
+      int currentPolls = taskCounts.get(pollerName).addAndGet(-1);
       if (currentPolls == 0) {
         balancerCondition.signalAll();
       }
       balancerLock.unlock();
     }
 
-    void addPoller(PollQueueTask task) {
+    void addPoller(String pollerName) {
       balancerLock.lock();
-      taskCounts.putIfAbsent(task, new AtomicInteger(0));
+      log.info("addPoller {}", pollerName);
+      taskCounts.put(pollerName, new AtomicInteger(0));
+      if (!taskCounts.containsKey(pollerName)) {
+        log.info("????? has null key for task {}", pollerName);
+        balancerLock.unlock();
+        return;
+      }
       balancerCondition.signalAll();
       balancerLock.unlock();
     }
 
-    void removePoller(PollQueueTask task) {
+    void removePoller(String pollerName) {
       balancerLock.lock();
-      taskCounts.remove(task);
+      taskCounts.remove(pollerName);
       balancerCondition.signalAll();
       balancerLock.unlock();
     }
 
     /** Ensure that at least one poller is running for each task. */
-    void balance(PollQueueTask t) throws InterruptedException {
+    void balance(String p) throws InterruptedException {
       while (!shouldTerminate()) {
         balancerLock.lock();
         // If this poller has no tasks then we can unblock immediately
-        if (taskCounts.get(t).get() == 0) {
+        if (taskCounts.get(p).get() == 0) {
           balancerLock.unlock();
           return;
         }
         // Check if all tasks have at least one poll request
-        for (PollQueueTask task : taskCounts.keySet()) {
-          if (taskCounts.get(task).get() == 0) {
-            balancerCondition.await();
+        boolean allOtherTasksHavePolls = true;
+        for (String task : taskCounts.keySet()) {
+          if (!Objects.equals(task, p) && taskCounts.get(task).get() == 0) {
+            allOtherTasksHavePolls = false;
+            break;
           }
         }
+        if (!allOtherTasksHavePolls) {
+          balancerCondition.await();
+        } else {
+          balancerLock.unlock();
+          return;
+        }
       }
     }
   }

temporal-sdk/src/main/java/io/temporal/internal/worker/AsyncWorkflowPollTask.java

@@ -173,4 +173,9 @@ public void disableNormalPoll() {
     }
     shutdown.set(true);
   }
+
+  @Override
+  public String getLabel() {
+    return stickyPoller ? "StickyWorkflowPollTask" : "NormalWorkflowPollTask";
+  }
 }

temporal-sdk/src/main/java/io/temporal/internal/worker/ShutdownManager.java

@@ -71,6 +71,7 @@ public CompletableFuture<Void> waitForStickyQueueBalancer(
     balancer.disableNormalPoll();
     scheduledExecutorService.schedule(
         () -> {
+          log.info("DRAIN TIMEOUT PASSED!!!!");
           future.complete(null);
         },
         timeout.toMillis(),
@@ -296,12 +297,18 @@ public SlotSupplierDelayShutdown(
 
     @Override
     boolean isTerminated() {
+      if (slotSupplier.getIssuedSlots() > 0) {
+        log.info(
+            "Wait for release of slots of {} is in progress {}",
+            name,
+            slotSupplier.getIssuedSlots());
+      }
       return slotSupplier.getIssuedSlots() == 0;
     }
 
     @Override
     void onSlowTermination() {
-      log.warn("Wait for release of slots of {} takes a long time", name);
+      log.info("Wait for release of slots of {} takes a long time", name);
     }
 
     @Override