(cherry picked from commit 256c44c)

Author: vivekmuniyandi

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

@@ -78,6 +78,7 @@ public class QueryBatcherImpl extends BatcherImpl implements QueryBatcher {
   private final AtomicBoolean started = new AtomicBoolean(false);
   private final Object lock = new Object();
   private final Map<Forest,List<QueryTask>> blackListedTasks = new HashMap<>();
+  private boolean isSingleThreaded = false;
   private JobTicket jobTicket;
   private Thread runJobCompletionListeners;
 
@@ -350,6 +351,12 @@ private synchronized void initialize() {
       // now we've set the threadCount
       threadCountSet = true;
     }
+    // If we are iterating and if we have the thread count to 1, we have a single thread acting as both
+    // consumer and producer of the ThreadPoolExecutor queue. Hence, we produce till the maximum and start
+    // consuming and produce again. Since the thread count is 1, there is no worry about thread utilization.
+    if(getThreadCount() == 1) {
+      isSingleThreaded = true;
+    }
     logger.info("Starting job batchSize={}, threadCount={}, onUrisReady listeners={}, failure listeners={}",
       getBatchSize(), getThreadCount(), urisReadyListeners.size(), failureListeners.size());
     threadPool = new QueryThreadPoolExecutor(getThreadCount(), this);
@@ -713,6 +720,91 @@ private void shutdownIfAllForestsAreDone() {
     threadPool.shutdown();
   }
 
+  private class IteratorTask implements Runnable {
+
+    private QueryBatcher batcher;
+
+    IteratorTask(QueryBatcher batcher) {
+      this.batcher = batcher;
+    }
+
+    @Override
+    public void run() {
+      try {
+        List<String> uriQueue = new ArrayList<>(getBatchSize());
+        while (iterator.hasNext()) {
+          uriQueue.add(iterator.next());
+          // if we've hit batchSize or the end of the iterator
+          if (uriQueue.size() == getBatchSize() || !iterator.hasNext()) {
+            final List<String> uris = uriQueue;
+            uriQueue = new ArrayList<>(getBatchSize());
+            Runnable processBatch = new Runnable() {
+              public void run() {
+                QueryBatchImpl batch = new QueryBatchImpl()
+                    .withBatcher(batcher)
+                    .withTimestamp(Calendar.getInstance())
+                    .withJobTicket(getJobTicket());
+                try {
+                  long currentBatchNumber = batchNumber.incrementAndGet();
+                  // round-robin from client 0 to (clientList.size() - 1);
+                  List<DatabaseClient> currentClientList = clientList.get();
+                  int clientIndex = (int) (currentBatchNumber % currentClientList.size());
+                  DatabaseClient client = currentClientList.get(clientIndex);
+                  batch = batch.withJobBatchNumber(currentBatchNumber)
+                      .withClient(client)
+                      .withJobResultsSoFar(resultsSoFar.addAndGet(uris.size()))
+                      .withItems(uris.toArray(new String[uris.size()]));
+                  logger.trace("batch size={}, jobBatchNumber={}, jobResultsSoFar={}", uris.size(),
+                      batch.getJobBatchNumber(), batch.getJobResultsSoFar());
+                  for (QueryBatchListener listener : urisReadyListeners) {
+                    try {
+                      listener.processEvent(batch);
+                    } catch (Throwable e) {
+                      logger.error("Exception thrown by an onUrisReady listener", e);
+                    }
+                  }
+                } catch (Throwable t) {
+                  batch = batch.withItems(uris.toArray(new String[uris.size()]));
+                  for (QueryFailureListener listener : failureListeners) {
+                    try {
+                      listener.processFailure(new QueryBatchException(batch, t));
+                    } catch (Throwable e) {
+                      logger.error("Exception thrown by an onQueryFailure listener", e);
+                    }
+                  }
+                  logger.warn("Error iterating to queue uris: {}", t.toString());
+                }
+              }
+            };
+            threadPool.execute(processBatch);
+            // If the queue is almost full, stop producing and add a task to continue later
+            if (isSingleThreaded && threadPool.getQueue().remainingCapacity() <= 2 && iterator.hasNext()) {
+              threadPool.execute(new IteratorTask(batcher));
+              return;
+            }
+          }
+        }
+      } catch (Throwable t) {
+        for (QueryFailureListener listener : failureListeners) {
+          QueryBatchImpl batch = new QueryBatchImpl()
+              .withItems(new String[0])
+              .withClient(clientList.get().get(0))
+              .withBatcher(batcher)
+              .withTimestamp(Calendar.getInstance())
+              .withJobResultsSoFar(0);
+
+          try {
+            listener.processFailure(new QueryBatchException(batch, t));
+          } catch (Throwable e) {
+            logger.error("Exception thrown by an onQueryFailure listener", e);
+          }
+        }
+        logger.warn("Error iterating to queue uris: {}", t.toString());
+      }
+      runJobCompletionListeners.start();
+      threadPool.shutdown();
+    }
+  }
 
   /* startIterating launches in a separate thread (actually a task handled by
    * threadPool) and just loops through the Iterator<String>, batching uris of
@@ -729,86 +821,7 @@ private void shutdownIfAllForestsAreDone() {
    * their listeners handled, they should use try-catch and handle them.
    */
   private void startIterating() {
-    final QueryBatcher batcher = this;
-    Runnable queueUris = new Runnable() {
-      public void run() {
-        try {
-          final AtomicLong batchNumber = new AtomicLong(0);
-          final AtomicLong resultsSoFar = new AtomicLong(0);
-          List<String> uriQueue = new ArrayList<>(getBatchSize());
-          while ( iterator.hasNext() ) {
-            try {
-              uriQueue.add(iterator.next());
-              // if we've hit batchSize or the end of the iterator
-              if ( uriQueue.size() == getBatchSize() || ! iterator.hasNext() ) {
-                final List<String> uris = uriQueue;
-                uriQueue = new ArrayList<>(getBatchSize());
-                Runnable processBatch = new Runnable() {
-                  public void run() {
-                    long currentBatchNumber = batchNumber.incrementAndGet();
-                    // round-robin from client 0 to (clientList.size() - 1);
-                    List<DatabaseClient> currentClientList = clientList.get();
-                    int clientIndex = (int) (currentBatchNumber % currentClientList.size());
-                    DatabaseClient client = currentClientList.get(clientIndex);
-                    QueryBatchImpl batch = new QueryBatchImpl()
-                      .withClient(client)
-                      .withBatcher(batcher)
-                      .withTimestamp(Calendar.getInstance())
-                      .withJobTicket(getJobTicket())
-                      .withJobBatchNumber(currentBatchNumber)
-                      .withJobResultsSoFar(resultsSoFar.addAndGet(uris.size()));
-                    batch = batch.withItems(uris.toArray(new String[uris.size()]));
-                    logger.trace("batch size={}, jobBatchNumber={}, jobResultsSoFar={}", uris.size(),
-                      batch.getJobBatchNumber(), batch.getJobResultsSoFar());
-                    for (QueryBatchListener listener : urisReadyListeners) {
-                      try {
-                        listener.processEvent(batch);
-                      } catch (Throwable e) {
-                        logger.error("Exception thrown by an onUrisReady listener", e);
-                      }
-                    }
-                  }
-                };
-                threadPool.execute(processBatch);
-              }
-            } catch (Throwable t) {
-              QueryBatchImpl batch = new QueryBatchImpl()
-                .withItems(new String[0])
-                .withClient(clientList.get().get(0))
-                .withBatcher(batcher)
-                .withTimestamp(Calendar.getInstance())
-                .withJobResultsSoFar(0);
-              for ( QueryFailureListener listener : failureListeners ) {
-                try {
-                  listener.processFailure(new QueryBatchException(batch, t));
-                } catch (Throwable e) {
-                  logger.error("Exception thrown by an onQueryFailure listener", e);
-                }
-              }
-              logger.warn("Error iterating to queue uris: {}", t.toString());
-            }
-          }
-        } catch (Throwable t) {
-          for ( QueryFailureListener listener : failureListeners ) {
-            try {
-              QueryBatchImpl batch = new QueryBatchImpl()
-                .withItems(new String[0])
-                .withClient(clientList.get().get(0))
-                .withBatcher(batcher)
-                .withTimestamp(Calendar.getInstance())
-                .withJobResultsSoFar(0);
-              listener.processFailure(new QueryBatchException(batch, t));
-            } catch (Throwable e) {
-              logger.error("Exception thrown by an onQueryFailure listener", e);
-            }
-          }
-          logger.warn("Error iterating to queue uris", t.toString());
-        }
-        runJobCompletionListeners.start();
-        threadPool.shutdown();
-      }
-    };
-    threadPool.execute(queueUris);
+    threadPool.execute(new IteratorTask(this));
   }
 
   public void stop() {

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

@@ -209,6 +209,22 @@ public void testRawValueQuery() throws Exception {
     runQueryBatcher(moveMgr.newQueryBatcher(query), query, matchesByForest, 17, 99);
   }
 
+  @Test
+  public void testIterator() throws Exception {
+    Map<String, String[]> matchesByForest = new HashMap<>();
+    matchesByForest.put("java-unittest-1", new String[] {uri1, uri3, uri4});
+    matchesByForest.put("java-unittest-2", new String[] {uri5});
+    matchesByForest.put("java-unittest-3", new String[] {uri2});
+    String[] uris = new String[] {uri1, uri2, uri3, uri4, uri5};
+    List<String> uriList = Arrays.asList(uris);
+    runQueryBatcher(moveMgr.newQueryBatcher(uriList.iterator()), null, matchesByForest, 1, 1, false);
+    runQueryBatcher(moveMgr.newQueryBatcher(uriList.iterator()), null, matchesByForest, 2, 2, false);
+    runQueryBatcher(moveMgr.newQueryBatcher(uriList.iterator()), null, matchesByForest, 2, 3, false);
+    runQueryBatcher(moveMgr.newQueryBatcher(uriList.iterator()), null, matchesByForest, 2, 10, false);
+    runQueryBatcher(moveMgr.newQueryBatcher(uriList.iterator()), null, matchesByForest, 10, 1, false);
+    runQueryBatcher(moveMgr.newQueryBatcher(uriList.iterator()), null, matchesByForest, 18, 33, false);
+  }
+
   @Test
   public void testRawCombinedQuery() throws Exception {
     StringHandle structuredQuery = new StringHandle(
@@ -233,9 +249,7 @@ public void testRawCombinedQuery() throws Exception {
   }
 
   public void runQueryBatcher(QueryBatcher queryBatcher, QueryDefinition query, Map<String,String[]> matchesByForest,
-      int batchSize, int threadCount)
-    throws Exception
-  {
+        int batchSize, int threadCount, boolean queryBatcherChecks) throws Exception {
     String queryBatcherJobId = "QueryBatcherJobId";
     String queryBatcherJobName = "QueryBatcherJobName";
     int numExpected = 0;
@@ -258,13 +272,15 @@ public void runQueryBatcher(QueryBatcher queryBatcher, QueryDefinition query, Ma
         batch -> {
           successfulBatchCount.incrementAndGet();
           totalResults.addAndGet(batch.getItems().length);
-          String forestName = batch.getForest().getForestName();
-          // atomically gets the set unless it's missing in which case it creates it
-          Set<String> matches = results.computeIfAbsent(forestName, k->ConcurrentHashMap.<String>newKeySet());
-          for ( String uri : batch.getItems() ) {
-            matches.add(uri);
+          if(queryBatcherChecks) {
+            String forestName = batch.getForest().getForestName();
+            // atomically gets the set unless it's missing in which case it creates it
+            Set<String> matches = results.computeIfAbsent(forestName, k->ConcurrentHashMap.<String>newKeySet());
+            for ( String uri : batch.getItems() ) {
+              matches.add(uri);
+            }
+            batchDatabaseName.set(batch.getForest().getDatabaseName());
           }
-          batchDatabaseName.set(batch.getForest().getDatabaseName());
           batchTicket.set(batch.getJobTicket());
           batchTimestamp.set(batch.getTimestamp());
         }
@@ -276,8 +292,8 @@ public void runQueryBatcher(QueryBatcher queryBatcher, QueryDefinition query, Ma
           failures.append("ERROR:[" + throwable + "]\n");
         }
       )
-        .withJobId(queryBatcherJobId)
-        .withJobName(queryBatcherJobName);
+      .withJobId(queryBatcherJobId)
+      .withJobName(queryBatcherJobName);
 
     assertEquals(batchSize, queryBatcher.getBatchSize());
     assertEquals(threadCount, queryBatcher.getThreadCount());
@@ -309,8 +325,6 @@ public void runQueryBatcher(QueryBatcher queryBatcher, QueryDefinition query, Ma
       fail(failures.toString());
     }
 
-    assertEquals("java-unittest", batchDatabaseName.get());
-
     // make sure we got the right number of results
     assertEquals(numExpected, totalResults.get());
 
@@ -327,27 +341,34 @@ public void runQueryBatcher(QueryBatcher queryBatcher, QueryDefinition query, Ma
     assertEquals("Job Report has incorrect failure events counts", failureBatchCount.get(), report.getFailureEventsCount());
     //assertEquals("Job Report has incorrect job completion information", true, report.isJobComplete());
 
-    // make sure we get the same number of results via search for the same query
-    SearchHandle searchResults = client.newQueryManager().search(query, new SearchHandle());
-    assertEquals(numExpected, searchResults.getTotalResults());
-
-    // if there are only the three expected forests, make sure we got the expected results per forest
-    if ( queryBatcher.getForestConfig().listForests().length == 3 ) {
-      for ( String forest : matchesByForest.keySet() ) {
-        String[] expected = matchesByForest.get(forest);
-        for ( String uri : expected ) {
-          if ( results.get(forest) == null || ! results.get(forest).contains(uri) ) {
-            for ( String resultsForest : results.keySet() ) {
-              logger.error("Results found for forest {}: {}, expected {}", resultsForest, results.get(resultsForest),
-                Arrays.asList(matchesByForest.get(resultsForest)));
+    if(queryBatcherChecks) {
+      assertEquals("java-unittest", batchDatabaseName.get());
+      // make sure we get the same number of results via search for the same query
+      SearchHandle searchResults = client.newQueryManager().search(query, new SearchHandle());
+      assertEquals(numExpected, searchResults.getTotalResults());
+      // if there are only the three expected forests, make sure we got the expected results per forest
+      if ( queryBatcher.getForestConfig().listForests().length == 3 ) {
+        for ( String forest : matchesByForest.keySet() ) {
+          String[] expected = matchesByForest.get(forest);
+          for ( String uri : expected ) {
+            if ( results.get(forest) == null || ! results.get(forest).contains(uri) ) {
+              for ( String resultsForest : results.keySet() ) {
+                logger.error("Results found for forest {}: {}, expected {}", resultsForest, results.get(resultsForest),
+                    Arrays.asList(matchesByForest.get(resultsForest)));
+              }
+              fail("Missing uri=[" + uri + "] from forest=[" + forest + "]");
             }
-            fail("Missing uri=[" + uri + "] from forest=[" + forest + "]");
           }
         }
       }
     }
   }
 
+  public void runQueryBatcher(QueryBatcher queryBatcher, QueryDefinition query, Map<String,String[]> matchesByForest,
+      int batchSize, int threadCount) throws Exception {
+    runQueryBatcher(queryBatcher, query, matchesByForest, batchSize, threadCount, true);
+  }
+
   @Test
   public void testMatchOneAndThrowException() {
     StructuredQueryDefinition query = new StructuredQueryBuilder().document(uri1);