MLE-12437 Allow for OkHttp timeouts to be configured

Author: rjrudin

build.gradle

@@ -33,6 +33,9 @@ dependencies {
     exclude module: 'jackson-dataformat-csv'
   }
 
+  // Need this so that an OkHttpClientConfigurator can be created.
+  implementation 'com.squareup.okhttp3:okhttp:4.12.0'
+
   // Makes it possible to use lambdas in Java 8 to implement Spark's Function1 and Function2 interfaces
   // See https://github.com/scala/scala-java8-compat for more information
   implementation("org.scala-lang.modules:scala-java8-compat_2.12:1.0.2") {

src/main/java/com/marklogic/spark/ContextSupport.java

@@ -17,6 +17,7 @@
 
 import com.marklogic.client.DatabaseClient;
 import com.marklogic.client.DatabaseClientFactory;
+import com.marklogic.client.extra.okhttpclient.OkHttpClientConfigurator;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -25,14 +26,22 @@
 import java.net.URLDecoder;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.concurrent.TimeUnit;
 
 public class ContextSupport implements Serializable {
 
     protected static final Logger logger = LoggerFactory.getLogger(ContextSupport.class);
     private final Map<String, String> properties;
+    private final boolean configuratorWasAdded;
+
+    // Java Client 6.5.0 has a bug in it (to be fixed in 6.5.1 or 6.6.0) where multiple threads that use a configurator
+    // can run into a ConcurrentModificationException. So need to synchronize adding a configurator and creating a
+    // client. Those two actions are rarely done, so the cost of synchronization will be negligible.
+    private static final Object CLIENT_LOCK = new Object();
 
     protected ContextSupport(Map<String, String> properties) {
         this.properties = properties;
+        this.configuratorWasAdded = addOkHttpConfiguratorIfNecessary();
     }
 
     public DatabaseClient connectToMarkLogic() {
@@ -50,10 +59,12 @@ public DatabaseClient connectToMarkLogic(String host) {
             connectionProps.put(Options.CLIENT_HOST, host);
         }
         DatabaseClient client;
-        try {
-            client = DatabaseClientFactory.newClient(propertyName -> connectionProps.get("spark." + propertyName));
-        } catch (Exception e) {
-            throw new ConnectorException(String.format("Unable to connect to MarkLogic; cause: %s", e.getMessage()), e);
+        if (configuratorWasAdded) {
+            synchronized (CLIENT_LOCK) {
+                client = connect(connectionProps);
+            }
+        } else {
+            client = connect(connectionProps);
         }
         DatabaseClient.ConnectionResult result = client.checkConnection();
         if (!result.isConnected()) {
@@ -62,6 +73,14 @@ public DatabaseClient connectToMarkLogic(String host) {
         return client;
     }
 
+    private DatabaseClient connect(Map<String, String> connectionProps) {
+        try {
+            return DatabaseClientFactory.newClient(propertyName -> connectionProps.get("spark." + propertyName));
+        } catch (Exception e) {
+            throw new ConnectorException(String.format("Unable to connect to MarkLogic; cause: %s", e.getMessage()), e);
+        }
+    }
+
     protected final Map<String, String> buildConnectionProperties() {
         Map<String, String> connectionProps = new HashMap<>();
         connectionProps.put("spark.marklogic.client.authType", "digest");
@@ -154,4 +173,37 @@ public final boolean hasOption(String... options) {
     public Map<String, String> getProperties() {
         return properties;
     }
+
+    /**
+     * @return true if a configurator was added
+     */
+    private boolean addOkHttpConfiguratorIfNecessary() {
+        final String prefix = "spark.marklogic.client.";
+        final long defaultValue = -1;
+        final long connectionTimeout = getNumericOption(prefix + "connectionTimeout", defaultValue, defaultValue);
+        final long callTimeout = getNumericOption(prefix + "callTimeout", defaultValue, defaultValue);
+        final long readTimeout = getNumericOption(prefix + "connectionTimeout", defaultValue, defaultValue);
+        final long writeTimeout = getNumericOption(prefix + "writeTimeout", defaultValue, defaultValue);
+
+        if (connectionTimeout > -1 || callTimeout > -1 || readTimeout > -1 || writeTimeout > -1) {
+            synchronized (CLIENT_LOCK) {
+                DatabaseClientFactory.addConfigurator((OkHttpClientConfigurator) builder -> {
+                    if (connectionTimeout > -1) {
+                        builder.connectTimeout(connectionTimeout, TimeUnit.SECONDS);
+                    }
+                    if (callTimeout > -1) {
+                        builder.callTimeout(callTimeout, TimeUnit.SECONDS);
+                    }
+                    if (readTimeout > -1) {
+                        builder.readTimeout(readTimeout, TimeUnit.SECONDS);
+                    }
+                    if (writeTimeout > -1) {
+                        builder.writeTimeout(writeTimeout, TimeUnit.SECONDS);
+                    }
+                });
+            }
+            return true;
+        }
+        return false;
+    }
 }

src/test/java/com/marklogic/spark/ContextSupportTest.java

@@ -5,8 +5,13 @@
 import org.junit.jupiter.api.Test;
 
 import java.net.UnknownHostException;
+import java.util.ArrayList;
 import java.util.HashMap;
+import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
 
 import static org.junit.jupiter.api.Assertions.*;
 
@@ -42,4 +47,32 @@ void isDirectConnection() {
         options.put(Options.CLIENT_CONNECTION_TYPE, "gateway");
         assertFalse(new ContextSupport(options).isDirectConnection());
     }
+
+    @Test
+    void createManyConnectionsWithOkHttpConfigurator() throws Exception {
+        options.put(Options.CLIENT_URI, makeClientUri());
+        options.put("spark.marklogic.client.callTimeout", "10");
+        options.put("spark.marklogic.client.connectionTimeout", "10");
+        options.put("spark.marklogic.client.readTimeout", "10");
+        options.put("spark.marklogic.client.writeTimeout", "10");
+
+        final int clientsToCreate = 100;
+
+        ExecutorService service = Executors.newFixedThreadPool(16);
+        List<Future> futures = new ArrayList<>();
+        for (int i = 0; i < clientsToCreate; i++) {
+            futures.add(service.submit(() -> new ContextSupport(options).connectToMarkLogic().release()));
+        }
+
+        int clientsCreated = 0;
+        for (Future f : futures) {
+            f.get();
+            clientsCreated++;
+        }
+        service.shutdown();
+
+        assertEquals(clientsToCreate, clientsCreated, "The expectation is that many clients can be created by " +
+            "multiple threads at the same time due to the use of synchronization in ContextSupport. This won't be " +
+            "necessary once we upgrade to Java Clien 6.5.1 or higher.");
+    }
 }

src/test/java/com/marklogic/spark/reader/customcode/ReadWithCustomCodeTest.java

@@ -1,8 +1,10 @@
 package com.marklogic.spark.reader.customcode;
 
 import com.marklogic.client.FailedRequestException;
+import com.marklogic.client.MarkLogicIOException;
 import com.marklogic.spark.AbstractIntegrationTest;
 import com.marklogic.spark.Options;
+import org.apache.spark.SparkException;
 import org.apache.spark.sql.DataFrameReader;
 import org.apache.spark.sql.Dataset;
 import org.apache.spark.sql.Row;
@@ -131,6 +133,22 @@ void badJavascriptForPartitions() {
         assertTrue(ex.getCause() instanceof FailedRequestException, "Unexpected cause: " + ex.getCause());
     }
 
+    @Test
+    void verifyTimeoutWorks() {
+        Dataset<Row> dataset = startRead()
+            .option(Options.READ_XQUERY, "(xdmp:sleep(1500), 'abc')")
+            .option("spark.marklogic.client.callTimeout", "1")
+            .load();
+
+        SparkException ex = assertThrows(SparkException.class, () -> dataset.count());
+        assertTrue(ex.getCause() instanceof MarkLogicIOException, "Unexpected cause: " + ex.getCause());
+        assertTrue(ex.getCause().getMessage().contains("timeout"),
+            "Expecting a timeout due to the callTimeout being set to 1sec. Curiously, the Java Client does not " +
+                "attempt a retry on this sort of exception; it only attempts a retry if an HTTP response is received. " +
+                "Actual message: " + ex.getCause().getMessage());
+    }
+
+
     private List<Row> readRows(String option, String value) {
         return startRead()
             .option(option, value)