Implement one-batch lookahead for index enumerators (opensearch-project#4345)

Author: Swiddis

async-query/src/main/java/org/opensearch/sql/spark/scheduler/job/ScheduledAsyncQueryJobRunner.java

@@ -12,7 +12,6 @@
 import org.opensearch.jobscheduler.spi.ScheduledJobParameter;
 import org.opensearch.jobscheduler.spi.ScheduledJobRunner;
 import org.opensearch.plugins.Plugin;
-import org.opensearch.sql.legacy.executor.AsyncRestExecutor;
 import org.opensearch.sql.spark.asyncquery.AsyncQueryExecutorService;
 import org.opensearch.sql.spark.asyncquery.model.NullAsyncQueryRequestContext;
 import org.opensearch.sql.spark.rest.model.CreateAsyncQueryRequest;
@@ -21,6 +20,8 @@
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.client.Client;
 
+import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME;
+
 /**
  * The job runner class for scheduling async query.
  *
@@ -37,7 +38,7 @@
 public class ScheduledAsyncQueryJobRunner implements ScheduledJobRunner {
   // Share SQL plugin thread pool
   private static final String ASYNC_QUERY_THREAD_POOL_NAME =
-      AsyncRestExecutor.SQL_WORKER_THREAD_POOL_NAME;
+      SQL_WORKER_THREAD_POOL_NAME;
   private static final Logger LOGGER = LogManager.getLogger(ScheduledAsyncQueryJobRunner.class);
 
   private static final ScheduledAsyncQueryJobRunner INSTANCE = new ScheduledAsyncQueryJobRunner();

async-query/src/test/java/org/opensearch/sql/spark/scheduler/job/ScheduledAsyncQueryJobRunnerTest.java

@@ -15,6 +15,7 @@
 import static org.mockito.Mockito.spy;
 import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
+import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME;
 
 import java.time.Instant;
 import org.apache.logging.log4j.LogManager;
@@ -87,7 +88,7 @@ public void testRunJobWithCorrectParameter() {
     spyJobRunner.runJob(request, context);
 
     ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
-    verify(threadPool.executor(AsyncRestExecutor.SQL_WORKER_THREAD_POOL_NAME))
+    verify(threadPool.executor(SQL_WORKER_THREAD_POOL_NAME))
         .submit(captor.capture());
 
     Runnable runnable = captor.getValue();
@@ -145,7 +146,7 @@ public void testDoRefreshThrowsException() {
     spyJobRunner.runJob(request, context);
 
     ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
-    verify(threadPool.executor(AsyncRestExecutor.SQL_WORKER_THREAD_POOL_NAME))
+    verify(threadPool.executor(SQL_WORKER_THREAD_POOL_NAME))
         .submit(captor.capture());
 
     Runnable runnable = captor.getValue();

datasources/src/main/java/org/opensearch/sql/datasources/utils/Scheduler.java

@@ -11,13 +11,11 @@
 import org.opensearch.common.unit.TimeValue;
 import org.opensearch.threadpool.ThreadPool;
 import org.opensearch.transport.client.node.NodeClient;
+import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME;
 
 /** The scheduler which schedule the task run in sql-worker thread pool. */
 @UtilityClass
 public class Scheduler {
-
-  public static final String SQL_WORKER_THREAD_POOL_NAME = "sql-worker";
-
   public static void schedule(NodeClient client, Runnable task) {
     ThreadPool threadPool = client.threadPool();
     threadPool.schedule(withCurrentContext(task), new TimeValue(0), SQL_WORKER_THREAD_POOL_NAME);

docs/user/admin/settings.rst

@@ -253,6 +253,27 @@ Result set::
       "transient": {}
     }
 
+Thread Pool Settings
+====================
+
+The SQL plugin is integrated with the `OpenSearch Thread Pool Settings <https://docs.opensearch.org/latest/install-and-configure/configuring-opensearch/thread-pool-settings/>`_.
+There are two thread pools which can be configured on cluster setup via `settings.yml`::
+
+    thread_pool:
+      sql-worker:
+        size: 30
+        queue_size: 100
+      sql_background_io:
+        size: 30
+        queue_size: 1000
+
+The ``sql-worker`` pool corresponds to compute resources related to running queries, such as compute-heavy evaluations on result sets.
+This directly maps to the number of queries that can be run concurrently.
+This is the primary pool you interact with externally.
+``sql_background_io`` is a low-footprint pool for IO requests the plugin makes,
+and can be used to limit indirect load that SQL places on your cluster for Calcite-enabled operations.
+A ``sql-worker`` thread may spawn multiple background threads.
+
 plugins.query.executionengine.spark.session.limit
 ==================================================
 

legacy/src/main/java/org/opensearch/sql/legacy/executor/AsyncRestExecutor.java

@@ -5,6 +5,8 @@
 
 package org.opensearch.sql.legacy.executor;
 
+import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME;
+
 import java.io.IOException;
 import java.time.Duration;
 import java.util.Map;
@@ -30,10 +32,6 @@
 
 /** A RestExecutor wrapper to execute request asynchronously to avoid blocking transport thread. */
 public class AsyncRestExecutor implements RestExecutor {
-
-  /** Custom thread pool name managed by OpenSearch */
-  public static final String SQL_WORKER_THREAD_POOL_NAME = "sql-worker";
-
   private static final Logger LOG = LogManager.getLogger(AsyncRestExecutor.class);
 
   /**

legacy/src/main/java/org/opensearch/sql/legacy/executor/cursor/CursorAsyncRestExecutor.java

@@ -5,6 +5,8 @@
 
 package org.opensearch.sql.legacy.executor.cursor;
 
+import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME;
+
 import java.io.IOException;
 import java.time.Duration;
 import java.util.Map;
@@ -24,9 +26,6 @@
 import org.opensearch.transport.client.Client;
 
 public class CursorAsyncRestExecutor {
-  /** Custom thread pool name managed by OpenSearch */
-  public static final String SQL_WORKER_THREAD_POOL_NAME = "sql-worker";
-
   private static final Logger LOG = LogManager.getLogger(CursorAsyncRestExecutor.class);
 
   /** Delegated rest executor to async */

legacy/src/main/java/org/opensearch/sql/legacy/plugin/RestSqlStatsAction.java

@@ -6,6 +6,7 @@
 package org.opensearch.sql.legacy.plugin;
 
 import static org.opensearch.core.rest.RestStatus.INTERNAL_SERVER_ERROR;
+import static org.opensearch.sql.opensearch.executor.OpenSearchQueryManager.SQL_WORKER_THREAD_POOL_NAME;
 
 import com.google.common.collect.ImmutableList;
 import java.util.Arrays;
@@ -90,7 +91,7 @@ protected Set<String> responseParams() {
 
   private void schedule(NodeClient client, Runnable task) {
     ThreadPool threadPool = client.threadPool();
-    threadPool.schedule(withCurrentContext(task), new TimeValue(0), "sql-worker");
+    threadPool.schedule(withCurrentContext(task), new TimeValue(0), SQL_WORKER_THREAD_POOL_NAME);
   }
 
   private Runnable withCurrentContext(final Runnable task) {

opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchClient.java

@@ -7,6 +7,7 @@
 
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import org.opensearch.action.search.CreatePitRequest;
 import org.opensearch.action.search.DeletePitRequest;
 import org.opensearch.sql.opensearch.mapping.IndexMapping;
@@ -97,7 +98,7 @@ public interface OpenSearchClient {
    */
   void schedule(Runnable task);
 
-  NodeClient getNodeClient();
+  Optional<NodeClient> getNodeClient();
 
   /**
    * Create PIT for given indices

opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchNodeClient.java

@@ -11,6 +11,7 @@
 import java.util.Collection;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.concurrent.ExecutionException;
 import java.util.function.Function;
 import java.util.function.Predicate;
@@ -223,8 +224,8 @@ public void schedule(Runnable task) {
   }
 
   @Override
-  public NodeClient getNodeClient() {
-    return client;
+  public Optional<NodeClient> getNodeClient() {
+    return Optional.of(client);
   }
 
   @Override

opensearch/src/main/java/org/opensearch/sql/opensearch/client/OpenSearchRestClient.java

@@ -13,6 +13,7 @@
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.Optional;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import lombok.RequiredArgsConstructor;
@@ -236,8 +237,8 @@ public void schedule(Runnable task) {
   }
 
   @Override
-  public NodeClient getNodeClient() {
-    throw new UnsupportedOperationException("Unsupported method.");
+  public Optional<NodeClient> getNodeClient() {
+    return Optional.empty();
   }
 
   @Override