Add option to use virtual threads in place of thread pools (#865)

Author: aprudhomme

docs/server_configuration.rst

@@ -22,7 +22,7 @@ Example server configuration
     search:
       maxThreads: 4
     index:
-      maxThreads: 18
+      useVirtualThreads: true
   botoCfgPath: "/user/app/boto.cfg"
   bucketName: "nrtsearch-bucket"
   serviceName: "nrtsearch-service-test"
@@ -141,6 +141,11 @@ Example server configuration
      - Name prefix for threads created by searcher threadpool executor
      - LuceneSearchExecutor
 
+   * - search.useVirtualThreads
+     - bool
+     - Whether to use virtual threads instead of a traditional thread pool for search operations
+     - false
+
    * - index.maxThreads
      - int
      - Size of indexing threadpool executor
@@ -156,6 +161,11 @@ Example server configuration
      - Name prefix for threads created by indexing threadpool executor
      - LuceneIndexingExecutor
 
+   * - index.useVirtualThreads
+     - bool
+     - Whether to use virtual threads instead of a traditional thread pool for indexing operations
+     - false
+
    * - server.maxThreads
      - int
      - Size of NrtsearchServer threadpool executor
@@ -171,6 +181,11 @@ Example server configuration
      - Name prefix for threads created by NrtsearchServer threadpool executor
      - GrpcServerExecutor
 
+   * - server.useVirtualThreads
+     - bool
+     - Whether to use virtual threads instead of a traditional thread pool for server operations
+     - false
+
    * - replicationserver.maxThreads
      - int
      - Size of ReplicationServer threadpool executor
@@ -186,6 +201,11 @@ Example server configuration
      - Name prefix for threads created by ReplicationServer threadpool executor
      - GrpcReplicationServerExecutor
 
+   * - replicationserver.useVirtualThreads
+     - bool
+     - Whether to use virtual threads instead of a traditional thread pool for replication server operations
+     - false
+
    * - fetch.maxThreads
      - int
      - Size of fetch threadpool executor
@@ -201,6 +221,11 @@ Example server configuration
      - Name prefix for threads created by fetch threadpool executor
      - LuceneFetchExecutor
 
+   * - fetch.useVirtualThreads
+     - bool
+     - Whether to use virtual threads instead of a traditional thread pool for fetch operations
+     - false
+
    * - grpc.maxThreads
      - int
      - Size of gRPC threadpool executor
@@ -216,6 +241,11 @@ Example server configuration
      - Name prefix for threads created by gRPC threadpool executor
      - GrpcExecutor
 
+   * - grpc.useVirtualThreads
+     - bool
+     - Whether to use virtual threads instead of a traditional thread pool for gRPC operations
+     - false
+
    * - metrics.maxThreads
      - int
      - Size of metrics threadpool executor
@@ -231,6 +261,11 @@ Example server configuration
      - Name prefix for threads created by metrics threadpool executor
      - MetricsExecutor
 
+   * - metrics.useVirtualThreads
+     - bool
+     - Whether to use virtual threads instead of a traditional thread pool for metrics operations
+     - false
+
    * - vectormerge.maxThreads
      - int
      - Size of vector merge threadpool executor
@@ -246,6 +281,11 @@ Example server configuration
      - Name prefix for threads created by vector merge threadpool executor
      - VectorMergeExecutor
 
+   * - vectormerge.useVirtualThreads
+     - bool
+     - Whether to use virtual threads instead of a traditional thread pool for vector merge operations
+     - false
+
    * - commit.maxThreads
      - int
      - Size of commit threadpool executor
@@ -261,6 +301,11 @@ Example server configuration
      - Name prefix for threads created by commit threadpool executor
      - CommitExecutor
 
+   * - commit.useVirtualThreads
+     - bool
+     - Whether to use virtual threads instead of a traditional thread pool for commit operations
+     - false
+
 .. list-table:: `Alternative Max Threads Config <https://github.com/Yelp/nrtsearch/blob/master/src/main/java/com/yelp/nrtsearch/server/config/ThreadPoolConfiguration.java>`_ (``threadPoolConfiguration.*.maxThreads.*``)
    :widths: 25 10 50 25
    :header-rows: 1

src/main/java/com/yelp/nrtsearch/server/concurrent/ExecutorFactory.java

@@ -21,6 +21,7 @@
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ThreadPoolExecutor;
 import java.util.concurrent.TimeUnit;
@@ -97,21 +98,30 @@ public ExecutorService getExecutor(ExecutorType executorType) {
   private ExecutorService createExecutor(ExecutorType executorType) {
     ThreadPoolConfiguration.ThreadPoolSettings threadPoolSettings =
         threadPoolConfiguration.getThreadPoolSettings(executorType);
-    logger.info(
-        "Creating {} of size {}",
-        threadPoolSettings.threadNamePrefix(),
-        threadPoolSettings.maxThreads());
-    BlockingQueue<Runnable> queue =
-        new LinkedBlockingQueue<>(threadPoolSettings.maxBufferedItems());
-    ThreadPoolExecutor threadPoolExecutor =
-        new ThreadPoolExecutor(
-            threadPoolSettings.maxThreads(),
-            threadPoolSettings.maxThreads(),
-            0L,
-            TimeUnit.SECONDS,
-            queue,
-            new NamedThreadFactory(threadPoolSettings.threadNamePrefix()));
-    ThreadPoolCollector.addPool(executorType.name(), threadPoolExecutor);
-    return threadPoolExecutor;
+    if (threadPoolSettings.useVirtualThreads()) {
+      logger.info("Creating virtual thread executor for {}", threadPoolSettings.threadNamePrefix());
+      ExecutorService virtualThreadExecutor = Executors.newVirtualThreadPerTaskExecutor();
+      ExecutorServiceStatsWrapper executorServiceStatsWrapper =
+          new ExecutorServiceStatsWrapper(virtualThreadExecutor);
+      ThreadPoolCollector.addVirtualPool(executorType.name(), executorServiceStatsWrapper);
+      return executorServiceStatsWrapper;
+    } else {
+      logger.info(
+          "Creating {} of size {}",
+          threadPoolSettings.threadNamePrefix(),
+          threadPoolSettings.maxThreads());
+      BlockingQueue<Runnable> queue =
+          new LinkedBlockingQueue<>(threadPoolSettings.maxBufferedItems());
+      ThreadPoolExecutor threadPoolExecutor =
+          new ThreadPoolExecutor(
+              threadPoolSettings.maxThreads(),
+              threadPoolSettings.maxThreads(),
+              0L,
+              TimeUnit.SECONDS,
+              queue,
+              new NamedThreadFactory(threadPoolSettings.threadNamePrefix()));
+      ThreadPoolCollector.addPool(executorType.name(), threadPoolExecutor);
+      return threadPoolExecutor;
+    }
   }
 }

src/main/java/com/yelp/nrtsearch/server/concurrent/ExecutorServiceStatsWrapper.java

@@ -0,0 +1,126 @@
+/*
+ * Copyright 2025 Yelp Inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.yelp.nrtsearch.server.concurrent;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Future;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.TimeoutException;
+import java.util.concurrent.atomic.AtomicLong;
+
+/**
+ * {@link ExecutorService} implementation that wraps another {@link ExecutorService} and tracks
+ * stats. Currently, only tracks the total number of tasks submitted to the executor.
+ */
+public class ExecutorServiceStatsWrapper implements ExecutorService {
+  private final ExecutorService delegate;
+  private final AtomicLong totalTasks = new AtomicLong(0);
+
+  public ExecutorServiceStatsWrapper(ExecutorService delegate) {
+    this.delegate = delegate;
+  }
+
+  /**
+   * Returns the total number of tasks submitted to this executor.
+   *
+   * @return total number of tasks
+   */
+  public long getTotalTasks() {
+    return totalTasks.get();
+  }
+
+  @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);
+  }
+
+  @Override
+  public <T> Future<T> submit(Callable<T> task) {
+    totalTasks.incrementAndGet();
+    return delegate.submit(task);
+  }
+
+  @Override
+  public <T> Future<T> submit(Runnable task, T result) {
+    totalTasks.incrementAndGet();
+    return delegate.submit(task, result);
+  }
+
+  @Override
+  public Future<?> submit(Runnable task) {
+    totalTasks.incrementAndGet();
+    return delegate.submit(task);
+  }
+
+  @Override
+  public <T> List<Future<T>> invokeAll(Collection<? extends Callable<T>> tasks)
+      throws InterruptedException {
+    totalTasks.addAndGet(tasks.size());
+    return delegate.invokeAll(tasks);
+  }
+
+  @Override
+  public <T> List<Future<T>> invokeAll(
+      Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit)
+      throws InterruptedException {
+    totalTasks.addAndGet(tasks.size());
+    return delegate.invokeAll(tasks, timeout, unit);
+  }
+
+  @Override
+  public <T> T invokeAny(Collection<? extends Callable<T>> tasks)
+      throws InterruptedException, ExecutionException {
+    totalTasks.addAndGet(tasks.size());
+    return delegate.invokeAny(tasks);
+  }
+
+  @Override
+  public <T> T invokeAny(Collection<? extends Callable<T>> tasks, long timeout, TimeUnit unit)
+      throws InterruptedException, ExecutionException, TimeoutException {
+    totalTasks.addAndGet(tasks.size());
+    return delegate.invokeAny(tasks, timeout, unit);
+  }
+
+  @Override
+  public void execute(Runnable command) {
+    totalTasks.incrementAndGet();
+    delegate.execute(command);
+  }
+}

