Fix Session state bug and improve Query Efficiency in REPL (#245) (#250)

* Fix Session state bug and improve Query Efficiency in REPL

This PR introduces a bug fix and enhancements to FlintREPL's session management and optimizes query execution methods. It addresses a specific issue where marking a session as 'dead' inadvertently triggered the creation of a new, unnecessary session. This behavior resulted in the new session entering a spin-wait state, leading to duplicate jobs.

The improvements include:
- **Introduction of `earlyExitFlag`**: A new flag, `earlyExitFlag`, has been introduced and is set to `true` under two conditions: when a job is excluded or when it is not associated with the current session's job run ID. This flag is evaluated in the shutdown hook to determine whether the session state should be marked as 'dead'. This change effectively prevents the unintended creation of duplicate sessions by the SQL plugin, ensuring resources are utilized more efficiently.
- **Query Method Optimization**: The method for executing queries has been shifted from scrolling to search, eliminating the need for creating unnecessary scroll contexts. This adjustment enhances the performance and efficiency of query operations.
- **Reversion of Previous Commit**: The PR reverts a previous change (be82024) following the resolution of the related issue in the SQL plugin (opensearch-project/sql#2436), further streamlining the operation and maintenance of the system.

**Testing**:
1. Integration tests were added to cover both REPL and streaming job functionalities, ensuring the robustness of the fixes.
2. Manual testing was conducted to validate the bug fix.



* added Java doc



* fix IT by restore env variable change



---------

Signed-off-by: Kaituo Li <[email protected]>

build.sbt

@@ -42,8 +42,9 @@ lazy val commonSettings = Seq(
   testScalastyle := (Test / scalastyle).toTask("").value,
   Test / test := ((Test / test) dependsOn testScalastyle).value)
 
+// running `scalafmtAll` includes all subprojects under root
 lazy val root = (project in file("."))
-  .aggregate(flintCore, flintSparkIntegration, pplSparkIntegration, sparkSqlApplication)
+  .aggregate(flintCore, flintSparkIntegration, pplSparkIntegration, sparkSqlApplication, integtest)
   .disablePlugins(AssemblyPlugin)
   .settings(name := "flint", publish / skip := true)
 
@@ -158,7 +159,7 @@ lazy val flintSparkIntegration = (project in file("flint-spark-integration"))
 
 // Test assembly package with integration test.
 lazy val integtest = (project in file("integ-test"))
-  .dependsOn(flintSparkIntegration % "test->test", pplSparkIntegration % "test->test" )
+  .dependsOn(flintSparkIntegration % "test->test", pplSparkIntegration % "test->test", sparkSqlApplication % "test->test")
   .settings(
     commonSettings,
     name := "integ-test",
@@ -174,7 +175,9 @@ lazy val integtest = (project in file("integ-test"))
       "org.opensearch.client" % "opensearch-java" % "2.6.0" % "test"
         exclude ("com.fasterxml.jackson.core", "jackson-databind")),
     libraryDependencies ++= deps(sparkVersion),
-    Test / fullClasspath ++= Seq((flintSparkIntegration / assembly).value, (pplSparkIntegration / assembly).value))
+    Test / fullClasspath ++= Seq((flintSparkIntegration / assembly).value, (pplSparkIntegration / assembly).value,
+      (sparkSqlApplication / assembly).value
+    ))
 
 lazy val standaloneCosmetic = project
   .settings(

flint-core/src/main/scala/org/opensearch/flint/core/FlintOptions.java

@@ -77,6 +77,8 @@ public class FlintOptions implements Serializable {
 
   public static final int DEFAULT_SOCKET_TIMEOUT_MILLIS = 60000;
 
+  public static final int DEFAULT_INACTIVITY_LIMIT_MILLIS = 10 * 60 * 1000;
+
   public FlintOptions(Map<String, String> options) {
     this.options = options;
     this.retryOptions = new FlintRetryOptions(options);

flint-core/src/main/scala/org/opensearch/flint/core/storage/OpenSearchQueryReader.java

@@ -0,0 +1,48 @@
+/*
+ * Copyright OpenSearch Contributors
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package org.opensearch.flint.core.storage;
+
+import org.opensearch.OpenSearchStatusException;
+import org.opensearch.action.search.ClearScrollRequest;
+import org.opensearch.action.search.SearchRequest;
+import org.opensearch.action.search.SearchResponse;
+import org.opensearch.action.search.SearchScrollRequest;
+import org.opensearch.client.RequestOptions;
+import org.opensearch.client.RestHighLevelClient;
+import org.opensearch.common.Strings;
+import org.opensearch.common.unit.TimeValue;
+import org.opensearch.flint.core.FlintOptions;
+import org.opensearch.client.RestHighLevelClient;
+import org.opensearch.search.builder.SearchSourceBuilder;
+
+import java.io.IOException;
+import java.util.Optional;
+import java.util.logging.Level;
+import java.util.logging.Logger;
+
+/**
+ * {@link OpenSearchReader} using search. https://opensearch.org/docs/latest/api-reference/search/
+ */
+public class OpenSearchQueryReader extends OpenSearchReader {
+
+  private static final Logger LOG = Logger.getLogger(OpenSearchQueryReader.class.getName());
+
+  public OpenSearchQueryReader(RestHighLevelClient client, String indexName, SearchSourceBuilder searchSourceBuilder) {
+    super(client, new SearchRequest().indices(indexName).source(searchSourceBuilder));
+  }
+
+  /**
+   * search.
+   */
+  Optional<SearchResponse> search(SearchRequest request) throws IOException {
+   return Optional.of(client.search(request, RequestOptions.DEFAULT));
+  }
+
+  /**
+   * nothing to clean
+   */
+  void clean() throws IOException {}
+}

flint-core/src/main/scala/org/opensearch/flint/core/storage/OpenSearchReader.java

@@ -5,6 +5,7 @@
 
 package org.opensearch.flint.core.storage;
 
+import org.opensearch.OpenSearchStatusException;
 import org.opensearch.action.search.SearchRequest;
 import org.opensearch.action.search.SearchResponse;
 import org.opensearch.client.RestHighLevelClient;
@@ -48,6 +49,13 @@ public OpenSearchReader(RestHighLevelClient client, SearchRequest searchRequest)
         iterator = searchHits.iterator();
       }
       return iterator.hasNext();
+    } catch (OpenSearchStatusException e) {
+      //  e.g., org.opensearch.OpenSearchStatusException: OpenSearch exception [type=index_not_found_exception, reason=no such index [query_results2]]
+      if (e.getMessage() != null && (e.getMessage().contains("index_not_found_exception"))) {
+        return false;
+      } else {
+        throw e;
+      }
     } catch (IOException e) {
       // todo. log error.
       throw new RuntimeException(e);

flint-spark-integration/src/main/scala/org/apache/spark/sql/flint/config/FlintSparkConf.scala

@@ -146,7 +146,30 @@ object FlintSparkConf {
       .datasourceOption()
       .doc("socket duration in milliseconds")
       .createWithDefault(String.valueOf(FlintOptions.DEFAULT_SOCKET_TIMEOUT_MILLIS))
-
+  val DATA_SOURCE_NAME =
+    FlintConfig(s"spark.flint.datasource.name")
+      .doc("data source name")
+      .createOptional()
+  val JOB_TYPE =
+    FlintConfig(s"spark.flint.job.type")
+      .doc("Flint job type. Including interactive and streaming")
+      .createWithDefault("interactive")
+  val SESSION_ID =
+    FlintConfig(s"spark.flint.job.sessionId")
+      .doc("Flint session id")
+      .createOptional()
+  val REQUEST_INDEX =
+    FlintConfig(s"spark.flint.job.requestIndex")
+      .doc("Request index")
+      .createOptional()
+  val EXCLUDE_JOB_IDS =
+    FlintConfig(s"spark.flint.deployment.excludeJobs")
+      .doc("Exclude job ids")
+      .createOptional()
+  val REPL_INACTIVITY_TIMEOUT_MILLIS =
+    FlintConfig(s"spark.flint.job.inactivityLimitMillis")
+      .doc("inactivity timeout")
+      .createWithDefault(String.valueOf(FlintOptions.DEFAULT_INACTIVITY_LIMIT_MILLIS))
 }
 
 /**
@@ -196,11 +219,18 @@ case class FlintSparkConf(properties: JMap[String, String]) extends Serializable
       CUSTOM_AWS_CREDENTIALS_PROVIDER,
       USERNAME,
       PASSWORD,
-      SOCKET_TIMEOUT_MILLIS)
+      SOCKET_TIMEOUT_MILLIS,
+      JOB_TYPE,
+      REPL_INACTIVITY_TIMEOUT_MILLIS)
       .map(conf => (conf.optionKey, conf.readFrom(reader)))
       .toMap
 
-    val optionsWithoutDefault = Seq(RETRYABLE_EXCEPTION_CLASS_NAMES)
+    val optionsWithoutDefault = Seq(
+      RETRYABLE_EXCEPTION_CLASS_NAMES,
+      DATA_SOURCE_NAME,
+      SESSION_ID,
+      REQUEST_INDEX,
+      EXCLUDE_JOB_IDS)
       .map(conf => (conf.optionKey, conf.readFrom(reader)))
       .flatMap {
         case (_, None) => None

flint-spark-integration/src/main/scala/org/opensearch/flint/app/FlintInstance.scala

@@ -25,7 +25,14 @@ class FlintInstance(
     val lastUpdateTime: Long,
     val jobStartTime: Long = 0,
     val excludedJobIds: Seq[String] = Seq.empty[String],
-    val error: Option[String] = None) {}
+    val error: Option[String] = None) {
+  override def toString: String = {
+    val excludedJobIdsStr = excludedJobIds.mkString("[", ", ", "]")
+    val errorStr = error.getOrElse("None")
+    s"FlintInstance(applicationId=$applicationId, jobId=$jobId, sessionId=$sessionId, state=$state, " +
+      s"lastUpdateTime=$lastUpdateTime, jobStartTime=$jobStartTime, excludedJobIds=$excludedJobIdsStr, error=$errorStr)"
+  }
+}
 
 object FlintInstance {