Use SQL thread pool

Signed-off-by: Louis Chu <[email protected]>
noCharger · Sep 4, 2024 · f9fe064 · f9fe064
1 parent 1e489b5
commit f9fe064
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 3 deletions.
diff --git a/...ry/src/main/java/org/opensearch/sql/spark/scheduler/job/ScheduledAsyncQueryJobRunner.java b/...ry/src/main/java/org/opensearch/sql/spark/scheduler/job/ScheduledAsyncQueryJobRunner.java
@@ -13,6 +13,7 @@
 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;
@@ -34,6 +35,9 @@
  * plugin.
  */
 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;
   private static final Logger LOGGER = LogManager.getLogger(ScheduledAsyncQueryJobRunner.class);
 
   private static ScheduledAsyncQueryJobRunner INSTANCE = new ScheduledAsyncQueryJobRunner();
@@ -96,7 +100,7 @@ public void runJob(ScheduledJobParameter jobParameter, JobExecutionContext conte
             LOGGER.error(throwable);
           }
         };
-    threadPool.generic().submit(runnable);
+    threadPool.executor(ASYNC_QUERY_THREAD_POOL_NAME).submit(runnable);
   }
 
   void doRefresh(ScheduledAsyncQueryJobRequest request) {

diff --git a/...rc/test/java/org/opensearch/sql/spark/scheduler/job/ScheduledAsyncQueryJobRunnerTest.java b/...rc/test/java/org/opensearch/sql/spark/scheduler/job/ScheduledAsyncQueryJobRunnerTest.java
@@ -31,6 +31,7 @@
 import org.opensearch.cluster.service.ClusterService;
 import org.opensearch.jobscheduler.spi.JobExecutionContext;
 import org.opensearch.jobscheduler.spi.ScheduledJobParameter;
+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;
@@ -86,7 +87,8 @@ public void testRunJobWithCorrectParameter() {
     spyJobRunner.runJob(request, context);
 
     ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
-    verify(threadPool.generic()).submit(captor.capture());
+    verify(threadPool.executor(AsyncRestExecutor.SQL_WORKER_THREAD_POOL_NAME))
+        .submit(captor.capture());
 
     Runnable runnable = captor.getValue();
     runnable.run();
@@ -143,7 +145,8 @@ public void testDoRefreshThrowsException() {
     spyJobRunner.runJob(request, context);
 
     ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
-    verify(threadPool.generic()).submit(captor.capture());
+    verify(threadPool.executor(AsyncRestExecutor.SQL_WORKER_THREAD_POOL_NAME))
+        .submit(captor.capture());
 
     Runnable runnable = captor.getValue();
     runnable.run();