Return on LA queue thread interrupts

Author: Sushisource

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

@@ -79,48 +79,63 @@ static final class QueuedLARequest {
   }
 
   private void processQueue() {
-    while (running || !requestQueue.isEmpty()) {
-      SlotPermit slotPermit = null;
-      QueuedLARequest request = null;
-      try {
-        request = requestQueue.take();
-
-        SlotSupplierFuture future = slotSupplier.reserveSlot(request.data);
+    try {
+      while (running || !requestQueue.isEmpty()) {
+        SlotPermit slotPermit = null;
+        QueuedLARequest request = null;
         try {
-          slotPermit = future.get();
-        } catch (InterruptedException e) {
-          SlotPermit maybePermitAnyway = future.abortReservation();
-          if (maybePermitAnyway != null) {
-            slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), maybePermitAnyway);
+          // Use poll with timeout instead of take() to avoid being blocked indefinitely
+          request = requestQueue.poll(100, TimeUnit.MILLISECONDS);
+          if (request == null) {
+            continue;
           }
+
+          SlotSupplierFuture future = slotSupplier.reserveSlot(request.data);
+          try {
+            slotPermit = future.get();
+          } catch (InterruptedException e) {
+            SlotPermit maybePermitAnyway = future.abortReservation();
+            if (maybePermitAnyway != null) {
+              slotSupplier.releaseSlot(SlotReleaseReason.neverUsed(), maybePermitAnyway);
+            }
+            Thread.currentThread().interrupt();
+            return;
+          } catch (ExecutionException e) {
+            log.error(
+                "Error reserving local activity slot, dropped activity id {}",
+                request.task.getActivityId(),
+                e);
+            continue;
+          }
+
+          request.task.getExecutionContext().setPermit(slotPermit);
+          afterReservedCallback.apply(request.task);
+        } catch (InterruptedException e) {
           Thread.currentThread().interrupt();
           return;
-        } catch (ExecutionException e) {
-          log.error(
-              "Error reserving local activity slot, dropped activity id {}",
-              request.task.getActivityId(),
-              e);
-          continue;
-        }
-
-        request.task.getExecutionContext().setPermit(slotPermit);
-        afterReservedCallback.apply(request.task);
-      } catch (InterruptedException e) {
-        Thread.currentThread().interrupt();
-      } catch (Throwable e) {
-        // Fail the workflow task if something went wrong executing the local activity (at the
-        // executor level, otherwise, the LA handler itself should be handling errors)
-        log.error("Unexpected error submitting local activity task to worker", e);
-        if (slotPermit != null) {
-          slotSupplier.releaseSlot(SlotReleaseReason.error(new RuntimeException(e)), slotPermit);
-        }
-        if (request != null) {
-          LocalActivityExecutionContext executionContext = request.task.getExecutionContext();
-          executionContext.callback(
-              LocalActivityResult.processingFailed(
-                  executionContext.getActivityId(), request.task.getAttemptTask().getAttempt(), e));
+        } catch (Throwable e) {
+          // Fail the workflow task if something went wrong executing the local activity (at the
+          // executor level, otherwise, the LA handler itself should be handling errors)
+          log.error("Unexpected error submitting local activity task to worker", e);
+          if (slotPermit != null) {
+            slotSupplier.releaseSlot(SlotReleaseReason.error(new RuntimeException(e)), slotPermit);
+          }
+          if (request != null) {
+            LocalActivityExecutionContext executionContext = request.task.getExecutionContext();
+            executionContext.callback(
+                LocalActivityResult.processingFailed(
+                    executionContext.getActivityId(),
+                    request.task.getAttemptTask().getAttempt(),
+                    e));
+          }
+          if (e.getCause() instanceof InterruptedException) {
+            Thread.currentThread().interrupt();
+            return;
+          }
         }
       }
+    } finally {
+      log.debug("LocalActivitySlotSupplierQueue thread exiting");
     }
   }
 
@@ -170,11 +185,10 @@ public boolean isTerminated() {
   @Override
   public CompletableFuture<Void> shutdown(ShutdownManager shutdownManager, boolean interruptTasks) {
     running = false;
-    if (requestQueue.isEmpty()) {
-      // Just interrupt the thread, so that if we're waiting on blocking take the thread will
-      // be interrupted and exit. Otherwise the loop will exit once the queue is empty.
-      queueThreadService.shutdownNow();
-    }
+
+    queueThreadService.shutdownNow();
+
+    log.debug("LocalActivitySlotSupplierQueue shutdown initiated");
 
     return interruptTasks
         ? shutdownManager.shutdownExecutorNowUntimed(
@@ -190,6 +204,15 @@ public void awaitTermination(long timeout, TimeUnit unit) {
       // timeout duration if no task was ever submitted.
       return;
     }
-    ShutdownManager.awaitTermination(queueThreadService, unit.toMillis(timeout));
+
+    try {
+      boolean terminated = queueThreadService.awaitTermination(timeout, unit);
+      if (!terminated) {
+        log.warn("LocalActivitySlotSupplierQueue did not terminate within the timeout period");
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+      log.warn("Interrupted while waiting for LocalActivitySlotSupplierQueue termination");
+    }
   }
 }