connect only to primary host for load balancer scenario #999

Author: ehennum

marklogic-client-api/src/main/java/com/marklogic/client/DatabaseClient.java

@@ -99,12 +99,22 @@ public interface DatabaseClient {
    */
   XMLDocumentManager newXMLDocumentManager();
 
+  /**
+   * Creates a manager for long-running asynchronous write or query jobs
+   * with the default ConnectionPolicy.FOREST_HOSTS policy of creating
+   * a new connection for each host that has forests for the database.
+   * Don't forget to call dataMovementManager.release() when you're done with it.
+   * @return	a manager supporting long-running asynchronous write or query jobs
+   */
+  DataMovementManager newDataMovementManager();
+
   /**
    * Creates a manager for long-running asynchronous write or query jobs.  Don't forget
    * to call dataMovementManager.release() when you're done with it.
+   * @param policy	how the DataMovementManager should create connections
    * @return	a manager supporting long-running asynchronous write or query jobs
    */
-  DataMovementManager newDataMovementManager();
+  DataMovementManager newDataMovementManager(ConnectionPolicy policy);
 
   /**
    * Creates a manager to query for database documents.
@@ -212,6 +222,12 @@ public interface DatabaseClient {
    */
   ServerEvaluationCall newServerEval();
 
+  /**
+   * Specifies whether to connect only using the host from the initial database connection
+   * or to create a new connection for each host that has forests for the database.
+   */
+  public enum ConnectionPolicy{PRIMARY_HOST, FOREST_HOSTS}
+
   String getHost();
 
   int getPort();

marklogic-client-api/src/main/java/com/marklogic/client/datamovement/DataMovementManager.java

@@ -184,4 +184,11 @@ public interface DataMovementManager {
    * @return the latest ForestConfiguration from the server
    */
   public ForestConfiguration readForestConfig();
+
+  /**
+   * Identify the DataMovementManager uses only the primary database connection
+   * or manages a connection for each host that has forests for the database.
+   * @return the connection policy
+   */
+  public DatabaseClient.ConnectionPolicy getConnectionPolicy();
 }

marklogic-client-api/src/main/java/com/marklogic/client/datamovement/FilteredForestConfiguration.java

@@ -16,6 +16,7 @@
 package com.marklogic.client.datamovement;
 
 import com.marklogic.client.datamovement.Forest.HostType;
+import com.marklogic.client.datamovement.impl.AnyForestConfiguration;
 import com.marklogic.client.datamovement.impl.ForestConfigurationImpl;
 import com.marklogic.client.datamovement.impl.ForestImpl;
 
@@ -65,6 +66,9 @@ public class FilteredForestConfiguration implements ForestConfiguration {
    * @param forestConfig the ForestConfiguration to wrap
    */
   public FilteredForestConfiguration(ForestConfiguration forestConfig) {
+    if (forestConfig instanceof AnyForestConfiguration) {
+      throw new IllegalArgumentException("cannot filter the forest hosts when using only the primary host");
+    }
     this.wrappedForestConfig = forestConfig;
   }
 

marklogic-client-api/src/main/java/com/marklogic/client/datamovement/HostAvailabilityListener.java

@@ -15,6 +15,7 @@
  */
 package com.marklogic.client.datamovement;
 
+import com.marklogic.client.DatabaseClient;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -38,7 +39,7 @@
 /** <p>HostAvailabilityListener is automatically registered with all QueryBatcher
  * and WriteBatcher instances to monitor for failover scenarios.  When
  * HostAvailabilityListener detects that a host is unavailable (matches one of
- * {@link #getHostUnavailableExceptions()}), it black-lists the host for a
+ * {@link #getHostUnavailableExceptions()}), it blacklists the MarkLogic host for a
  * period of time equal to {@link #getSuspendTimeForHostUnavailable()}.  After
  * that time, it calls {@link DataMovementManager#readForestConfig()} then
  * passes that updated ForestConfiguration to batcher.withForestConfig() so the
@@ -101,6 +102,8 @@ public void processFailure(QueryBatch batch, Throwable throwable) {
   }
 
   /**
+   * Manages refreshing the connection hosts and retrying events after a host
+   * becomes unavailable.
    * @param moveMgr the DataMovementManager (used to call readForestConfig to reset after black-listing an unavailable host)
    */
   public HostAvailabilityListener(DataMovementManager moveMgr) {
@@ -129,15 +132,22 @@ public HostAvailabilityListener withSuspendTimeForHostUnavailable(Duration
    * @return this instance (for method chaining)
    */
   public HostAvailabilityListener withMinHosts(int numHosts) {
-    if (numHosts <= 0) throw new IllegalArgumentException("numHosts must be > 0");
-    int numConfigHosts = moveMgr.readForestConfig().getPreferredHosts().length;
-    if (numHosts > numConfigHosts) throw new IllegalArgumentException
-      ("numHosts must be less than or equal to the number of hosts in the cluster");
+    if (moveMgr.getConnectionPolicy() == DatabaseClient.ConnectionPolicy.PRIMARY_HOST) {
+      if (numHosts != 0) {
+        throw new IllegalArgumentException("numHosts must be 1 when using only the primary host for the connection");
+      }
+    } else {
+      if (numHosts <= 0) throw new IllegalArgumentException("numHosts must be > 0");
+      int numConfigHosts = moveMgr.readForestConfig().getPreferredHosts().length;
+      if (numHosts > numConfigHosts) throw new IllegalArgumentException
+          ("numHosts must be less than or equal to the number of hosts in the cluster");
+    }
     this.minHosts = numHosts;
     return this;
   }
 
-  /** Overwrites the list of exceptions for which a host will be blacklisted
+  /** Overwrites the list of exceptions for which a request can be retried and
+   * a MarkLogic host can be blacklisted
    *
    * @param exceptionTypes the list of types of Throwable, any of which constitute a host that's unavailable
    *
@@ -152,7 +162,7 @@ public HostAvailabilityListener withHostUnavailableExceptions(Class<Throwable>..
   }
 
   /**
-   * @return the list of types of Throwable, any of which constitute a host that's unavailable
+   * @return the list of types of Throwable, any of which constitute a MarkLogic host that's unavailable
    */
   public Throwable[] getHostUnavailableExceptions() {
     return hostUnavailableExceptions.toArray(new Throwable[hostUnavailableExceptions.size()]);
@@ -217,6 +227,19 @@ private synchronized boolean processException(Batcher batcher, Throwable throwab
     // we only do something if this throwable is on our list of exceptions
     // which we consider marking a host as unavilable
     boolean isHostUnavailableException = isHostUnavailableException(throwable, new HashSet<>());
+    return (moveMgr.getConnectionPolicy() == DatabaseClient.ConnectionPolicy.PRIMARY_HOST) ?
+           processPrimaryHostException(batcher, throwable, host, isHostUnavailableException) :
+           processForestHostException(batcher, throwable, host, isHostUnavailableException);
+  }
+
+  private boolean processPrimaryHostException(Batcher batcher, Throwable throwable, String host, boolean isHostUnavailableException) {
+    // TODO:  currently, cancelling the job because the load balancer is down would require
+    // a custom listener; determine whether the approach could be simplified with an array
+    // of exceptions different from the exceptions that indicate the MarkLogic host is unavailable
+    return isHostUnavailableException;
+  }
+
+  private boolean processForestHostException(Batcher batcher, Throwable throwable, String host, boolean isHostUnavailableException) {
     boolean shouldWeRetry = isHostUnavailableException;
     if ( isHostUnavailableException == true ) {
       ForestConfiguration existingForestConfig = batcher.getForestConfig();

marklogic-client-api/src/main/java/com/marklogic/client/datamovement/impl/AnyForestConfiguration.java

@@ -0,0 +1,36 @@
+package com.marklogic.client.datamovement.impl;
+
+import com.marklogic.client.DatabaseClient;
+import com.marklogic.client.datamovement.Forest;
+import com.marklogic.client.datamovement.ForestConfiguration;
+
+public class AnyForestConfiguration implements ForestConfiguration {
+   private Forest[] anyForests;
+   private String[] primaryHosts;
+
+   AnyForestConfiguration(DatabaseClient client) {
+      if (client == null) throw new IllegalArgumentException("client argument must not be null");
+
+      String primaryHost = client.getHost();
+      String database    = client.getDatabase();
+
+      this.anyForests   = new Forest[]{new AnyForest(primaryHost, database,true, false)};
+      this.primaryHosts = new String[]{primaryHost};
+   }
+
+   @Override
+   public Forest[] listForests() {
+      return anyForests;
+   }
+   @Override
+   public String[] getPreferredHosts() {
+      return primaryHosts;
+   }
+
+   static public class AnyForest extends ForestImpl {
+      private AnyForest(String primaryHost, String databaseName, boolean isUpdateable, boolean isDeleteOnly) {
+         super(primaryHost, primaryHost, primaryHost, primaryHost, (databaseName == null) ? "" : databaseName,
+               "*ANY*", "", isUpdateable, isDeleteOnly);
+      }
+   }
+}

marklogic-client-api/src/main/java/com/marklogic/client/datamovement/impl/DataMovementManagerImpl.java

@@ -54,12 +54,17 @@ public class DataMovementManagerImpl implements DataMovementManager {
   private static ConcurrentHashMap<String, JobTicket> activeJobs = new ConcurrentHashMap<>();
   private ForestConfiguration forestConfig;
   private DatabaseClient primaryClient;
+  private DatabaseClient.ConnectionPolicy connectPolicy;
   // clientMap key is the hostname_database
   private Map<String,DatabaseClient> clientMap = new HashMap<>();
 
-  public DataMovementManagerImpl(DatabaseClient client) {
+  public DataMovementManagerImpl(DatabaseClient client, DatabaseClient.ConnectionPolicy connectPolicy) {
     setPrimaryClient(client);
     clientMap.put(primaryClient.getHost(), primaryClient);
+    this.connectPolicy = connectPolicy;
+    if (connectPolicy == DatabaseClient.ConnectionPolicy.PRIMARY_HOST) {
+      forestConfig = new AnyForestConfiguration(client);
+    }
   }
 
   @Override
@@ -137,19 +142,16 @@ public QueryBatcher newQueryBatcher(RawCombinedQueryDefinition query) {
 
   private QueryBatcher newQueryBatcherImpl(QueryDefinition query) {
     if ( query == null ) throw new IllegalArgumentException("query must not be null");
-    QueryBatcherImpl batcher = new QueryBatcherImpl(query, this, getForestConfig());
-    batcher.onQueryFailure(new HostAvailabilityListener(this));
-    QueryJobReportListener queryJobListener = new QueryJobReportListener();
-    batcher.onQueryFailure(queryJobListener);
-    batcher.onQueryFailure(new NoResponseListener(this));
-    batcher.onUrisReady(queryJobListener);
-    return batcher;
+    return newQueryBatcher(new QueryBatcherImpl(query, this, getForestConfig()));
   }
 
   @Override
   public QueryBatcher newQueryBatcher(Iterator<String> iterator) {
     if ( iterator == null ) throw new IllegalArgumentException("iterator must not be null");
-    QueryBatcherImpl batcher = new QueryBatcherImpl(iterator, this, getForestConfig());
+    return newQueryBatcher(new QueryBatcherImpl(iterator, this, getForestConfig()));
+  }
+
+  private QueryBatcher newQueryBatcher(QueryBatcherImpl batcher) {
     // add a default listener to handle host failover scenarios
     batcher.onQueryFailure(new HostAvailabilityListener(this));
     QueryJobReportListener queryJobListener = new QueryJobReportListener();
@@ -166,7 +168,9 @@ private ForestConfiguration getForestConfig() {
 
   @Override
   public ForestConfiguration readForestConfig() {
-    forestConfig = service.readForestConfig();
+    if (connectPolicy == DatabaseClient.ConnectionPolicy.FOREST_HOSTS) {
+     forestConfig = service.readForestConfig();
+    }
     return forestConfig;
   }
 
@@ -203,6 +207,11 @@ public JobTicket getActiveJob(String jobId) {
     }
   }
 
+  @Override
+  public DatabaseClient.ConnectionPolicy getConnectionPolicy() {
+    return connectPolicy;
+  }
+
   public DataMovementServices getDataMovementServices() {
     return service;
   }

marklogic-client-api/src/main/java/com/marklogic/client/datamovement/impl/QueryBatcherImpl.java

@@ -627,7 +627,9 @@ public void run() {
         // this try-with-resources block will call results.close() once the block is done
         // here we call the /v1/internal/uris endpoint to get the text/uri-list of documents
         // matching this structured or string query
-        try ( UrisHandle results = queryMgr.uris(query, handle, start, null, forest.getForestName()) ) {
+        try ( UrisHandle results = queryMgr.uris(
+              query, handle, start, null, (forest instanceof AnyForestConfiguration.AnyForest) ? null : forest.getForestName()
+        ) ) {
           // if we're doing consistentSnapshot and this is the first result set, let's capture the
           // serverTimestamp so we can use it for all future queries
           if ( consistentSnapshot == true && serverTimestamp.get() == -1 ) {
@@ -712,10 +714,8 @@ private void shutdownIfAllForestsAreDone() {
     }
     // if we made it this far, all forests are done. let's run the Job
     // completion listeners and shutdown.
-    if ( !runJobCompletionListeners.isAlive() ) {
-      synchronized(this) {
-        if ( !runJobCompletionListeners.isAlive() ) runJobCompletionListeners.start();
-      }
+    synchronized(this) {
+      if ( !runJobCompletionListeners.isAlive() ) runJobCompletionListeners.start();
     }
     threadPool.shutdown();
   }

marklogic-client-api/src/main/java/com/marklogic/client/impl/DatabaseClientImpl.java

@@ -136,7 +136,14 @@ public QueryManager newQueryManager() {
   }
   @Override
   public DataMovementManager newDataMovementManager() {
-    DataMovementManagerImpl moveMgr = new DataMovementManagerImpl(this);
+    return newDataMovementManager(null);
+  }
+  @Override
+  public DataMovementManager newDataMovementManager(ConnectionPolicy policy) {
+    DataMovementManagerImpl moveMgr = new DataMovementManagerImpl(
+        this,
+        (policy == null) ? ConnectionPolicy.FOREST_HOSTS : policy
+    );
     return moveMgr;
   }
   @Override

marklogic-client-api/src/test/java/com/marklogic/client/test/Common.java

@@ -61,11 +61,14 @@ public class Common {
   final public static String WRITE_PRIVILIGED_PASS = "x";
 
   final public static String  HOST          = System.getProperty("TEST_HOST", "localhost");
+
   final public static int     PORT          = Integer.parseInt(System.getProperty("TEST_PORT", "8012"));
-  final public static boolean PLACE_DIRECT  = Boolean.parseBoolean(System.getProperty("TEST_PLACE_DIRECT", "true"));
   final public static boolean WITH_WAIT     = Boolean.parseBoolean(System.getProperty("TEST_WAIT", "false"));
   final public static int     PROPERTY_WAIT = Integer.parseInt(System.getProperty("TEST_PROPERTY_WAIT", WITH_WAIT ? "8200" : "0"));
 
+  final public static DatabaseClient.ConnectionPolicy CONNECT_POLICY =
+      DatabaseClient.ConnectionPolicy.valueOf(System.getProperty("TEST_CONNECT_POLICY", "FOREST_HOSTS"));
+
   public static DatabaseClient client;
   public static DatabaseClient adminClient;
   public static DatabaseClient serverAdminClient;
@@ -197,20 +200,6 @@ public static Document testStringToDocument(String document) {
       throw new RuntimeException(e);
     }
   }
-  public static <T extends Batcher> T initBatcher(DataMovementManager moveMgr, T batcher) {
-    if (moveMgr == null || batcher == null || PLACE_DIRECT == true) {
-      return batcher;
-    }
-    return (T) batcher.withForestConfig(
-          new FilteredForestConfiguration(moveMgr.readForestConfig()).withWhiteList(HOST)
-    );
-  }
-  public static ForestConfiguration initForestConfig(ForestConfiguration forestConfig) {
-    if (forestConfig == null || PLACE_DIRECT == true) {
-      return forestConfig;
-    }
-    return new FilteredForestConfiguration(forestConfig).withWhiteList(HOST);
-  }
   public static void propertyWait() {
     waitFor(PROPERTY_WAIT);
   }

marklogic-client-api/src/test/java/com/marklogic/client/test/datamovement/ApplyTransformTest.java

@@ -43,7 +43,7 @@
 
 public class ApplyTransformTest {
   private static DatabaseClient client = Common.connect();
-  private static DataMovementManager moveMgr = client.newDataMovementManager();
+  private static DataMovementManager moveMgr = client.newDataMovementManager(Common.CONNECT_POLICY);
   private static GenericDocumentManager docMgr = client.newDocumentManager();
   private static QueryManager queryMgr = client.newQueryManager();
   private static StructuredQueryBuilder sqb = new StructuredQueryBuilder();
@@ -85,7 +85,7 @@ public void testResultReplace() throws Exception {
     ApplyTransformListener listener = new ApplyTransformListener()
       .withTransform(transform)
       .withApplyResult(ApplyResult.REPLACE);
-    QueryBatcher batcher = Common.initBatcher(moveMgr, moveMgr.newQueryBatcher(query))
+    QueryBatcher batcher = moveMgr.newQueryBatcher(query)
       .onUrisReady(listener);
     JobTicket ticket = moveMgr.startJob( batcher );
     batcher.awaitCompletion();
@@ -108,7 +108,7 @@ public void testResultIgnore() throws Exception {
     ApplyTransformListener listener = new ApplyTransformListener()
       .withTransform(transform)
       .withApplyResult(ApplyResult.IGNORE);
-    QueryBatcher batcher = Common.initBatcher(moveMgr, moveMgr.newQueryBatcher(query))
+    QueryBatcher batcher = moveMgr.newQueryBatcher(query)
       .onUrisReady(listener);
     JobTicket ticket = moveMgr.startJob( batcher );
     batcher.awaitCompletion();
@@ -124,7 +124,7 @@ public void testManyDocs() throws Exception {
     DocumentMetadataHandle meta = new DocumentMetadataHandle().withCollections(collection);
     int numDocs = 1000;
     // write the documents
-    WriteBatcher batcher1 = Common.initBatcher(moveMgr, moveMgr.newWriteBatcher())
+    WriteBatcher batcher1 = moveMgr.newWriteBatcher()
       .withBatchSize(100)
       .onBatchFailure((batch, throwable) -> throwable.printStackTrace());
     JobTicket ticket1 = moveMgr.startJob( batcher1 );
@@ -145,7 +145,7 @@ public void testManyDocs() throws Exception {
       .withApplyResult(ApplyResult.REPLACE)
       .onSuccess(batch -> count2.addAndGet(batch.getItems().length))
       .onBatchFailure((batch, throwable) -> throwable.printStackTrace());
-    QueryBatcher batcher = Common.initBatcher(moveMgr, moveMgr.newQueryBatcher(query2))
+    QueryBatcher batcher = moveMgr.newQueryBatcher(query2)
       .onUrisReady(listener)
       .withConsistentSnapshot();
     JobTicket ticket2 = moveMgr.startJob( batcher );
@@ -157,7 +157,7 @@ public void testManyDocs() throws Exception {
     StructuredQueryDefinition query3 = sqb.value(sqb.jsonProperty("testProperty"), "test3a");
     query3.setCollections(collection);
     final AtomicInteger count3 = new AtomicInteger(0);
-    QueryBatcher batcher3 = Common.initBatcher(moveMgr, moveMgr.newQueryBatcher(query3))
+    QueryBatcher batcher3 = moveMgr.newQueryBatcher(query3)
       .withBatchSize(100)
       .onUrisReady(batch -> count3.addAndGet(batch.getItems().length))
       .onQueryFailure((throwable) -> throwable.printStackTrace());
@@ -183,7 +183,7 @@ public void testOnSkipped() throws Exception {
     ServerTransform transform = new ServerTransform(transformName1);
     List skippedUris = new ArrayList<>();
     StringBuilder failures = new StringBuilder();
-    QueryBatcher batcher = Common.initBatcher(moveMgr, moveMgr.newQueryBatcher(uris.iterator()))
+    QueryBatcher batcher = moveMgr.newQueryBatcher(uris.iterator())
       .withBatchSize(1)
       .onUrisReady(
         new ApplyTransformListener()