#44 Figured out how to use Spring's support for waiting for tasks to finish

Author: Rob Rudin

.gitignore

@@ -7,3 +7,4 @@
 /.project
 .idea
 *.iml
+out

build.gradle

@@ -17,10 +17,10 @@ dependencies {
   compile 'com.marklogic:java-client-api:3.0.5'
   compile 'com.marklogic:marklogic-xcc:8.0.5'
   compile 'org.jdom:jdom2:2.0.5'
-  compile 'org.springframework:spring-context:4.1.5.RELEASE'
+  compile 'org.springframework:spring-context:4.3.5.RELEASE'
 
   testCompile 'junit:junit:4+'
-  testCompile 'org.springframework:spring-test:4.1.5.RELEASE'
+  testCompile 'org.springframework:spring-test:4.3.5.RELEASE'
 
   // Used for testing loading modules from the classpath
   testRuntime files("lib/modules.jar")

gradle.properties

@@ -1,3 +1,3 @@
 group=com.marklogic
 javadocsDir=../gh-pages-marklogic-java/javadocs
-version=2.11.0
+version=2.12.DEV

src/main/java/com/marklogic/client/batch/BatchWriterSupport.java

@@ -4,6 +4,7 @@
 import org.springframework.core.task.AsyncTaskExecutor;
 import org.springframework.core.task.SyncTaskExecutor;
 import org.springframework.core.task.TaskExecutor;
+import org.springframework.scheduling.concurrent.ExecutorConfigurationSupport;
 import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
 
 import java.util.ArrayList;
@@ -13,27 +14,44 @@
 
 /**
  * Support class for BatchWriter implementations that uses Spring's TaskExecutor interface for parallelizing writes to
- * MarkLogic.
+ * MarkLogic. Allows for setting a TaskExecutor instance, and if one is not set, a default one will be created based
+ * on the threadCount attribute. That attribute is ignored if a TaskExecutor is set.
  */
 public abstract class BatchWriterSupport extends LoggingObject implements BatchWriter {
 
 	private TaskExecutor taskExecutor;
 	private int threadCount = 16;
 
-	/**
-	 * Seems necessary to keep track of each Future instance so that we can properly wait for each one to finish.
-	 * Spring's TaskExecutor library doesn't seem to provide a better way of doing this.
-	 */
-	private List<Future<?>> futures = new ArrayList<>();
-
 	@Override
 	public void initialize() {
+		if (taskExecutor == null) {
+			initializeDefaultTaskExecutor();
+		}
+	}
+
+	@Override
+	public void waitForCompletion() {
+		if (taskExecutor instanceof ExecutorConfigurationSupport) {
+			if (logger.isInfoEnabled()) {
+				logger.info("Calling shutdown on thread pool");
+			}
+			((ExecutorConfigurationSupport) taskExecutor).shutdown();
+			if (logger.isInfoEnabled()) {
+				logger.info("Thread pool finished shutdown");
+			}
+		}
+	}
+
+	protected void initializeDefaultTaskExecutor() {
 		if (threadCount > 1) {
 			if (logger.isInfoEnabled()) {
 				logger.info("Initializing thread pool with a count of " + threadCount);
 			}
 			ThreadPoolTaskExecutor tpte = new ThreadPoolTaskExecutor();
 			tpte.setCorePoolSize(threadCount);
+			// By default, wait for tasks to finish, and wait up to an hour
+			tpte.setWaitForTasksToCompleteOnShutdown(true);
+			tpte.setAwaitTerminationSeconds(60 * 60);
 			tpte.afterPropertiesSet();
 			this.taskExecutor = tpte;
 		} else {
@@ -44,33 +62,8 @@ public void initialize() {
 		}
 	}
 
-	@Override
-	public void waitForCompletion() {
-		int size = futures.size();
-		if (logger.isDebugEnabled()) {
-			logger.debug("Waiting for threads to finish document processing; futures count: " + size);
-		}
-
-		for (int i = 0; i < size; i++) {
-			Future<?> f = futures.get(i);
-			if (f.isDone()) {
-				continue;
-			}
-			try {
-				// Wait up to 1 hour for a write to ML to finish (should never happen)
-				f.get(1, TimeUnit.HOURS);
-			} catch (Exception ex) {
-				logger.warn("Unable to wait for last set of documents to be processed: " + ex.getMessage(), ex);
-			}
-		}
-	}
-
-	protected void execute(Runnable runnable) {
-		if (taskExecutor instanceof AsyncTaskExecutor) {
-			futures.add(((AsyncTaskExecutor) taskExecutor).submit(runnable));
-		} else {
-			taskExecutor.execute(runnable);
-		}
+	protected TaskExecutor getTaskExecutor() {
+		return taskExecutor;
 	}
 
 	public void setTaskExecutor(TaskExecutor taskExecutor) {

src/main/java/com/marklogic/client/batch/RestBatchWriter.java

@@ -7,10 +7,16 @@
 
 import java.util.List;
 
+/**
+ * REST API-based implementation, using the Java Client API. By default, this will call release() on each of the
+ * DatabaseClient objects that are passed in. Be sure to disable this if you want to keep using those DatabaseClient
+ * objects.
+ */
 public class RestBatchWriter extends BatchWriterSupport {
 
 	private List<DatabaseClient> databaseClients;
 	private int clientIndex = 0;
+	private boolean releaseDatabaseClients = true;
 
 	public RestBatchWriter(List<DatabaseClient> databaseClients) {
 		this.databaseClients = databaseClients;
@@ -24,7 +30,7 @@ public void write(final List<? extends DocumentWriteOperation> items) {
 		final DatabaseClient client = databaseClients.get(clientIndex);
 		clientIndex++;
 
-		execute(new Runnable() {
+		getTaskExecutor().execute(new Runnable() {
 			@Override
 			public void run() {
 				GenericDocumentManager mgr = client.newDocumentManager();
@@ -43,4 +49,21 @@ public void run() {
 			}
 		});
 	}
+
+	@Override
+	public void waitForCompletion() {
+		super.waitForCompletion();
+
+		if (databaseClients != null && releaseDatabaseClients) {
+			logger.info("Releasing DatabaseClient instances...");
+			for (DatabaseClient client : databaseClients) {
+				client.release();
+			}
+			logger.info("Finished releasing DatabaseClient instances");
+		}
+	}
+
+	public void setReleaseDatabaseClients(boolean releaseDatabaseClients) {
+		this.releaseDatabaseClients = releaseDatabaseClients;
+	}
 }

src/main/java/com/marklogic/client/batch/XccBatchWriter.java

@@ -35,7 +35,7 @@ public void write(final List<? extends DocumentWriteOperation> items) {
 		final ContentSource contentSource = contentSources.get(contentSourceIndex);
 		contentSourceIndex++;
 
-		execute(new Runnable() {
+		getTaskExecutor().execute(new Runnable() {
 			@Override
 			public void run() {
 				Session session = contentSource.newSession();