Adding connection timeout override for 10 seconds (#1273)

* Adding connection timeout override for 10 seconds

Signed-off-by: Sachin Sriramagiri <srirasac@amazon.com>

* chaining request config builder options

Signed-off-by: Sachin Sriramagiri <srirasac@amazon.com>

---------

Signed-off-by: Sachin Sriramagiri <srirasac@amazon.com>
(cherry picked from commit 9b5c6ed)
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>

Author: github-actions[bot]

docs/index.md

@@ -556,6 +556,7 @@ In the index mapping, the `_meta` and `properties`field stores meta and schema i
 - `spark.flint.index.checkpointLocation.rootDir`: default is None. Flint will create a default checkpoint location in format of '<rootDir>/<indexName>/<UUID>' to isolate checkpoint data.
 - `spark.flint.index.checkpoint.mandatory`: default is true.
 - `spark.datasource.flint.socket_timeout_millis`: default value is 60000.
+- `spark.datasource.flint.connection_timeout_millis`: default value is 10000.
 - `spark.datasource.flint.request.completionDelayMillis`: Time to wait in milliseconds after request is complete. Applied after index creation. Default value is 2000 if using aoss service, otherwise 0.
 - `spark.flint.monitor.initialDelaySeconds`: Initial delay in seconds before starting the monitoring task. Default value is 15.
 - `spark.flint.monitor.intervalSeconds`: Interval in seconds for scheduling the monitoring task. Default value is 60.

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

@@ -87,6 +87,10 @@ public class FlintOptions implements Serializable {
 
   public static final int DEFAULT_SOCKET_TIMEOUT_MILLIS = 60000;
 
+  public static final String CONNECTION_TIMEOUT_MILLIS = "connection_timeout_millis";
+
+  public static final int DEFAULT_CONNECTION_TIMEOUT_MILLIS = 10000;
+
   public static final int DEFAULT_INACTIVITY_LIMIT_MILLIS = 3 * 60 * 1000;
 
   public static final String REQUEST_COMPLETION_DELAY_MILLIS = "request.completionDelayMillis";
@@ -207,6 +211,10 @@ public int getSocketTimeoutMillis() {
     return Integer.parseInt(options.getOrDefault(SOCKET_TIMEOUT_MILLIS, String.valueOf(DEFAULT_SOCKET_TIMEOUT_MILLIS)));
   }
 
+  public int getConnectionTimeoutMillis() {
+    return Integer.parseInt(options.getOrDefault(CONNECTION_TIMEOUT_MILLIS, String.valueOf(DEFAULT_CONNECTION_TIMEOUT_MILLIS)));
+  }
+
   public int getRequestCompletionDelayMillis() {
     int defaultValue = SERVICE_NAME_AOSS.equals(getServiceName())
         ? DEFAULT_AOSS_REQUEST_COMPLETION_DELAY_MILLIS

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

@@ -11,6 +11,7 @@
 
 /**
  * allows override default socket timeout in RestClientBuilder.DEFAULT_SOCKET_TIMEOUT_MILLIS
+ * allows override of default connection timeout in RestClientBuilder.DEFAULT_CONNECT_TIMEOUT_MILLIS
  */
 public class RequestConfigurator implements RestClientBuilder.RequestConfigCallback {
 
@@ -22,8 +23,9 @@ public RequestConfigurator(FlintOptions options) {
 
     @Override
     public RequestConfig.Builder customizeRequestConfig(RequestConfig.Builder requestConfigBuilder) {
-        // Set the socket timeout in milliseconds
-        return requestConfigBuilder.setSocketTimeout(options.getSocketTimeoutMillis());
+        // Set the connection and socket timeouts in milliseconds
+        return requestConfigBuilder.setConnectTimeout(options.getConnectionTimeoutMillis())
+                .setSocketTimeout(options.getSocketTimeoutMillis());
     }
 }
 

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

@@ -251,6 +251,13 @@ object FlintSparkConf {
       .datasourceOption()
       .doc("socket duration in milliseconds")
       .createWithDefault(String.valueOf(FlintOptions.DEFAULT_SOCKET_TIMEOUT_MILLIS))
+
+  val CONNECTION_TIMEOUT_MILLIS =
+    FlintConfig(s"spark.datasource.flint.${FlintOptions.CONNECTION_TIMEOUT_MILLIS}")
+      .datasourceOption()
+      .doc("connection duration in milliseconds")
+      .createWithDefault(String.valueOf(FlintOptions.DEFAULT_CONNECTION_TIMEOUT_MILLIS))
+
   val REQUEST_COMPLETION_DELAY_MILLIS =
     FlintConfig(s"spark.datasource.flint.${FlintOptions.REQUEST_COMPLETION_DELAY_MILLIS}")
       .datasourceOption()
@@ -408,6 +415,7 @@ case class FlintSparkConf(properties: JMap[String, String]) extends Serializable
       USERNAME,
       PASSWORD,
       SOCKET_TIMEOUT_MILLIS,
+      CONNECTION_TIMEOUT_MILLIS,
       JOB_TYPE,
       REPL_INACTIVITY_TIMEOUT_MILLIS,
       BATCH_BYTES)