#51 Added WriteListener for catching failures

Author: Rob Rudin

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

@@ -1,11 +1,15 @@
 package com.marklogic.client.batch;
 
+import com.marklogic.client.document.DocumentWriteOperation;
 import com.marklogic.client.helper.LoggingObject;
+import org.springframework.core.task.AsyncListenableTaskExecutor;
 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 org.springframework.util.concurrent.ListenableFuture;
+import org.springframework.util.concurrent.ListenableFutureCallback;
 
 import java.util.ArrayList;
 import java.util.List;
@@ -21,6 +25,7 @@ public abstract class BatchWriterSupport extends LoggingObject implements BatchW
 
 	private TaskExecutor taskExecutor;
 	private int threadCount = 16;
+	private WriteListener writeListener;
 
 	@Override
 	public void initialize() {
@@ -64,6 +69,31 @@ protected void initializeDefaultTaskExecutor() {
 		}
 	}
 
+	/**
+	 * Will use the WriteListener if the TaskExecutor is an instance of AsyncListenableTaskExecutor. The WriteListener
+	 * will then be used to listen for failures.
+	 *
+	 * @param runnable
+	 * @param items
+	 */
+	protected void executeRunnable(Runnable runnable, final List<? extends DocumentWriteOperation> items) {
+		if (writeListener != null && taskExecutor instanceof AsyncListenableTaskExecutor) {
+			AsyncListenableTaskExecutor asyncListenableTaskExecutor = (AsyncListenableTaskExecutor)taskExecutor;
+			ListenableFuture<?> future = asyncListenableTaskExecutor.submitListenable(runnable);
+			future.addCallback(new ListenableFutureCallback<Object>() {
+				@Override
+				public void onFailure(Throwable ex) {
+					writeListener.onWriteFailure(ex, items);
+				}
+				@Override
+				public void onSuccess(Object result) {
+				}
+			});
+		} else {
+			taskExecutor.execute(runnable);
+		}
+	}
+
 	protected TaskExecutor getTaskExecutor() {
 		return taskExecutor;
 	}
@@ -75,4 +105,12 @@ public void setTaskExecutor(TaskExecutor taskExecutor) {
 	public void setThreadCount(int threadCount) {
 		this.threadCount = threadCount;
 	}
+
+	protected WriteListener getWriteListener() {
+		return writeListener;
+	}
+
+	public void setWriteListener(WriteListener writeListener) {
+		this.writeListener = writeListener;
+	}
 }

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

@@ -6,6 +6,7 @@
 import com.marklogic.client.document.DocumentWriteSet;
 import com.marklogic.client.document.ServerTransform;
 
+import java.util.Arrays;
 import java.util.List;
 
 /**
@@ -20,19 +21,32 @@ public class RestBatchWriter extends BatchWriterSupport {
 	private boolean releaseDatabaseClients = true;
 	private ServerTransform serverTransform;
 
+	public RestBatchWriter(DatabaseClient databaseClient) {
+		this(Arrays.asList(databaseClient));
+	}
+
 	public RestBatchWriter(List<DatabaseClient> databaseClients) {
 		this.databaseClients = databaseClients;
 	}
 
 	@Override
-	public void write(final List<? extends DocumentWriteOperation> items) {
+	public void write(List<? extends DocumentWriteOperation> items) {
+		DatabaseClient client = determineDatabaseClientToUse();
+		Runnable runnable = buildRunnable(client, items);
+		executeRunnable(runnable, items);
+	}
+
+	protected DatabaseClient determineDatabaseClientToUse() {
 		if (clientIndex >= databaseClients.size()) {
 			clientIndex = 0;
 		}
-		final DatabaseClient client = databaseClients.get(clientIndex);
+		DatabaseClient client = databaseClients.get(clientIndex);
 		clientIndex++;
+		return client;
+	}
 
-		getTaskExecutor().execute(new Runnable() {
+	protected Runnable buildRunnable(final DatabaseClient client, final List<? extends DocumentWriteOperation> items) {
+		return new Runnable() {
 			@Override
 			public void run() {
 				DocumentManager<?, ?> mgr = buildDocumentManager(client);
@@ -53,7 +67,7 @@ public void run() {
 					logger.info("Wrote " + count + " documents to MarkLogic");
 				}
 			}
-		});
+		};
 	}
 
 	/**

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

@@ -0,0 +1,13 @@
+package com.marklogic.client.batch;
+
+import com.marklogic.client.document.DocumentWriteOperation;
+
+import java.util.List;
+
+/**
+ * Callback interface for when a list of DocumentWriteOperation instances cannot be written to MarkLogic.
+ */
+public interface WriteListener {
+
+	void onWriteFailure(Throwable ex, List<? extends DocumentWriteOperation> items);
+}

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

