Merge pull request #165 from marklogic/feature/12357-metrics

MLE-12357 Improved logging of success/failure counts

Author: rjrudin

src/main/java/com/marklogic/spark/writer/CommitMessage.java

@@ -4,32 +4,34 @@
 
 public class CommitMessage implements WriterCommitMessage {
 
-    private final String action;
-    private final int docCount;
+    private final int successItemCount;
+    private final int failedItemCount;
     private final int partitionId;
     private final long taskId;
     private final long epochId;
 
-    public CommitMessage(String action, int docCount, int partitionId, long taskId, long epochId) {
-        this.action = action;
-        this.docCount = docCount;
+    public CommitMessage(int successItemCount, int failedItemCount, int partitionId, long taskId, long epochId) {
+        this.successItemCount = successItemCount;
+        this.failedItemCount = failedItemCount;
         this.partitionId = partitionId;
         this.taskId = taskId;
         this.epochId = epochId;
     }
 
-    public String getAction() {
-        return action;
+    public int getSuccessItemCount() {
+        return successItemCount;
     }
 
-    public int getDocCount() {
-        return docCount;
+    public int getFailedItemCount() {
+        return failedItemCount;
     }
 
     @Override
     public String toString() {
         return epochId != 0L ?
-            String.format("[partitionId: %d; taskId: %d; epochId: %d]; docCount: %d", partitionId, taskId, epochId, docCount) :
-            String.format("[partitionId: %d; taskId: %d]; docCount: %d", partitionId, taskId, docCount);
+            String.format("[partitionId: %d; taskId: %d; epochId: %d]; docCount: %d; failedItemCount: %d",
+                partitionId, taskId, epochId, successItemCount, failedItemCount) :
+            String.format("[partitionId: %d; taskId: %d]; docCount: %d; failedItemCount: %d",
+                partitionId, taskId, successItemCount, failedItemCount);
     }
 }

src/main/java/com/marklogic/spark/writer/MarkLogicWrite.java

@@ -29,13 +29,18 @@
 import org.slf4j.LoggerFactory;
 
 import java.util.Arrays;
+import java.util.function.Consumer;
 
-class MarkLogicWrite implements BatchWrite, StreamingWrite {
+public class MarkLogicWrite implements BatchWrite, StreamingWrite {
 
     private static final Logger logger = LoggerFactory.getLogger("com.marklogic.spark");
 
     private WriteContext writeContext;
 
+    // Used solely for testing. Will never be populated in a real world scenario.
+    public static Consumer<Integer> successCountConsumer;
+    public static Consumer<Integer> failureCountConsumer;
+
     MarkLogicWrite(WriteContext writeContext) {
         this.writeContext = writeContext;
     }
@@ -57,13 +62,24 @@ public void commit(WriterCommitMessage[] messages) {
             if (logger.isDebugEnabled()) {
                 logger.debug("Commit messages received: {}", Arrays.asList(messages));
             }
-            final String action = ((CommitMessage) messages[0]).getAction();
+            int successCount = 0;
+            int failureCount = 0;
+            for (WriterCommitMessage message : messages) {
+                CommitMessage msg = (CommitMessage)message;
+                successCount += msg.getSuccessItemCount();
+                failureCount += msg.getFailedItemCount();
+            }
+            if (successCountConsumer != null) {
+                successCountConsumer.accept(successCount);
+            }
+            if (failureCountConsumer != null) {
+                failureCountConsumer.accept(failureCount);
+            }
             if (Util.MAIN_LOGGER.isInfoEnabled()) {
-                int count = 0;
-                for (WriterCommitMessage message : messages) {
-                    count += ((CommitMessage) message).getDocCount();
-                }
-                Util.MAIN_LOGGER.info("{} {} documents", action, count);
+                Util.MAIN_LOGGER.info("Success count: {}", successCount);
+            }
+            if (failureCount > 0) {
+                Util.MAIN_LOGGER.error("Failure count: {}", failureCount);
             }
         }
     }

src/main/java/com/marklogic/spark/writer/WriteBatcherDataWriter.java

@@ -27,6 +27,7 @@
 import org.slf4j.LoggerFactory;
 
 import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicReference;
 import java.util.function.Function;
 
@@ -51,7 +52,9 @@ class WriteBatcherDataWriter implements DataWriter<InternalRow> {
 
     private final Function<InternalRow, DocBuilder.DocumentInputs> rowToDocumentFunction;
 
-    private int docCount;
+    // Updated as batches are processed.
+    private AtomicInteger successItemCount = new AtomicInteger(0);
+    private AtomicInteger failedItemCount = new AtomicInteger(0);
 
     WriteBatcherDataWriter(WriteContext writeContext, int partitionId, long taskId, long epochId) {
         this.writeContext = writeContext;
@@ -63,12 +66,7 @@ class WriteBatcherDataWriter implements DataWriter<InternalRow> {
         this.databaseClient = writeContext.connectToMarkLogic();
         this.dataMovementManager = this.databaseClient.newDataMovementManager();
         this.writeBatcher = writeContext.newWriteBatcher(this.dataMovementManager);
-        if (writeContext.isAbortOnFailure()) {
-            this.writeBatcher.onBatchFailure((batch, failure) ->
-                // Logging not needed here, as WriteBatcherImpl already logs this at the warning level.
-                this.writeFailure.compareAndSet(null, failure)
-            );
-        }
+        addBatchListeners();
         this.dataMovementManager.startJob(this.writeBatcher);
 
         if (writeContext.isUsingFileSchema()) {
@@ -85,16 +83,15 @@ public void write(InternalRow row) throws IOException {
         throwWriteFailureIfExists();
         DocBuilder.DocumentInputs inputs = rowToDocumentFunction.apply(row);
         this.writeBatcher.add(this.docBuilder.build(inputs));
-        this.docCount++;
     }
 
     @Override
     public WriterCommitMessage commit() throws IOException {
-        CommitMessage message = new CommitMessage("Wrote", docCount, partitionId, taskId, epochId);
+        this.writeBatcher.flushAndWait();
+        CommitMessage message = new CommitMessage(successItemCount.get(), failedItemCount.get(), partitionId, taskId, epochId);
         if (logger.isDebugEnabled()) {
             logger.debug("Committing {}", message);
         }
-        this.writeBatcher.flushAndWait();
         throwWriteFailureIfExists();
         return message;
     }
@@ -113,6 +110,20 @@ public void close() {
         stopJobAndRelease();
     }
 
+    private void addBatchListeners() {
+        this.writeBatcher.onBatchSuccess(batch -> this.successItemCount.getAndAdd(batch.getItems().length));
+        if (writeContext.isAbortOnFailure()) {
+            // Logging not needed here, as WriteBatcherImpl already logs this at the warning level.
+            this.writeBatcher.onBatchFailure((batch, failure) ->
+                this.writeFailure.compareAndSet(null, failure)
+            );
+        } else {
+            this.writeBatcher.onBatchFailure((batch, failure) ->
+                this.failedItemCount.getAndAdd(batch.getItems().length)
+            );
+        }
+    }
+
     private synchronized void throwWriteFailureIfExists() throws IOException {
         if (writeFailure.get() != null) {
             // Only including the message seems sufficient here, as Spark is logging the stacktrace. And the user

src/main/java/com/marklogic/spark/writer/customcode/CustomCodeWriter.java

@@ -42,7 +42,10 @@ class CustomCodeWriter implements DataWriter<InternalRow> {
     private final int partitionId;
     private final long taskId;
     private final long epochId;
-    private int itemCount;
+
+    // Updated after each call to MarkLogic.
+    private int successItemCount;
+    private int failedItemCount;
 
     CustomCodeWriter(CustomCodeContext customCodeContext, int partitionId, long taskId, long epochId) {
         this.customCodeContext = customCodeContext;
@@ -68,7 +71,6 @@ public void write(InternalRow row) {
         String rowValue = customCodeContext.isCustomSchema() ?
             convertRowToJSONString(row) :
             row.getString(0);
-        itemCount++;
 
         this.currentBatch.add(rowValue);
 
@@ -80,7 +82,7 @@ public void write(InternalRow row) {
     @Override
     public WriterCommitMessage commit() {
         flush();
-        CommitMessage message = new CommitMessage("Processed", itemCount, partitionId, taskId, epochId);
+        CommitMessage message = new CommitMessage(successItemCount, failedItemCount, partitionId, taskId, epochId);
         if (logger.isDebugEnabled()) {
             logger.debug("Committing {}", message);
         }
@@ -114,13 +116,14 @@ private void flush() {
         if (logger.isDebugEnabled()) {
             logger.debug("Calling custom code in MarkLogic");
         }
+        final int itemCount = currentBatch.size();
         ServerEvaluationCall call = customCodeContext.buildCall(
             this.databaseClient,
             new CustomCodeContext.CallOptions(Options.WRITE_INVOKE, Options.WRITE_JAVASCRIPT, Options.WRITE_XQUERY)
         );
         call.addVariable(determineExternalVariableName(), makeVariableValue());
         currentBatch.clear();
-        executeCall(call);
+        executeCall(call, itemCount);
     }
 
     private String determineExternalVariableName() {
@@ -165,13 +168,15 @@ private String convertRowToJSONString(InternalRow row) {
         return jsonObjectWriter.toString();
     }
 
-    private void executeCall(ServerEvaluationCall call) {
+    private void executeCall(ServerEvaluationCall call, int itemCount) {
         try {
             call.evalAs(String.class);
+            this.successItemCount += itemCount;
         } catch (RuntimeException ex) {
             if (customCodeContext.isAbortOnFailure()) {
                 throw ex;
             }
+            this.failedItemCount += itemCount;
             Util.MAIN_LOGGER.error(String.format("Unable to process row; cause: %s", ex.getMessage()));
         }
     }

src/test/java/com/marklogic/spark/writer/WriteRowsTest.java

@@ -24,6 +24,7 @@
 import org.junit.jupiter.api.Test;
 
 import java.io.IOException;
+import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.junit.jupiter.api.Assertions.*;
 
@@ -177,6 +178,11 @@ void invalidPermissionsConfig() {
 
     @Test
     void dontAbortOnFailure() {
+        AtomicInteger successCount = new AtomicInteger();
+        AtomicInteger failureCount = new AtomicInteger();
+        MarkLogicWrite.successCountConsumer = count -> successCount.set(count);
+        MarkLogicWrite.failureCountConsumer = count -> failureCount.set(count);
+
         newWriterWithDefaultConfig("temporal-data-with-invalid-rows.csv", 1)
             .option(Options.WRITE_TEMPORAL_COLLECTION, TEMPORAL_COLLECTION)
             // Force each row in the CSV to be written in its own batch, ensuring that the one row that should succeed
@@ -186,6 +192,8 @@ void dontAbortOnFailure() {
             .save();
 
         assertCollectionSize("9 of the batches should have failed, with the 10th batch succeeding", COLLECTION, 1);
+        assertEquals(9, failureCount.get());
+        assertEquals(1, successCount.get());
     }
 
     private void verifyFailureIsDueToLackOfPermission(SparkException ex) {

src/test/java/com/marklogic/spark/writer/customcode/ProcessWithCustomCodeTest.java

@@ -4,10 +4,12 @@
 import com.marklogic.junit5.XmlNode;
 import com.marklogic.spark.Options;
 import com.marklogic.spark.writer.AbstractWriteTest;
+import com.marklogic.spark.writer.MarkLogicWrite;
 import org.apache.spark.SparkException;
 import org.apache.spark.sql.SaveMode;
 import org.junit.jupiter.api.Test;
 
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.stream.Stream;
 
 import static org.junit.jupiter.api.Assertions.*;
@@ -121,15 +123,22 @@ void abortOnFailure() {
                 "Unexpected error message: " + ex.getMessage());
     }
 
-    @SuppressWarnings("java:S2699") // The absence of an assertion is fine for this test.
     @Test
     void dontAbortOnFailure() {
+        AtomicInteger successCount = new AtomicInteger();
+        AtomicInteger failureCount = new AtomicInteger();
+        MarkLogicWrite.successCountConsumer = count -> successCount.set(count);
+        MarkLogicWrite.failureCountConsumer = count -> failureCount.set(count);
+
         // The lack of an error here indicates that the job did not abort. The connector is expected to have logged
         // each error instead.
         newWriterWithDefaultConfig("three-uris.csv", 2)
             .option(Options.WRITE_JAVASCRIPT, "var URI; throw Error('Boom!');")
             .option(Options.WRITE_ABORT_ON_FAILURE, "false")
             .save();
+
+        assertEquals(3, failureCount.get());
+        assertEquals(0, successCount.get());
     }
 
     private void verifyThreeJsonDocumentsWereWritten() {