handle duplicate events

flume-mongodb-sink/pom.xml

@@ -67,6 +67,12 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.testcontainers</groupId>
+      <artifactId>mongodb</artifactId>
+      <scope>test</scope>
+    </dependency>
+
     <dependency>
       <groupId>org.apache.logging.log4j</groupId>
       <artifactId>log4j-api</artifactId>
@@ -78,5 +84,34 @@
       <scope>test</scope>
     </dependency>
 
+    <dependency>
+      <groupId>org.apache.logging.log4j</groupId>
+      <artifactId>log4j-slf4j2-impl</artifactId>
+      <scope>test</scope>
+    </dependency>
+
   </dependencies>
+
+  <build>
+    <plugins>
+      <!--
+        ~ MongoDbSinkIT starts its own MongoDB container via Testcontainers and
+        ~ skips itself (via org.junit.Assume) when Docker is not available, so
+        ~ the failsafe plugin can run unconditionally on every platform/CI job.
+        -->
+      <plugin>
+        <groupId>org.apache.maven.plugins</groupId>
+        <artifactId>maven-failsafe-plugin</artifactId>
+        <executions>
+          <execution>
+            <goals>
+              <goal>integration-test</goal>
+              <goal>verify</goal>
+            </goals>
+          </execution>
+        </executions>
+      </plugin>
+    </plugins>
+  </build>
+
 </project>

flume-mongodb-sink/src/main/java/org/apache/flume/sink/mongodb/DefaultMongoDbWriter.java

@@ -16,10 +16,17 @@
  */
 package org.apache.flume.sink.mongodb;
 
+import com.mongodb.DuplicateKeyException;
+import com.mongodb.ErrorCategory;
+import com.mongodb.MongoBulkWriteException;
+import com.mongodb.MongoWriteException;
 import com.mongodb.WriteConcern;
