Copying data - ensure the queue blocks if full

KAFKA-85
mongodb · Feb 14, 2020 · 7e6bf97 · 7e6bf97
1 parent 5e68296
commit 7e6bf97
Show file tree

Hide file tree

Showing 2 changed files with 56 additions and 3 deletions.
diff --git a/src/main/java/com/mongodb/kafka/connect/source/MongoCopyDataManager.java b/src/main/java/com/mongodb/kafka/connect/source/MongoCopyDataManager.java
@@ -120,13 +120,21 @@ private void copyDataFrom(final MongoNamespace namespace) {
             mongoClient.getDatabase(namespace.getDatabaseName())
                     .getCollection(namespace.getCollectionName(), BsonDocument.class)
                     .aggregate(createPipeline(namespace))
-                    .forEach((Consumer<? super BsonDocument>) queue::add);
+                    .forEach((Consumer<? super BsonDocument>) this::putToQueue);
             namespacesToCopy.decrementAndGet();
         } catch (Exception e) {
             errorException = e;
         }
     }
 
+    private void putToQueue(final BsonDocument bsonDocument) {
+        try {
+            queue.put(bsonDocument);
+        } catch (InterruptedException e) {
+            throw new RuntimeException(e);
+        }
+    }
+
     private List<Bson> createPipeline(final MongoNamespace namespace) {
         List<Bson> pipeline = new ArrayList<>();
         pipeline.add(BsonDocument.parse("{$replaceRoot: {newRoot: {"

diff --git a/src/test/java/com/mongodb/kafka/connect/source/MongoCopyDataManagerTest.java b/src/test/java/com/mongodb/kafka/connect/source/MongoCopyDataManagerTest.java
@@ -17,10 +17,12 @@
 
 
 import static com.mongodb.kafka.connect.source.MongoSourceConfig.COLLECTION_CONFIG;
+import static com.mongodb.kafka.connect.source.MongoSourceConfig.COPY_EXISTING_QUEUE_SIZE_CONFIG;
 import static com.mongodb.kafka.connect.source.SourceTestHelper.TEST_COLLECTION;
 import static com.mongodb.kafka.connect.source.SourceTestHelper.TEST_DATABASE;
 import static com.mongodb.kafka.connect.source.SourceTestHelper.createConfigMap;
 import static com.mongodb.kafka.connect.source.SourceTestHelper.createSourceConfig;
+import static java.lang.String.format;
 import static java.util.Arrays.asList;
 import static junit.framework.TestCase.assertTrue;
 import static org.junit.jupiter.api.Assertions.assertEquals;
@@ -31,11 +33,14 @@
 import static org.mockito.Mockito.when;
 
 import java.util.ArrayList;
+import java.util.Arrays;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
 
 import org.junit.jupiter.api.DisplayName;
 import org.junit.jupiter.api.Test;
@@ -111,6 +116,42 @@ void testReturnsTheExpectedCollectionResults() {
         assertEquals(expected, results);
     }
 
+    @Test
+    @DisplayName("test blocks adding docs to the queue")
+    void testBlocksAddingResultsToTheQueue() {
+        List<BsonDocument> docs = IntStream.range(0, 10).mapToObj(i ->
+                BsonDocument.parse(format("{'_id': {'_id': %s, 'copy': true}}", i))
+        ).collect(Collectors.toList());
+
+        when(mongoClient.getDatabase(TEST_DATABASE)).thenReturn(mongoDatabase);
+        when(mongoDatabase.getCollection(TEST_COLLECTION, BsonDocument.class)).thenReturn(mongoCollection);
+        when(mongoCollection.aggregate(anyList())).thenReturn(aggregateIterable);
+        doCallRealMethod().when(aggregateIterable).forEach(any(Consumer.class));
+        when(aggregateIterable.iterator()).thenReturn(cursor);
+
+        Boolean[] hasNextResponses = new Boolean[docs.size()];
+        Arrays.fill(hasNextResponses, true);
+        hasNextResponses[hasNextResponses.length - 1] = false;
+
+        when(cursor.hasNext()).thenReturn(true, hasNextResponses);
+        when(cursor.next()).thenReturn(docs.get(0), docs.subList(1, docs.size()).toArray(new BsonDocument[docs.size() - 1]));
+
+        List<Optional<BsonDocument>> results;
+        try (MongoCopyDataManager copyExistingDataManager =
+                     new MongoCopyDataManager(createSourceConfig(COPY_EXISTING_QUEUE_SIZE_CONFIG, "1"), mongoClient)) {
+            sleep();
+            results = IntStream.range(0, 11).mapToObj(i -> {
+                        sleep(200);
+                        return copyExistingDataManager.poll();
+                    }
+            ).collect(Collectors.toList());
+        }
+
+        List<Optional<BsonDocument>> expected = docs.stream().map(Optional::of).collect(Collectors.toList());
+        expected.add(Optional.empty());
+        assertEquals(expected, results);
+    }
+
     @Test
     @DisplayName("test returns the expected database results")
     void testReturnsTheExpectedDatabaseResults() {
@@ -231,11 +272,15 @@ void testReturnsTheExpectedClientResults() {
         assertEquals(results.get(results.size() - 1), Optional.empty());
     }
 
-    private void sleep() {
+    private void sleep(final int millis) {
         try {
-            Thread.sleep(500);
+            Thread.sleep(millis);
         } catch (InterruptedException e) {
             // ignore
         }
     }
+
+    private void sleep() {
+       sleep(500);
+    }
 }