#838 - Calling retry with failure listeners rather than a single retry

Author: vivekmuniyandi

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

@@ -182,7 +182,7 @@ public void processFailure(WriteBatch batch, Throwable throwable) {
         logger.warn("Retrying failed batch: {}, results so far: {}, uris: {}",
           batch.getJobBatchNumber(), batch.getJobWritesSoFar(),
           Stream.of(batch.getItems()).map(event->event.getTargetUri()).collect(Collectors.toList()));
-        batch.getBatcher().retry(batch);
+        batch.getBatcher().retryWithFailureListeners(batch);
       } catch (RuntimeException e) {
         logger.error("Exception during retry", e);
         processFailure(batch, e);
@@ -202,7 +202,7 @@ public void processFailure(QueryBatchException queryBatch) {
         logger.warn("Retrying failed batch: {}, results so far: {}, forest: {}, forestBatch: {}, forest results so far: {}",
           queryBatch.getJobBatchNumber(), queryBatch.getJobResultsSoFar(), queryBatch.getForest().getForestName(),
           queryBatch.getForestBatchNumber(), queryBatch.getForestResultsSoFar());
-        queryBatch.getBatcher().retry(queryBatch);
+        queryBatch.getBatcher().retryWithFailureListeners(queryBatch);
       } catch (RuntimeException e) {
         logger.error("Exception during retry", e);
         processFailure(new QueryBatchException(queryBatch, e));

src/main/java/com/marklogic/client/datamovement/QueryBatcher.java

@@ -310,4 +310,18 @@ public interface QueryBatcher extends Batcher {
    * @param queryBatchListener  the QueryBatchListener which needs to be applied
    */
   void retryListener(QueryBatch batch, QueryBatchListener queryBatchListener);
+
+  /**
+   * Retry in the same thread to query a batch that failed. If it fails again,
+   * all the failure listeners associated with the batcher using onQueryFailure
+   * method would be processed.
+   * 
+   * Note : Use this method with caution as there is a possibility of infinite
+   * loops. If a batch fails and one of the failure listeners calls this method
+   * to retry with failure listeners and if the batch again fails, this would go
+   * on as an infinite loop until the batch succeeds.
+   * 
+   * @param queryEvent the information about the batch that failed
+   */
+  void retryWithFailureListeners(QueryEvent queryEvent);
 }

src/main/java/com/marklogic/client/datamovement/WriteBatcher.java

@@ -321,4 +321,18 @@ WriteBatcher addAs(String uri, DocumentMetadataWriteHandle metadataHandle,
    * @throws IllegalStateException if this job has not yet been started
    */
   JobTicket getJobTicket();
+
+  /**
+   * Retry in the same thread to send a batch that failed. If it fails again,
+   * all the failure listeners associated with the batcher using onBatchFailure
+   * method would be processed.
+   * 
+   * Note : Use this method with caution as there is a possibility of infinite
+   * loops. If a batch fails and one of the failure listeners calls this method
+   * to retry with failure listeners and if the batch again fails, this would go
+   * on as an infinite loop until the batch succeeds.
+   * 
+   * @param queryEvent the information about the batch that failed
+   */
+  public void retryWithFailureListeners(WriteBatch writeBatch);
 }

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

@@ -51,7 +51,7 @@
 public class DataMovementManagerImpl implements DataMovementManager {
   private static Logger logger = LoggerFactory.getLogger(DataMovementManager.class);
   private DataMovementServices service = new DataMovementServices();
-  private ConcurrentHashMap<String, JobTicket> activeJobs = new ConcurrentHashMap<>();
+  private static ConcurrentHashMap<String, JobTicket> activeJobs = new ConcurrentHashMap<>();
   private ForestConfiguration forestConfig;
   private DatabaseClient primaryClient;
   // clientMap key is the hostname_database

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

@@ -124,11 +124,19 @@ public QueryBatcherImpl onQueryFailure(QueryFailureListener listener) {
    */
   @Override
   public void retry(QueryEvent queryEvent) {
+    retry(queryEvent, false);
+  }
+
+  @Override
+  public void retryWithFailureListeners(QueryEvent queryEvent) {
+    retry(queryEvent, true);
+  }
+
+  private void retry(QueryEvent queryEvent, boolean callFailListeners) {
     if ( isStopped() == true ) {
       logger.warn("Job is now stopped, aborting the retry");
       return;
     }
-    boolean callFailListeners = false;
     Forest retryForest = null;
     for ( Forest forest : getForestConfig().listForests() ) {
       if ( forest.equals(queryEvent.getForest()) ) {
@@ -151,7 +159,6 @@ public void retry(QueryEvent queryEvent) {
       queryEvent.getForestBatchNumber(), start, queryEvent.getJobBatchNumber(), callFailListeners);
     runnable.run();
   }
-
   /*
    * Accepts a QueryBatch which was successfully retrieved from the server and a 
    * QueryBatchListener which was failed to apply and retry that listener on the batch. 

src/main/java/com/marklogic/client/datamovement/impl/WriteBatcherImpl.java

@@ -461,26 +461,38 @@ public WriteBatcher onBatchFailure(WriteFailureListener listener) {
     return this;
   }
 
+  @Override
+  public void retryWithFailureListeners(WriteBatch batch) {
+    retry(batch, true);
+  }
+
   @Override
   public void retry(WriteBatch batch) {
+    retry(batch, false);
+  }
+
+  private void retry(WriteBatch batch, boolean callFailListeners) {
     if ( isStopped() == true ) {
       logger.warn("Job is now stopped, aborting the retry");
       return;
     }
     if ( batch == null ) throw new IllegalArgumentException("batch must not be null");
     boolean forceNewTransaction = true;
     BatchWriteSet writeSet = newBatchWriteSet(forceNewTransaction, batch.getJobBatchNumber());
-    writeSet.onFailure(throwable -> {
-      if ( throwable instanceof RuntimeException ) throw (RuntimeException) throwable;
-      else throw new DataMovementException("Failed to retry batch", throwable);
-    });
-    for ( WriteEvent doc : batch.getItems() ) {
+    if ( !callFailListeners ) {
+      writeSet.onFailure(throwable -> {
+        if ( throwable instanceof RuntimeException )
+          throw (RuntimeException) throwable;
+        else
+          throw new DataMovementException("Failed to retry batch", throwable);
+      });
+    }
+    for (WriteEvent doc : batch.getItems()) {
       writeSet.getWriteSet().add(doc.getTargetUri(), doc.getMetadata(), doc.getContent());
     }
     BatchWriter runnable = new BatchWriter(writeSet);
     runnable.run();
   }
-
   @Override
   public WriteBatchListener[]        getBatchSuccessListeners() {
     return successListeners.toArray(new WriteBatchListener[successListeners.size()]);