+import com.mongodb.bulk.BulkWriteError;
 import com.mongodb.client.MongoCollection;
 import com.mongodb.client.MongoDatabase;
 import java.util.List;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 import org.bson.Document;
 
 /**
@@ -28,6 +35,8 @@
  */
 public class DefaultMongoDbWriter implements MongoDbWriter {
 
+    private static final Logger logger = LogManager.getLogger(DefaultMongoDbWriter.class);
+
     private final MongoDatabase mongoDatabase;
     private final WriteConcern writeConcern;
 
@@ -37,12 +46,89 @@ public DefaultMongoDbWriter(MongoDatabase mongoDatabase, WriteConcern writeConce
     }
 
     @Override
-    public void write(String collectionName, List<Document> documents) {
+    public MongoDbWriteResult write(String collectionName, List<Document> documents) {
         MongoCollection<Document> collection = mongoDatabase.getCollection(collectionName);
         if (writeConcern != null) {
             collection = collection.withWriteConcern(writeConcern);
         }
-        collection.insertMany(documents);
+
+        // First try to insert the whole batch in a single operation, which is
+        // far more efficient than one insert per document. Only fall back to
+        // inserting one document at a time if the batch insert fails because
+        // of duplicate keys.
+        try {
+            collection.insertMany(documents);
+            return new MongoDbWriteResult(documents.size(), 0);
+        } catch (MongoBulkWriteException ex) {
+            if (isDuplicateKeyOnly(ex)) {
+                // insertMany() is ordered by default, so it stops at the first
+                // failing document; everything before that point was already
+                // successfully persisted. Skip those already-written documents
+                // before retrying the remainder one at a time, otherwise they
+                // would be re-attempted and incorrectly counted as duplicates.
+                int alreadyInserted = ex.getWriteResult().getInsertedCount();
+                logger.warn(
+                        "Duplicate key(s) while batch inserting into collection {}, "
+                                + "retrying remaining documents one at a time: {}",
+                        collectionName,
+                        ex.getMessage());
+                MongoDbWriteResult retryResult = writeOneAtATime(
+                        collection, collectionName, documents.subList(alreadyInserted, documents.size()));
+                return new MongoDbWriteResult(
+                        alreadyInserted + retryResult.getInsertedCount(), retryResult.getDuplicateCount());
+            }
+            throw ex;
+        }
+    }
+
+    /**
+     * Returns {@code true} if every error reported by the bulk write failure
+     * is a duplicate key error.
+     */
+    private boolean isDuplicateKeyOnly(MongoBulkWriteException ex) {
+        List<BulkWriteError> errors = ex.getWriteErrors();
+        if (errors.isEmpty()) {
+            return false;
+        }
+        for (BulkWriteError error : errors) {
+            if (ErrorCategory.fromErrorCode(error.getCode()) != ErrorCategory.DUPLICATE_KEY) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    /**
+     * Inserts documents one at a time so that a duplicate key on any single
+     * document does not prevent the rest of the batch from being inserted.
+     */
+    private MongoDbWriteResult writeOneAtATime(
+            MongoCollection<Document> collection, String collectionName, List<Document> documents) {
+        long insertedCount = 0;
+        long duplicateCount = 0;
+        for (Document document : documents) {
+            try {
+                collection.insertOne(document);
+                insertedCount++;
+            } catch (DuplicateKeyException ex) {
+                logger.warn(
+                        "Duplicate key while inserting into collection {}, skipping event: {}",
+                        collectionName,
+                        ex.getMessage());
+                duplicateCount++;
+            } catch (MongoWriteException ex) {
+                if (ex.getError().getCategory() == ErrorCategory.DUPLICATE_KEY) {
+                    logger.warn(
+                            "Duplicate key while inserting into collection {}, skipping event: {}",
+                            collectionName,
+                            ex.getMessage());
+                    duplicateCount++;
+                } else {
+                    throw ex;
+                }
+            }
+        }
+        return new MongoDbWriteResult(insertedCount, duplicateCount);
     }
 
     @Override

flume-mongodb-sink/src/main/java/org/apache/flume/sink/mongodb/MongoDbSink.java

@@ -36,7 +36,6 @@
 import org.apache.flume.conf.BatchSizeSupported;
 import org.apache.flume.conf.Configurable;
 import org.apache.flume.conf.ConfigurationException;
-import org.apache.flume.instrumentation.SinkCounter;
 import org.apache.flume.sink.AbstractSink;
 import org.apache.logging.log4j.LogManager;
 import org.apache.logging.log4j.Logger;
@@ -86,7 +85,7 @@ public class MongoDbSink extends AbstractSink implements Configurable, BatchSize
 
     private MongoClient mongoClient;
     private MongoDbWriter writer;
-    private SinkCounter counter;
+    private MongoDbSinkCounter counter;
 
     // For testing
     public String getDatabaseName() {
@@ -97,6 +96,11 @@ public String getDefaultCollection() {
         return defaultCollection;
     }
 
+    /** For testing: number of events skipped due to duplicate keys. */
+    public long getDuplicateEventCount() {
+        return counter.getDuplicateEventCount();
+    }
+
     @Override
     public long getBatchSize() {
         return batchSize;
@@ -141,12 +145,19 @@ public Status process() throws EventDeliveryException {
                         .add(document);
             }
 
+            long insertedEvents = 0;
+            long duplicateEvents = 0;
             for (Map.Entry<String, List<Document>> entry : documentsByCollection.entrySet()) {
-                writer.write(entry.getKey(), entry.getValue());
+                MongoDbWriteResult writeResult = writer.write(entry.getKey(), entry.getValue());
+                insertedEvents += writeResult.getInsertedCount();
+                duplicateEvents += writeResult.getDuplicateCount();
             }
 
-            if (processedEvents > 0) {
-                counter.addToEventDrainSuccessCount(processedEvents);
+            if (insertedEvents > 0) {
+                counter.addToEventDrainSuccessCount(insertedEvents);
+            }
+            if (duplicateEvents > 0) {
+                counter.addToDuplicateEventCount(duplicateEvents);
             }
 
             transaction.commit();
@@ -284,7 +295,7 @@ public void configure(Context context) {
         }
 
         if (counter == null) {
-            counter = new SinkCounter(getName());
+            counter = new MongoDbSinkCounter(getName());
         }
     }
 

flume-mongodb-sink/src/main/java/org/apache/flume/sink/mongodb/MongoDbSinkCounter.java

@@ -0,0 +1,50 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flume.sink.mongodb;
+
+import org.apache.flume.instrumentation.SinkCounter;
+
+/**
+ * {@link SinkCounter} extension that additionally tracks the number of
+ * events skipped because they duplicated a document already present in
+ * MongoDB (i.e. resulted in a {@code com.mongodb.DuplicateKeyException}).
+ * These events are not counted towards {@code eventDrainSuccessCount} since
+ * they were not actually inserted, but they should also not be treated as a
+ * batch failure.
+ */
+public class MongoDbSinkCounter extends SinkCounter {
+
+    private static final String COUNTER_DUPLICATE_EVENT = "sink.event.duplicate";
+
+    private static final String[] ATTRIBUTES = {COUNTER_DUPLICATE_EVENT};
+
+    public MongoDbSinkCounter(String name) {
+        super(name, ATTRIBUTES);
+    }
+
+    public long getDuplicateEventCount() {
+        return get(COUNTER_DUPLICATE_EVENT);
+    }
+
+    public long incrementDuplicateEventCount() {
+        return increment(COUNTER_DUPLICATE_EVENT);
+    }
+
+    public long addToDuplicateEventCount(long delta) {
+        return addAndGet(COUNTER_DUPLICATE_EVENT, delta);
+    }
+}

flume-mongodb-sink/src/main/java/org/apache/flume/sink/mongodb/MongoDbWriteResult.java

@@ -0,0 +1,45 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to you under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.apache.flume.sink.mongodb;
+
+/**
+ * Outcome of writing a batch of documents to a single MongoDB collection.
+ */
+public final class MongoDbWriteResult {
+
+    private final long insertedCount;
+    private final long duplicateCount;
+
+    public MongoDbWriteResult(long insertedCount, long duplicateCount) {
+        this.insertedCount = insertedCount;
+        this.duplicateCount = duplicateCount;
+    }
+
+    /** Number of documents that were successfully inserted. */
+    public long getInsertedCount() {
+        return insertedCount;
+    }
+
+    /**
+     * Number of documents that were skipped because they duplicated a
+     * document that already existed in the collection (i.e. triggered a
+     * {@code com.mongodb.DuplicateKeyException}).
+     */
+    public long getDuplicateCount() {
+        return duplicateCount;
+    }
+}

flume-mongodb-sink/src/main/java/org/apache/flume/sink/mongodb/MongoDbWriter.java

@@ -28,12 +28,18 @@
 public interface MongoDbWriter {
 
     /**
-     * Writes the given documents to the named collection.
+     * Writes the given documents to the named collection. Documents that
+     * fail to insert because they duplicate an existing document (i.e.
+     * trigger a {@code com.mongodb.DuplicateKeyException}) are skipped
+     * rather than causing the whole batch to fail; they are reported via
+     * {@link MongoDbWriteResult#getDuplicateCount()}.
      *
      * @param collectionName the target collection name
      * @param documents      the documents to insert, in order
+     * @return the number of documents inserted and the number skipped as
+     *         duplicates
      */
-    void write(String collectionName, List<Document> documents);
+    MongoDbWriteResult write(String collectionName, List<Document> documents);
 
     /**
      * Releases any resources (e.g. the underlying MongoDB client) held by