get forest configuration even in load balancer case #999

Author: ehennum

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

@@ -16,18 +16,13 @@
 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.GatewayForestConfiguration;
 import com.marklogic.client.datamovement.impl.ForestImpl;
 
-import java.util.ArrayList;
-import java.util.Arrays;
 import java.util.HashMap;
 import java.util.HashSet;
-import java.util.List;
 import java.util.Map;
 import java.util.Set;
-import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 /**
@@ -66,8 +61,8 @@ 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");
+    if (forestConfig instanceof GatewayForestConfiguration) {
+      throw new IllegalArgumentException("cannot filter the forest hosts when connecting using a gateway");
     }
     this.wrappedForestConfig = forestConfig;
   }

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

@@ -138,6 +138,7 @@ public HostAvailabilityListener withMinHosts(int numHosts) {
       }
     } else {
       if (numHosts <= 0) throw new IllegalArgumentException("numHosts must be > 0");
+// TODO: use existing forest configuration instead of refreshing?
       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");
@@ -274,23 +275,7 @@ private boolean processForestHostException(Batcher batcher, Throwable throwable,
           filteredForestConfig = filteredForestConfig.withRenamedHost(host, randomAvailableHost);
         }
         batcher.withForestConfig(filteredForestConfig);
-        // cancel any previously scheduled re-sync
-        if ( future != null ) future.cancel(false);
-        // schedule a re-sync with the server forest config
-        future = Executors.newScheduledThreadPool(1)
-          .schedule( () -> {
-              if ( batcher.isStopped() ) {
-                logger.debug("Job \"{}\" is stopped, so cancelling re-sync with the server forest config",
-                  batcher.getJobName());
-              } else {
-                ForestConfiguration updatedForestConfig = moveMgr.readForestConfig();
-                logger.info("it's been {} since host {} failed, opening communication to all server hosts [{}]",
-                  suspendTimeForHostUnavailable.toString(), host, Arrays.asList(updatedForestConfig.getPreferredHosts()));
-                // set the forestConfig back to whatever the server says it is
-                batcher.withForestConfig(updatedForestConfig);
-              }
-            }
-            , suspendTimeForHostUnavailable.toMillis(), TimeUnit.MILLISECONDS);
+        scheduleForestResynch(batcher, host);
       } else {
         // by black-listing this host we'd move below minHosts, so it's time to
         // stop this job
@@ -304,6 +289,26 @@ private boolean processForestHostException(Batcher batcher, Throwable throwable,
     return shouldWeRetry;
   }
 
+  private void scheduleForestResynch(Batcher batcher, String host) {
+    // cancel any previously scheduled re-sync
+    if ( future != null ) future.cancel(false);
+    // schedule a re-sync with the server forest config
+    future = Executors.newScheduledThreadPool(1)
+      .schedule( () -> {
+          if ( batcher.isStopped() ) {
+            logger.debug("Job \"{}\" is stopped, so cancelling re-sync with the server forest config",
+              batcher.getJobName());
+          } else {
+            ForestConfiguration updatedForestConfig = moveMgr.readForestConfig();
+            logger.info("it's been {} since host {} failed, opening communication to all server hosts [{}]",
+              suspendTimeForHostUnavailable.toString(), host, Arrays.asList(updatedForestConfig.getPreferredHosts()));
+            // set the forestConfig back to whatever the server says it is
+            batcher.withForestConfig(updatedForestConfig);
+          }
+        }
+        , suspendTimeForHostUnavailable.toMillis(), TimeUnit.MILLISECONDS);
+  }
+
   protected boolean isHostUnavailableException(Throwable throwable, Set<Throwable> path) {
     for ( Class<?> type : hostUnavailableExceptions ) {
       if ( type.isInstance(throwable) ) {

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

@@ -47,16 +47,10 @@ public class DataMovementManagerImpl implements DataMovementManager {
   private static ConcurrentHashMap<String, JobTicket> activeJobs = new ConcurrentHashMap<>();
   private ForestConfiguration forestConfig;
   private DatabaseClient primaryClient;
-  private DatabaseClient.ConnectionType connectionType;
   // clientMap key is the hostname_database
   private Map<String,DatabaseClient> clientMap = new HashMap<>();
 
   public DataMovementManagerImpl(DatabaseClient client) {
-    connectionType = client.getConnectionType();
-    if (connectionType == DatabaseClient.ConnectionType.GATEWAY) {
-      forestConfig = new AnyForestConfiguration(client);
-    }
-
     setPrimaryClient(client);
 
     clientMap.put(primaryClient.getHost(), primaryClient);
@@ -163,9 +157,7 @@ private ForestConfiguration getForestConfig() {
 
   @Override
   public ForestConfiguration readForestConfig() {
-    if (getConnectionType() == DatabaseClient.ConnectionType.DIRECT) {
-     forestConfig = service.readForestConfig();
-    }
+    forestConfig = service.readForestConfig();
     return forestConfig;
   }
 
@@ -204,7 +196,7 @@ public JobTicket getActiveJob(String jobId) {
 
   @Override
   public DatabaseClient.ConnectionType getConnectionType() {
-    return connectionType;
+    return primaryClient.getConnectionType();
   }
 
   public DataMovementServices getDataMovementServices() {

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

@@ -21,15 +21,10 @@
 
 import com.fasterxml.jackson.databind.JsonNode;
 import com.marklogic.client.DatabaseClient;
+import com.marklogic.client.datamovement.*;
 import com.marklogic.client.impl.DatabaseClientImpl;
 import com.marklogic.client.io.JacksonHandle;
-import com.marklogic.client.datamovement.Batcher;
-import com.marklogic.client.datamovement.DataMovementException;
-import com.marklogic.client.datamovement.JobTicket;
 import com.marklogic.client.datamovement.JobTicket.JobType;
-import com.marklogic.client.datamovement.QueryBatcher;
-import com.marklogic.client.datamovement.WriteBatcher;
-import com.marklogic.client.datamovement.JobReport;
 import com.marklogic.client.datamovement.impl.ForestConfigurationImpl;
 import java.util.List;
 
@@ -46,6 +41,8 @@ public DataMovementServices setClient(DatabaseClient client) {
   }
 
   public ForestConfigurationImpl readForestConfig() {
+    DatabaseClient.ConnectionType connectionType = client.getConnectionType();
+
     List<ForestImpl> forests = new ArrayList<>();
     JsonNode results = ((DatabaseClientImpl) client).getServices()
       .getResource(null, "internal/forestinfo", null, null, new JacksonHandle())
@@ -54,24 +51,34 @@ public ForestConfigurationImpl readForestConfig() {
       String id = forestNode.get("id").asText();
       String name = forestNode.get("name").asText();
       String database = forestNode.get("database").asText();
-      String host = forestNode.get("host").asText();
-      String openReplicaHost = null;
-      if ( forestNode.get("openReplicaHost") != null ) openReplicaHost = forestNode.get("openReplicaHost").asText();
-      String requestHost = null;
-      if ( forestNode.get("requestHost") != null ) requestHost = forestNode.get("requestHost").asText();
-      String alternateHost = null;
-      if ( forestNode.get("alternateHost") != null ) alternateHost = forestNode.get("alternateHost").asText();
-      // Since we added the forestinfo end point to populate both alternateHost and requestHost
-      // in case we have a requestHost so that we don't break the existing API code, we will make the 
-      // alternateHost as null if both alternateHost and requestHost is set.
-      if ( requestHost != null && alternateHost != null )
-        alternateHost = null;
       boolean isUpdateable = "all".equals(forestNode.get("updatesAllowed").asText());
       boolean isDeleteOnly = false; // TODO: get this for real after we start using a REST endpoint
-      forests.add(new ForestImpl(host, openReplicaHost, requestHost, alternateHost, database, name, id, isUpdateable, isDeleteOnly));
+      if (connectionType == DatabaseClient.ConnectionType.GATEWAY) {
+        forests.add(
+            new GatewayForestConfiguration.GatewayForest(client.getHost(), database, name, id, isUpdateable, isDeleteOnly)
+        );
+      } else {
+        String host = forestNode.get("host").asText();
+        String openReplicaHost = null;
+        if ( forestNode.get("openReplicaHost") != null ) openReplicaHost = forestNode.get("openReplicaHost").asText();
+        String requestHost = null;
+        if ( forestNode.get("requestHost") != null ) requestHost = forestNode.get("requestHost").asText();
+        String alternateHost = null;
+        if ( forestNode.get("alternateHost") != null ) alternateHost = forestNode.get("alternateHost").asText();
+        // Since we added the forestinfo end point to populate both alternateHost and requestHost
+        // in case we have a requestHost so that we don't break the existing API code, we will make the
+        // alternateHost as null if both alternateHost and requestHost is set.
+        if ( requestHost != null && alternateHost != null )
+          alternateHost = null;
+        forests.add(
+            new ForestImpl(host, openReplicaHost, requestHost, alternateHost, database, name, id, isUpdateable, isDeleteOnly)
+        );
+      }
     }
 
-    return new ForestConfigurationImpl(forests.toArray(new ForestImpl[forests.size()]));
+    Forest[] forestdefs = forests.toArray(new ForestImpl[forests.size()]);
+    return (connectionType == DatabaseClient.ConnectionType.GATEWAY) ?
+          new GatewayForestConfiguration(client.getHost(), forestdefs) : new ForestConfigurationImpl(forestdefs);
   }
 
   public JobTicket startJob(WriteBatcher batcher, ConcurrentHashMap<String, JobTicket> activeJobs) {

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

@@ -0,0 +1,29 @@
+package com.marklogic.client.datamovement.impl;
+
+import com.marklogic.client.DatabaseClient;
+import com.marklogic.client.datamovement.Forest;
+import com.marklogic.client.datamovement.ForestConfiguration;
+
+import java.util.stream.Stream;
+
+public class GatewayForestConfiguration extends ForestConfigurationImpl {
+   private String[] preferredHosts;
+
+   GatewayForestConfiguration(String primaryHost, Forest[] forests) {
+      super(forests);
+      if (primaryHost == null) throw new IllegalArgumentException("primaryHost argument must not be null");
+
+      this.preferredHosts = new String[]{primaryHost};
+   }
+
+   @Override
+   public String[] getPreferredHosts() {
+      return preferredHosts;
+   }
+
+   static public class GatewayForest extends ForestImpl {
+      public GatewayForest(String primaryHost, String databaseName, String forestName, String forestId, boolean isUpdateable, boolean isDeleteOnly) {
+         super(primaryHost, primaryHost, primaryHost, primaryHost, databaseName, forestName, forestId, isUpdateable, isDeleteOnly);
+      }
+   }
+}

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

@@ -627,9 +627,7 @@ 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 instanceof AnyForestConfiguration.AnyForest) ? null : forest.getForestName()
-        ) ) {
+        try ( UrisHandle results = queryMgr.uris(query, handle, start, 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 ) {