Fix Virtual Threads unbounded (#26013)

* Fix Virtual Threads unbounded

* Bound all AsyncService paths with semaphore-wrapped executor

---------

Co-authored-by: Adrià Manero <adria.estivill@getcollate.io>

Author: 2 people

openmetadata-service/src/main/java/org/openmetadata/service/OpenMetadataApplication.java

@@ -167,6 +167,7 @@
 import org.openmetadata.service.socket.WebSocketManager;
 import org.openmetadata.service.swagger.SwaggerBundle;
 import org.openmetadata.service.swagger.SwaggerBundleConfiguration;
+import org.openmetadata.service.util.AsyncService;
 import org.openmetadata.service.util.CustomParameterNameProvider;
 import org.openmetadata.service.util.incidentSeverityClassifier.IncidentSeverityClassifierInterface;
 import org.quartz.SchedulerException;
@@ -1099,6 +1100,7 @@ public void start() {
     public void stop() throws InterruptedException, SchedulerException {
       LOG.info("Cache with Id Stats {}", EntityRepository.CACHE_WITH_ID.stats());
       LOG.info("Cache with name Stats {}", EntityRepository.CACHE_WITH_NAME.stats());
+      AsyncService.getInstance().shutdown();
       EventPubSub.shutdown();
       AppScheduler.shutDown();
       EventSubscriptionScheduler.shutDown();

openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedJobParticipant.java

@@ -63,6 +63,8 @@ public class DistributedJobParticipant implements Managed {
    */
   @Getter private UUID currentJobId;
 
+  private volatile Thread participantThread;
+
   public DistributedJobParticipant(
       CollectionDAO collectionDAO,
       SearchRepository searchRepository,
@@ -116,6 +118,18 @@ public void start() {
   @Override
   public void stop() {
     if (running.compareAndSet(true, false)) {
+      Thread thread = participantThread;
+      if (thread != null) {
+        thread.interrupt();
+        try {
+          thread.join(10_000);
+          if (thread.isAlive()) {
+            LOG.warn("Participant thread did not terminate within 10s after interrupt");
+          }
+        } catch (InterruptedException e) {
+          Thread.currentThread().interrupt();
+        }
+      }
       notifier.stop();
       LOG.info("Stopped distributed job participant on server {}", serverId);
     }
@@ -190,22 +204,22 @@ private void joinAndProcessJob(SearchIndexJob job) {
       pollingNotifier.setParticipating(true);
     }
 
-    Thread.ofVirtual()
-        .name("job-participant-" + job.getId().toString().substring(0, 8))
-        .start(
-            () -> {
-              try {
-                processJobPartitions(job);
-              } finally {
-                participating.set(false);
-                currentJobId = null;
-
-                // Reset polling notifier to idle interval
-                if (notifier instanceof PollingJobNotifier pollingNotifier) {
-                  pollingNotifier.setParticipating(false);
-                }
-              }
-            });
+    participantThread =
+        Thread.ofVirtual()
+            .name("job-participant-" + job.getId().toString().substring(0, 8))
+            .start(
+                () -> {
+                  try {
+                    processJobPartitions(job);
+                  } finally {
+                    currentJobId = null;
+                    if (notifier instanceof PollingJobNotifier pollingNotifier) {
+                      pollingNotifier.setParticipating(false);
+                    }
+                    participating.set(false);
+                    participantThread = null;
+                  }
+                });
   }
 
   /** Process partitions for a job. */

openmetadata-service/src/main/java/org/openmetadata/service/apps/bundles/searchIndex/distributed/DistributedSearchIndexExecutor.java

@@ -105,6 +105,7 @@ public static boolean isCoordinatingJob(UUID jobId) {
   private final List<PartitionWorker> activeWorkers = new ArrayList<>();
   private Thread lockRefreshThread;
   private Thread partitionHeartbeatThread;
+  private Thread staleReclaimerThread;
 
   // App context for WebSocket broadcasts
   private UUID appId;
@@ -442,8 +443,7 @@ public ExecutionResult execute(
           });
     }
 
-    // Start stale partition reclaimer in a separate thread
-    Thread staleReclaimer =
+    staleReclaimerThread =
         Thread.ofVirtual()
             .name("stale-reclaimer-" + jobId.toString().substring(0, 8))
             .start(() -> runStaleReclaimerLoop(jobId));
@@ -462,14 +462,9 @@ public ExecutionResult execute(
       Thread.currentThread().interrupt();
       LOG.warn("Execution interrupted for job {}", jobId);
     } finally {
-      // Stop all background threads
-      staleReclaimer.interrupt();
-      if (lockRefreshThread != null) {
-        lockRefreshThread.interrupt();
-      }
-      if (partitionHeartbeatThread != null) {
-        partitionHeartbeatThread.interrupt();
-      }
+      interruptAndJoin(staleReclaimerThread, "stale-reclaimer");
+      interruptAndJoin(lockRefreshThread, "lock-refresh");
+      interruptAndJoin(partitionHeartbeatThread, "partition-heartbeat");
 
       // Shutdown executor
       workerExecutor.shutdown();
@@ -849,6 +844,26 @@ private void markJobAsFailedDueToLostLock(UUID jobId) {
     }
   }
 
+  private static final long THREAD_JOIN_TIMEOUT_MS = 5000;
+
+  private void interruptAndJoin(Thread thread, String name) {
+    if (thread == null) {
+      return;
+    }
+    thread.interrupt();
+    try {
+      thread.join(THREAD_JOIN_TIMEOUT_MS);
+      if (thread.isAlive()) {
+        LOG.warn(
+            "Thread {} did not terminate within {}ms after interrupt",
+            name,
+            THREAD_JOIN_TIMEOUT_MS);
+      }
+    } catch (InterruptedException e) {
+      Thread.currentThread().interrupt();
+    }
+  }
+
   /**
    * Request to stop the current job execution.
    */

openmetadata-service/src/main/java/org/openmetadata/service/audit/AuditLogRepository.java

@@ -18,6 +18,7 @@
 import org.openmetadata.schema.utils.ResultList;
 import org.openmetadata.service.Entity;
 import org.openmetadata.service.jdbi3.CollectionDAO;
+import org.openmetadata.service.util.AsyncService;
 import org.openmetadata.service.util.FullyQualifiedName;
 import org.openmetadata.service.util.RestUtil;
 
@@ -122,18 +123,19 @@ public void write(ChangeEvent changeEvent, boolean isBot) {
    * populated. Runs asynchronously using a virtual thread to avoid blocking the caller.
    */
   public void writeAuthEvent(EventType eventType, String userName, UUID userId) {
-    Thread.startVirtualThread(
-        () -> {
-          ChangeEvent changeEvent =
-              new ChangeEvent()
-                  .withId(UUID.randomUUID())
-                  .withEventType(eventType)
-                  .withEntityType(Entity.USER)
-                  .withEntityId(userId)
-                  .withUserName(userName)
-                  .withTimestamp(System.currentTimeMillis());
-          write(changeEvent);
-        });
+    AsyncService.getInstance()
+        .execute(
+            () -> {
+              ChangeEvent changeEvent =
+                  new ChangeEvent()
+                      .withId(UUID.randomUUID())
+                      .withEventType(eventType)
+                      .withEntityType(Entity.USER)
+                      .withEntityId(userId)
+                      .withUserName(userName)
+                      .withTimestamp(System.currentTimeMillis());
+              write(changeEvent);
+            });
   }
 
   /** Determine actor type from username pattern - agents, bots, or regular users. */

openmetadata-service/src/main/java/org/openmetadata/service/util/AsyncService.java

@@ -1,25 +1,66 @@
 package org.openmetadata.service.util;
 
+import java.util.List;
+import java.util.concurrent.AbstractExecutorService;
+import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.Executors;
+import java.util.concurrent.Semaphore;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.TimeoutException;
 import java.util.function.Supplier;
+import lombok.Getter;
 import lombok.extern.slf4j.Slf4j;
+import org.openmetadata.service.OpenMetadataApplicationConfigHolder;
 
 @Slf4j
 public class AsyncService {
   private static AsyncService instance;
   private final ExecutorService executorService;
+  private final Semaphore concurrencyLimiter;
+  @Getter private final int maxConcurrency;
 
-  // Default retry configuration
   private static final int DEFAULT_MAX_RETRIES = 3;
-  private static final long DEFAULT_INITIAL_RETRY_DELAY_MS = 1000; // 1 second
-  private static final long DEFAULT_OPERATION_TIMEOUT_SECONDS = 60; // 60 seconds
+  private static final long DEFAULT_INITIAL_RETRY_DELAY_MS = 1000;
+  private static final long DEFAULT_OPERATION_TIMEOUT_SECONDS = 60;
+
+  private static final long SHUTDOWN_TIMEOUT_SECONDS = 30;
 
   private AsyncService() {
-    executorService = Executors.newVirtualThreadPerTaskExecutor();
+    maxConcurrency = resolveMaxConcurrency();
+    concurrencyLimiter = new Semaphore(maxConcurrency);
+    executorService =
+        new BoundedExecutorService(Executors.newVirtualThreadPerTaskExecutor(), concurrencyLimiter);
+    LOG.info("AsyncService initialized with max concurrency: {}", maxConcurrency);
+  }
+
+  private static int resolveMaxConcurrency() {
+    String env = System.getenv("ASYNC_SERVICE_MAX_CONCURRENCY");
+    if (env != null) {
+      try {
+        int value = Integer.parseInt(env.trim());
+        if (value > 0) {
+          return value;
+        }
+      } catch (NumberFormatException ignored) {
+      }
+    }
+    int cpuBudget = Runtime.getRuntime().availableProcessors() * 2;
+    try {
+      if (OpenMetadataApplicationConfigHolder.isInitialized()) {
+        int poolSize =
+            OpenMetadataApplicationConfigHolder.getInstance().getDataSourceFactory().getMaxSize();
+        if (poolSize > 0) {
+          return Math.max(4, Math.min(cpuBudget, poolSize / 3));
+        }
+      }
+    } catch (Exception e) {
+      LOG.warn(
+          "Could not determine database pool size, using CPU-based concurrency budget: {}",
+          e.getMessage());
+    }
+    return Math.max(4, cpuBudget);
   }
 
   public static synchronized AsyncService getInstance() {
@@ -33,9 +74,38 @@ public ExecutorService getExecutorService() {
     return executorService;
   }
 
-  // Optionally, provide a method to shut down the executor service
+  public void execute(Runnable task) {
+    executorService.execute(task);
+  }
+
+  public <T> CompletableFuture<T> submit(Callable<T> task) {
+    return CompletableFuture.supplyAsync(
+        () -> {
+          try {
+            return task.call();
+          } catch (RuntimeException e) {
+            throw e;
+          } catch (Exception e) {
+            throw new RuntimeException(e);
+          }
+        },
+        executorService);
+  }
+
   public void shutdown() {
+    LOG.info("Shutting down AsyncService executor (max concurrency: {})", maxConcurrency);
     executorService.shutdown();
+    try {
+      if (!executorService.awaitTermination(SHUTDOWN_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
+        LOG.warn(
+            "AsyncService executor did not terminate within {}s, forcing shutdown",
+            SHUTDOWN_TIMEOUT_SECONDS);
+        executorService.shutdownNow();
+      }
+    } catch (InterruptedException e) {
+      executorService.shutdownNow();
+      Thread.currentThread().interrupt();
+    }
   }
 
   /**
@@ -161,4 +231,63 @@ private static <T> T executeWithRetry(
     throw new RuntimeException(
         String.format("Failed to %s %s", operationName.toLowerCase(), context), lastException);
   }
+
+  /**
+   * ExecutorService wrapper that enforces concurrency limits via a semaphore. Every task submitted
+   * through any method (execute, submit, invokeAll, invokeAny) acquires a permit before running and
+   * releases it on completion. This ensures ALL callers — including those using getExecutorService()
+   * directly — are bounded.
+   */
+  private static class BoundedExecutorService extends AbstractExecutorService {
+    private final ExecutorService delegate;
+    private final Semaphore semaphore;
+
+    BoundedExecutorService(ExecutorService delegate, Semaphore semaphore) {
+      this.delegate = delegate;
+      this.semaphore = semaphore;
+    }
+
+    @Override
+    public void execute(Runnable command) {
+      delegate.execute(
+          () -> {
+            try {
+              semaphore.acquire();
+            } catch (InterruptedException e) {
+              Thread.currentThread().interrupt();
+              throw new RuntimeException("Interrupted waiting for concurrency permit", e);
+            }
+            try {
+              command.run();
+            } finally {
+              semaphore.release();
+            }
+          });
+    }
+
+    @Override
+    public void shutdown() {
+      delegate.shutdown();
+    }
+
+    @Override
+    public List<Runnable> shutdownNow() {
+      return delegate.shutdownNow();
+    }
+
+    @Override
+    public boolean isShutdown() {
+      return delegate.isShutdown();
+    }
+
+    @Override
+    public boolean isTerminated() {
+      return delegate.isTerminated();
+    }
+
+    @Override
+    public boolean awaitTermination(long timeout, TimeUnit unit) throws InterruptedException {
+      return delegate.awaitTermination(timeout, unit);
+    }
+  }
 }