Incremental Co-operative Rebalancing Support for HDFS Connector

"TopicPartitionWriters" Map is being cleared on every close call, which is correct if we consider the stop the world scenario. But in Incremental mode there can be case that only few of the partitions have been revoked and task has some partitions with it to read from, In that scenario it will be throwing NullPointerException at the following part of the code.

As a consequence the add or the removal of an instance create a high pressure on all instances as they need to recover from last commited file for all the task of the cluster

Applying unmerged commit from original project : confluentinc#712

Author: HerveRiviere

src/main/java/io/confluent/connect/hdfs/DataWriter.java

@@ -452,15 +452,12 @@ public void open(Collection<TopicPartition> partitions) {
   }
 
   public void close() {
-    // Close any writers we have. We may get assigned the same partitions and end up duplicating
-    // some effort since we'll have to reprocess those messages. It may be possible to hold on to
-    // the TopicPartitionWriter and continue to use the temp file, but this can get significantly
-    // more complex due to potential failures and network partitions. For example, we may get
-    // this close, then miss a few generations of group membership, during which
-    // data may have continued to be processed and we'd have to restart from the recovery stage,
-    // make sure we apply the WAL, and only reuse the temp file if the starting offset is still
-    // valid. For now, we prefer the simpler solution that may result in a bit of wasted effort.
-    for (TopicPartitionWriter writer : topicPartitionWriters.values()) {
+    close(new HashSet<>(topicPartitionWriters.keySet()));
+  }
+
+  public void close(Collection<TopicPartition> partitions) {
+    for (TopicPartition partition: partitions) {
+      TopicPartitionWriter writer = topicPartitionWriters.get(partition);
       try {
         if (writer != null) {
           // In some failure modes, the writer might not have been created for all assignments
@@ -469,8 +466,8 @@ public void close() {
       } catch (ConnectException e) {
         log.warn("Unable to close writer for topic partition {}: ", writer.topicPartition(), e);
       }
+      topicPartitionWriters.remove(partition);
     }
-    topicPartitionWriters.clear();
   }
 
   public void stop() {

src/main/java/io/confluent/connect/hdfs/HdfsSinkTask.java

@@ -166,7 +166,7 @@ public void open(Collection<TopicPartition> partitions) {
   public void close(Collection<TopicPartition> partitions) {
     log.debug("Closing HDFS Sink Task {}", connectorNameAndTaskId);
     if (hdfsWriter != null) {
-      hdfsWriter.close();
+      hdfsWriter.close(partitions);
     }
   }
 

src/test/java/io/confluent/connect/hdfs/HdfsSinkTaskTest.java

@@ -28,6 +28,7 @@
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collection;
+import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
@@ -341,6 +342,43 @@ final int record = 12;
     }
   }
 
+  @Test
+  public void testPartialRevocation() throws Exception {
+    setUp();
+
+    Collection<TopicPartition> initialAssignment = new ArrayList<>();
+    initialAssignment.add(TOPIC_PARTITION);
+    initialAssignment.add(TOPIC_PARTITION2);
+    initialAssignment.add(TOPIC_PARTITION3);
+
+    Collection<TopicPartition> revokedPartitions = new ArrayList<>();
+    revokedPartitions.add(TOPIC_PARTITION3);
+
+    String key = "key";
+    Schema schema = createSchema();
+    Struct record = createRecord(schema);
+    Collection<SinkRecord> sinkRecords = Collections.singleton(
+            new SinkRecord(TOPIC_PARTITION.topic(), TOPIC_PARTITION.partition(),
+                    Schema.STRING_SCHEMA, key, schema, record, 0));
+
+    HdfsSinkTask task = new HdfsSinkTask();
+    task.initialize(context);
+    task.start(properties);
+
+    // Given 3 owned partitions
+    task.open(initialAssignment);
+
+    // When 1 partition revoked (partial revocation)
+    task.close(revokedPartitions);
+
+    try {
+      // Should continue processing messages from the 2 left partitions (should succeed)
+      task.put(sinkRecords);
+    } finally {
+      task.stop();
+    }
+  }
+
   private void createCommittedFiles() throws IOException {
     String topicsDir = this.topicsDir.get(TOPIC_PARTITION.topic());
     String file1 = FileUtils.committedFileName(url, topicsDir, DIRECTORY1, TOPIC_PARTITION, 0,