src/main/java/com/yelp/nrtsearch/server/config/ThreadPoolConfiguration.java

@@ -63,47 +63,62 @@ public class ThreadPoolConfiguration {
    * @param maxThreads max number of threads
    * @param maxBufferedItems max number of buffered items
    * @param threadNamePrefix prefix for thread names
+   * @param useVirtualThreads whether to use virtual threads, instead of a thread pool
    */
-  public record ThreadPoolSettings(int maxThreads, int maxBufferedItems, String threadNamePrefix) {}
+  public record ThreadPoolSettings(
+      int maxThreads, int maxBufferedItems, String threadNamePrefix, boolean useVirtualThreads) {}
 
   private static final Map<ExecutorFactory.ExecutorType, ThreadPoolSettings>
       defaultThreadPoolSettings =
           Map.of(
               ExecutorFactory.ExecutorType.SEARCH,
               new ThreadPoolSettings(
-                  DEFAULT_SEARCHING_THREADS, DEFAULT_SEARCH_BUFFERED_ITEMS, "LuceneSearchExecutor"),
+                  DEFAULT_SEARCHING_THREADS,
+                  DEFAULT_SEARCH_BUFFERED_ITEMS,
+                  "LuceneSearchExecutor",
+                  false),
               ExecutorFactory.ExecutorType.INDEX,
               new ThreadPoolSettings(
                   DEFAULT_INDEXING_THREADS,
                   DEFAULT_INDEXING_BUFFERED_ITEMS,
-                  "LuceneIndexingExecutor"),
+                  "LuceneIndexingExecutor",
+                  false),
               ExecutorFactory.ExecutorType.SERVER,
               new ThreadPoolSettings(
                   DEFAULT_GRPC_SERVER_THREADS,
                   DEFAULT_GRPC_SERVER_BUFFERED_ITEMS,
-                  "GrpcServerExecutor"),
+                  "GrpcServerExecutor",
+                  false),
               ExecutorFactory.ExecutorType.REPLICATIONSERVER,
               new ThreadPoolSettings(
                   DEFAULT_GRPC_REPLICATIONSERVER_THREADS,
                   DEFAULT_GRPC_REPLICATIONSERVER_BUFFERED_ITEMS,
-                  "GrpcReplicationServerExecutor"),
+                  "GrpcReplicationServerExecutor",
+                  false),
               ExecutorFactory.ExecutorType.FETCH,
               new ThreadPoolSettings(
-                  DEFAULT_FETCH_THREADS, DEFAULT_FETCH_BUFFERED_ITEMS, "LuceneFetchExecutor"),
+                  DEFAULT_FETCH_THREADS,
+                  DEFAULT_FETCH_BUFFERED_ITEMS,
+                  "LuceneFetchExecutor",
+                  false),
               ExecutorFactory.ExecutorType.GRPC,
               new ThreadPoolSettings(
-                  DEFAULT_GRPC_THREADS, DEFAULT_GRPC_BUFFERED_ITEMS, "GrpcExecutor"),
+                  DEFAULT_GRPC_THREADS, DEFAULT_GRPC_BUFFERED_ITEMS, "GrpcExecutor", false),
               ExecutorFactory.ExecutorType.METRICS,
               new ThreadPoolSettings(
-                  DEFAULT_METRICS_THREADS, DEFAULT_METRICS_BUFFERED_ITEMS, "MetricsExecutor"),
+                  DEFAULT_METRICS_THREADS,
+                  DEFAULT_METRICS_BUFFERED_ITEMS,
+                  "MetricsExecutor",
+                  false),
               ExecutorFactory.ExecutorType.VECTORMERGE,
               new ThreadPoolSettings(
                   DEFAULT_VECTOR_MERGE_THREADS,
                   DEFAULT_VECTOR_MERGE_BUFFERED_ITEMS,
-                  "VectorMergeExecutor"),
+                  "VectorMergeExecutor",
+                  false),
               ExecutorFactory.ExecutorType.COMMIT,
               new ThreadPoolSettings(
-                  DEFAULT_COMMIT_THREADS, DEFAULT_COMMIT_BUFFERED_ITEMS, "CommitExecutor"));
+                  DEFAULT_COMMIT_THREADS, DEFAULT_COMMIT_BUFFERED_ITEMS, "CommitExecutor", false));
 
   private final Map<ExecutorFactory.ExecutorType, ThreadPoolSettings> threadPoolSettings;
 
@@ -121,8 +136,13 @@ public ThreadPoolConfiguration(YamlConfigReader configReader) {
       String threadNamePrefix =
           configReader.getString(
               poolConfigPrefix + "threadNamePrefix", defaultSettings.threadNamePrefix());
+      boolean useVirtualThreads =
+          configReader.getBoolean(
+              poolConfigPrefix + "useVirtualThreads", defaultSettings.useVirtualThreads());
       threadPoolSettings.put(
-          executorType, new ThreadPoolSettings(maxThreads, maxBufferedItems, threadNamePrefix));
+          executorType,
+          new ThreadPoolSettings(
+              maxThreads, maxBufferedItems, threadNamePrefix, useVirtualThreads));
     }
   }