@@ -8,7 +8,6 @@
 import com.marklogic.xcc.Session;
 import com.marklogic.xcc.exceptions.RequestException;
 
-import java.util.ArrayList;
 import java.util.List;
 
 /**
@@ -28,14 +27,22 @@ public XccBatchWriter(List<ContentSource> contentSources) {
 
 	@Override
 	public void write(final List<? extends DocumentWriteOperation> items) {
+		ContentSource contentSource = determineContentSourceToUse();
+		Runnable runnable = buildRunnable(contentSource, items);
+		executeRunnable(runnable, items);
+	}
+
+	protected ContentSource determineContentSourceToUse() {
 		if (contentSourceIndex >= contentSources.size()) {
 			contentSourceIndex = 0;
 		}
-
-		final ContentSource contentSource = contentSources.get(contentSourceIndex);
+		ContentSource contentSource = contentSources.get(contentSourceIndex);
 		contentSourceIndex++;
+		return contentSource;
+	}
 
-		getTaskExecutor().execute(new Runnable() {
+	protected Runnable buildRunnable(final ContentSource contentSource, final List<? extends DocumentWriteOperation> items) {
+		return new Runnable() {
 			@Override
 			public void run() {
 				Session session = contentSource.newSession();
@@ -58,7 +65,7 @@ public void run() {
 					session.close();
 				}
 			}
-		});
+		};
 	}
 
 	public void setDocumentWriteOperationAdapter(DocumentWriteOperationAdapter documentWriteOperationAdapter) {

src/test/java/com/marklogic/client/batch/RestBatchWriterTest.java

@@ -0,0 +1,41 @@
+package com.marklogic.client.batch;
+
+import com.marklogic.client.AbstractIntegrationTest;
+import com.marklogic.client.FailedRequestException;
+import com.marklogic.client.document.DocumentWriteOperation;
+import com.marklogic.client.impl.DocumentWriteOperationImpl;
+import com.marklogic.client.io.StringHandle;
+import org.junit.Test;
+
+import java.util.Arrays;
+import java.util.List;
+
+public class RestBatchWriterTest extends AbstractIntegrationTest {
+
+	@Test
+	public void failureTest() {
+		RestBatchWriter writer = new RestBatchWriter(newClient("Documents"));
+		TestWriteListener testWriteListener = new TestWriteListener();
+		writer.setWriteListener(testWriteListener);
+
+		DocumentWriteOperation op = new DocumentWriteOperationImpl(DocumentWriteOperation.OperationType.DOCUMENT_WRITE,
+			"/test.xml", null, new StringHandle("<hello>world</hello>asdf"));
+
+		writer.initialize();
+		writer.write(Arrays.asList(op));
+		writer.waitForCompletion();
+
+		Throwable caughtError = testWriteListener.caughtError;
+		assertNotNull("An error should have been thrown due to the invalid XML", caughtError);
+		assertTrue(caughtError instanceof FailedRequestException);
+	}
+}
+
+class TestWriteListener implements WriteListener {
+	public Throwable caughtError;
+
+	@Override
+	public void onWriteFailure(Throwable ex, List<? extends DocumentWriteOperation> items) {
+		this.caughtError = ex;
+	}
+}