[Kafka][CosmosSourceConnector]FixDuplicateRecordProcessing (Azure#47308)

xinlian12 · annie-mac · Copilot · web-flow · commit a52ae3a097bb · 2025-11-18T09:23:23.000-08:00
* fix duplicate record processing
---------

Co-authored-by: annie-mac &lt;xinlian@microsoft.com&gt;
Co-authored-by: Copilot &lt;175728472+Copilot@users.noreply.github.com&gt;
diff --git a/sdk/cosmos/azure-cosmos-kafka-connect/CHANGELOG.md b/sdk/cosmos/azure-cosmos-kafka-connect/CHANGELOG.md
@@ -7,6 +7,7 @@
 #### Breaking Changes
 
 #### Bugs Fixed
+* Fixed duplicate processing in `CosmosSourceConnector` when task rebalance(pause/resume) happens - See [PR 47308](https://github.com/Azure/azure-sdk-for-java/pull/47308)
 
 #### Other Changes
 
diff --git a/sdk/cosmos/azure-cosmos-kafka-connect/src/main/java/com/azure/cosmos/kafka/connect/implementation/source/CosmosSourceTask.java b/sdk/cosmos/azure-cosmos-kafka-connect/src/main/java/com/azure/cosmos/kafka/connect/implementation/source/CosmosSourceTask.java
@@ -47,6 +47,8 @@ public class CosmosSourceTask extends SourceTask {
     private CosmosSourceTaskConfig taskConfig;
     private CosmosClientCacheItem cosmosClientItem;
     private CosmosClientCacheItem throughputControlCosmosClientItem;
+    private MetadataKafkaStorageManager kafkaOffsetStorageReader;
+
     private final Queue<ITaskUnit> taskUnitsQueue = new LinkedList<>();
 
     private long lastLogTimeMs = System.currentTimeMillis();
@@ -76,7 +78,13 @@ public void start(Map<String, String> map) {
                 this.taskConfig.getFeedRangeTaskUnits().size(),
                 this.taskConfig.getTaskId());
 
-            this.taskUnitsQueue.addAll(this.taskConfig.getFeedRangeTaskUnits());
+            this.kafkaOffsetStorageReader = new MetadataKafkaStorageManager(this.context.offsetStorageReader());
+
+            for (FeedRangeTaskUnit taskUnit : this.taskConfig.getFeedRangeTaskUnits()) {
+                this.applyContinuationStateFromOffsetIfExists(taskUnit);
+                this.taskUnitsQueue.add(taskUnit);
+            }
+
             LOGGER.info("Creating the cosmos client");
 
             this.cosmosClientItem =
@@ -92,6 +100,55 @@ public void start(Map<String, String> map) {
         }
     }
 
+    private void applyContinuationStateFromOffsetIfExists(FeedRangeTaskUnit feedRangeTaskUnit) {
+        // populate the changeFeedState if offset exists for the feed range
+
+        // In task level, we will try to find the EXACT offset matching for the feed range.
+        // The task config used here is only a snapshot of the config assigned; when there is a task level rebalance,
+        // the new task will keep using the snapshot which could contain stale feed range or change feed state
+
+        KafkaCosmosChangeFeedState changeFeedState =
+            this.getContinuationStateFromOffset(
+                feedRangeTaskUnit.getDatabaseName(),
+                feedRangeTaskUnit.getContainerRid(),
+                feedRangeTaskUnit.getFeedRange());
+
+        if (changeFeedState != null) {
+            LOGGER.debug(
+                "Find continuation state from offset {} for feed range {} in task {}. ",
+                changeFeedState,
+                feedRangeTaskUnit.getFeedRange(),
+                this.taskConfig.getTaskId());
+            feedRangeTaskUnit.setContinuationState(changeFeedState);
+        } else {
+            LOGGER.debug(
+                "Can not find continuation state from offset for feed range {} in task {}. ",
+                feedRangeTaskUnit.getFeedRange(),
+                this.taskConfig.getTaskId());
+        }
+    }
+
+    private KafkaCosmosChangeFeedState getContinuationStateFromOffset(
+        String databaseName,
+        String containerRid,
+        FeedRange feedRange) {
+
+        FeedRangeContinuationTopicOffset feedRangeContinuationTopicOffset =
+            this.kafkaOffsetStorageReader.getFeedRangeContinuationOffset(
+                databaseName,
+                containerRid,
+                feedRange);
+
+        if (feedRangeContinuationTopicOffset == null) {
+            return null;
+        }
+
+        return new KafkaCosmosChangeFeedState(
+            feedRangeContinuationTopicOffset.getResponseContinuation(),
+            feedRange,
+            feedRangeContinuationTopicOffset.getItemLsn());
+    }
+
     private CosmosClientCacheItem getThroughputControlCosmosClientItem() {
         if (this.taskConfig.getThroughputControlConfig().isThroughputControlEnabled()
             && this.taskConfig.getThroughputControlConfig().getThroughputControlAccountConfig() != null) {
@@ -147,7 +204,7 @@ public List<SourceRecord> poll() {
                     return loggingContext;
                 });
 
-                logFeedRangeCounts();
+                logFeedRangeCounts(false);
             }
 
             return results;
@@ -160,10 +217,10 @@ public List<SourceRecord> poll() {
         }
     }
 
-    private void logFeedRangeCounts() {
+    private void logFeedRangeCounts(boolean flushLogs) {
         long currentTimeInMs = System.currentTimeMillis();
         long durationInMs = currentTimeInMs - lastLogTimeMs;
-        if (durationInMs >= CosmosSourceTaskConfig.LOG_INTERVAL_MS) {
+        if (flushLogs || durationInMs >= CosmosSourceTaskConfig.LOG_INTERVAL_MS) {
             // Log accumulated counts for all feed ranges
             for (Map.Entry<String, FeedRangeLoggingContext> entry : feedRangeCounts.entrySet()) {
                 LOGGER.info(
@@ -370,6 +427,8 @@ private Mono<Boolean> handleFeedRangeGone(FeedRangeTaskUnit feedRangeTaskUnit) {
                                 pkRange,
                                 getChildRangeChangeFeedState(feedRangeTaskUnit.getContinuationState(), pkRange),
                                 feedRangeTaskUnit.getTopic());
+                        applyContinuationStateFromOffsetIfExists(childTaskUnit);
+
                         this.taskUnitsQueue.add(childTaskUnit);
                     }
 
@@ -476,6 +535,7 @@ private void cleanup() {
     @Override
     public void stop() {
         LOGGER.info("Stopping CosmosSourceTask");
+        this.logFeedRangeCounts(true);
         this.cleanup();
     }
 
diff --git a/sdk/cosmos/azure-cosmos-kafka-connect/src/test/java/com/azure/cosmos/kafka/connect/CosmosSourceConnectorITest.java b/sdk/cosmos/azure-cosmos-kafka-connect/src/test/java/com/azure/cosmos/kafka/connect/CosmosSourceConnectorITest.java
@@ -34,6 +34,7 @@
 import reactor.core.publisher.Mono;
 
 import java.time.Duration;
+import java.time.Instant;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashMap;
@@ -201,20 +202,16 @@ public void readFromSingleContainer(boolean useMasterKey, CosmosMetadataStorageT
             List<ConsumerRecord<String, JsonNode>> metadataRecords = new ArrayList<>();
             List<ConsumerRecord<String, JsonNode>> itemRecords = new ArrayList<>();
             int expectedMetadataRecordsCount = metadataStorageType == CosmosMetadataStorageType.COSMOS ? 0 : 2;
-            int expectedItemRecords = createdItems.size();
+            int expectedItemRecordsCount = createdItems.size();
 
-            Unreliables.retryUntilTrue(30, TimeUnit.SECONDS, () -> {
-                kafkaConsumer.poll(Duration.ofMillis(1000))
-                    .iterator()
-                    .forEachRemaining(consumerRecord -> {
-                        if (consumerRecord.topic().equals(topicName)) {
-                            itemRecords.add(consumerRecord);
-                        } else if (consumerRecord.topic().equals(sourceConfig.getMetadataConfig().getStorageName())) {
-                            metadataRecords.add(consumerRecord);
-                        }
-                    });
-                return metadataRecords.size() >= expectedMetadataRecordsCount && itemRecords.size() >= expectedItemRecords;
-            });
+            pollChangesForSingleTopic(
+                kafkaConsumer,
+                topicName,
+                sourceConfig.getMetadataConfig().getStorageName(),
+                itemRecords,
+                metadataRecords,
+                expectedItemRecordsCount,
+                expectedMetadataRecordsCount);
 
             assertThat(metadataRecords.size()).isEqualTo(expectedMetadataRecordsCount);
             if (metadataStorageType == CosmosMetadataStorageType.KAFKA) {
@@ -270,15 +267,7 @@ public void readFromSingleContainer(boolean useMasterKey, CosmosMetadataStorageT
                 assertThat(feedRangesMetadataTopicOffsetOffset.getFeedRanges().size()).isEqualTo(1);
             }
 
-            // validate the item records
-            assertThat(itemRecords.size()).isEqualTo(createdItems.size());
-            List<String> receivedItems =
-                itemRecords.stream().map(consumerRecord -> {
-                    JsonNode jsonNode = consumerRecord.value();
-                    return jsonNode.get("payload").get("id").asText();
-                }).collect(Collectors.toList());
-
-            assertThat(receivedItems.containsAll(createdItems)).isTrue();
+            validateFeedRangeItemRecords(itemRecords, createdItems);
 
         } finally {
             if (client != null) {
@@ -680,4 +669,150 @@ public void readFromAllContainer(boolean useMasterKey, CosmosMetadataStorageType
             }
         }
     }
+
+    @Test(groups = { "kafka-integration" }, timeOut = 2 * TIMEOUT)
+    public void readFromSingleContainer_pause_and_resume() {
+        logger.info("Pause and resume connector for single container ");
+        String topicName = singlePartitionContainerName + "-" + UUID.randomUUID();
+
+        Map<String, String> sourceConnectorConfig = new HashMap<>();
+        sourceConnectorConfig.put("connector.class", "com.azure.cosmos.kafka.connect.CosmosSourceConnector");
+        sourceConnectorConfig.put("azure.cosmos.account.endpoint", KafkaCosmosTestConfigurations.HOST);
+        sourceConnectorConfig.put("azure.cosmos.application.name", "Test");
+        sourceConnectorConfig.put("azure.cosmos.source.database.name", databaseName);
+        sourceConnectorConfig.put("azure.cosmos.source.containers.includeAll", "false");
+        sourceConnectorConfig.put("azure.cosmos.source.containers.includedList", singlePartitionContainerName);
+        sourceConnectorConfig.put("azure.cosmos.source.containers.topicMap", topicName + "#" + singlePartitionContainerName);
+        sourceConnectorConfig.put("azure.cosmos.account.key", KafkaCosmosTestConfigurations.MASTER_KEY);
+
+        // Create topic ahead of time
+        kafkaCosmosConnectContainer.createTopic(topicName, 1);
+
+        CosmosSourceConfig sourceConfig = new CosmosSourceConfig(sourceConnectorConfig);
+        CosmosAsyncContainer container = client.getDatabase(databaseName).getContainer(singlePartitionContainerName);
+
+        String connectorName = "simpleTest-" + UUID.randomUUID();
+
+        try {
+            // create few items in the container
+            logger.info("creating items in container {}", singlePartitionContainerName);
+            List<String> createdItems = new ArrayList<>();
+            for (int i = 0; i < 10; i++) {
+                TestItem testItem = TestItem.createNewItem();
+                container.createItem(testItem).block();
+                createdItems.add(testItem.getId());
+            }
+
+            kafkaCosmosConnectContainer.registerConnector(connectorName, sourceConnectorConfig);
+
+            logger.info("Getting consumer and subscribe to topic {}", singlePartitionContainerName);
+
+            Properties consumerProperties = kafkaCosmosConnectContainer.getConsumerProperties();
+            consumerProperties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class.getName());
+            consumerProperties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, JsonDeserializer.class.getName());
+            KafkaConsumer<String, JsonNode> kafkaConsumer = new KafkaConsumer<>(consumerProperties);
+
+            kafkaConsumer.subscribe(
+                Arrays.asList(
+                    topicName,
+                    sourceConfig.getMetadataConfig().getStorageName()));
+
+            List<ConsumerRecord<String, JsonNode>> metadataRecords = new ArrayList<>();
+            List<ConsumerRecord<String, JsonNode>> itemRecords = new ArrayList<>();
+            int expectedMetadataRecordsCount = 2;
+            int expectedItemRecordsCount = createdItems.size();
+
+            pollChangesForSingleTopic(
+                kafkaConsumer,
+                topicName,
+                sourceConfig.getMetadataConfig().getStorageName(),
+                itemRecords,
+                metadataRecords,
+                expectedItemRecordsCount,
+                expectedMetadataRecordsCount);
+
+            assertThat(metadataRecords.size()).isEqualTo(expectedMetadataRecordsCount);
+            validateFeedRangeItemRecords(itemRecords, createdItems);
+
+            // now pause the connector
+            kafkaCosmosConnectContainer.pauseConnector(connectorName);
+
+            // create few items
+            createdItems.clear();
+            metadataRecords.clear();
+            itemRecords.clear();
+
+            for (int i = 0; i < 5; i++) {
+                TestItem testItem = TestItem.createNewItem();
+                container.createItem(testItem).block();
+                createdItems.add(testItem.getId());
+            }
+
+            // resume the connector
+            kafkaCosmosConnectContainer.resumeConnector(connectorName);
+            // poll again, poll a little bit longer to make sure no duplicate records are being returned
+            Instant startPollTime = Instant.now();
+            while (Duration.between(startPollTime, Instant.now()).toMillis() < 60 * 1000 ) {
+                kafkaConsumer.poll(Duration.ofMillis(1000))
+                    .iterator()
+                    .forEachRemaining(consumerRecord -> {
+                        if (consumerRecord.topic().equals(topicName)) {
+                            itemRecords.add(consumerRecord);
+                        } else if (consumerRecord.topic().equals(sourceConfig.getMetadataConfig().getStorageName())) {
+                            metadataRecords.add(consumerRecord);
+                        }
+                    });
+            }
+
+            assertThat(metadataRecords.size()).isEqualTo(expectedMetadataRecordsCount);
+            validateFeedRangeItemRecords(itemRecords, createdItems);
+        } finally {
+            if (client != null) {
+                logger.info("cleaning container {}", singlePartitionContainerName);
+                cleanUpContainer(client, databaseName, singlePartitionContainerName);
+            }
+
+            // IMPORTANT: remove the connector after use
+            if (kafkaCosmosConnectContainer != null) {
+                kafkaCosmosConnectContainer.deleteConnector(connectorName);
+            }
+        }
+    }
+
+    private void pollChangesForSingleTopic(
+        KafkaConsumer<String, JsonNode> kafkaConsumer,
+        String topicName,
+        String storageName,
+        List<ConsumerRecord<String, JsonNode>> itemRecords,
+        List<ConsumerRecord<String, JsonNode>> metadataRecords,
+        int expectedItemRecords,
+        int expectedMetadataRecordsCount) {
+
+        Unreliables.retryUntilTrue(30, TimeUnit.SECONDS, () -> {
+            kafkaConsumer.poll(Duration.ofMillis(1000))
+                .iterator()
+                .forEachRemaining(consumerRecord -> {
+                    if (consumerRecord.topic().equals(topicName)) {
+                        itemRecords.add(consumerRecord);
+                    } else if (consumerRecord.topic().equals(storageName)) {
+                        metadataRecords.add(consumerRecord);
+                    }
+                });
+            return metadataRecords.size() >= expectedMetadataRecordsCount && itemRecords.size() >= expectedItemRecords;
+        });
+    }
+
+    private void validateFeedRangeItemRecords(
+        List<ConsumerRecord<String, JsonNode>> itemRecords,
+        List<String> expectedItems) {
+        // validate the item records
+        assertThat(itemRecords.size()).isEqualTo(expectedItems.size());
+        List<String> receivedItems =
+            itemRecords.stream().map(consumerRecord -> {
+                JsonNode jsonNode = consumerRecord.value();
+                return jsonNode.get("payload").get("id").asText();
+            }).collect(Collectors.toList());
+
+        assertThat(receivedItems.containsAll(expectedItems)).isTrue();
+    }
 }
diff --git a/sdk/cosmos/azure-cosmos-kafka-connect/src/test/java/com/azure/cosmos/kafka/connect/KafkaCosmosConnectContainer.java b/sdk/cosmos/azure-cosmos-kafka-connect/src/test/java/com/azure/cosmos/kafka/connect/KafkaCosmosConnectContainer.java
@@ -172,7 +172,7 @@ public void deleteConnector(String name) {
         KafkaConnectClient kafkaConnectClient = new KafkaConnectClient(new Configuration(getTarget()));
         try {
             kafkaConnectClient.deleteConnector(name);
-            logger.info("Deleting container {} succeeded.", name);
+            logger.info("Deleting connector {} succeeded.", name);
         } catch (Exception exception) {
             if (exception instanceof ResourceNotFoundException) {
                 logger.info("Connector {} not found");
@@ -182,6 +182,34 @@ public void deleteConnector(String name) {
         }
     }
 
+    public void pauseConnector(String name) {
+        KafkaConnectClient kafkaConnectClient = new KafkaConnectClient(new Configuration(getTarget()));
+        try {
+            kafkaConnectClient.pauseConnector(name);
+            logger.info("Pause connector {} succeeded.", name);
+        } catch (Exception exception) {
+            if (exception instanceof ResourceNotFoundException) {
+                logger.info("Connector {} not found");
+            }
+
+            logger.warn("Failed to pause connector {}", name);
+        }
+    }
+
+    public void resumeConnector(String name) {
+        KafkaConnectClient kafkaConnectClient = new KafkaConnectClient(new Configuration(getTarget()));
+        try {
+            kafkaConnectClient.resumeConnector(name);
+            logger.info("Resume connector {} succeeded.", name);
+        } catch (Exception exception) {
+            if (exception instanceof ResourceNotFoundException) {
+                logger.info("Connector {} not found");
+            }
+
+            logger.warn("Failed to Resume connector {}", name);
+        }
+    }
+
     public ConnectorStatus getConnectorStatus(String name) {
         KafkaConnectClient kafkaConnectClient = new KafkaConnectClient(new Configuration(getTarget()));
         return kafkaConnectClient.getConnectorStatus(name);
diff --git a/sdk/cosmos/azure-cosmos-kafka-connect/src/test/java/com/azure/cosmos/kafka/connect/implementation/source/CosmosSourceTaskTest.java b/sdk/cosmos/azure-cosmos-kafka-connect/src/test/java/com/azure/cosmos/kafka/connect/implementation/source/CosmosSourceTaskTest.java
