Refactor internals into Writer and Committer (#234)

kafka-connect-runtime/src/test/java/io/tabular/iceberg/connect/IntegrationCdcTest.java

@@ -120,7 +120,7 @@ private void runTest(String branch, boolean useSchema) {
             .config("value.converter.schemas.enable", useSchema)
             .config("iceberg.tables", String.format("%s.%s", TEST_DB, TEST_TABLE))
             .config("iceberg.tables.cdc-field", "op")
-            .config("iceberg.control.commit.interval-ms", 1000)
+            .config("iceberg.control.commit.interval-ms", 5000)
             .config("iceberg.control.commit.timeout-ms", Integer.MAX_VALUE)
             .config("iceberg.kafka.auto.offset.reset", "earliest");
 

kafka-connect-runtime/src/test/java/io/tabular/iceberg/connect/IntegrationDynamicTableTest.java

@@ -97,7 +97,7 @@ private void runTest(String branch, boolean useSchema) {
             .config("value.converter.schemas.enable", useSchema)
             .config("iceberg.tables.dynamic-enabled", true)
             .config("iceberg.tables.route-field", "payload")
-            .config("iceberg.control.commit.interval-ms", 1000)
+            .config("iceberg.control.commit.interval-ms", 5000)
             .config("iceberg.control.commit.timeout-ms", Integer.MAX_VALUE)
             .config("iceberg.kafka.auto.offset.reset", "earliest");
 

kafka-connect-runtime/src/test/java/io/tabular/iceberg/connect/IntegrationMultiTableTest.java

@@ -101,7 +101,7 @@ private void runTest(String branch, boolean useSchema) {
             .config("iceberg.tables.route-field", "type")
             .config(String.format("iceberg.table.%s.%s.route-regex", TEST_DB, TEST_TABLE1), "type1")
             .config(String.format("iceberg.table.%s.%s.route-regex", TEST_DB, TEST_TABLE2), "type2")
-            .config("iceberg.control.commit.interval-ms", 1000)
+            .config("iceberg.control.commit.interval-ms", 5000)
             .config("iceberg.control.commit.timeout-ms", Integer.MAX_VALUE)
             .config("iceberg.kafka.auto.offset.reset", "earliest");
 

kafka-connect-runtime/src/test/java/io/tabular/iceberg/connect/IntegrationTest.java

@@ -229,7 +229,7 @@ private void runTest(String branch, boolean useSchema, Map<String, String> extra
             .config("value.converter", "org.apache.kafka.connect.json.JsonConverter")
             .config("value.converter.schemas.enable", useSchema)
             .config("iceberg.tables", String.format("%s.%s", TEST_DB, TEST_TABLE))
-            .config("iceberg.control.commit.interval-ms", 1000)
+            .config("iceberg.control.commit.interval-ms", 5000)
             .config("iceberg.control.commit.timeout-ms", Integer.MAX_VALUE)
             .config("iceberg.kafka.auto.offset.reset", "earliest");
 

kafka-connect/src/main/java/io/tabular/iceberg/connect/IcebergSinkTask.java

@@ -18,27 +18,14 @@
  */
 package io.tabular.iceberg.connect;
 
-import io.tabular.iceberg.connect.channel.Coordinator;
-import io.tabular.iceberg.connect.channel.CoordinatorThread;
-import io.tabular.iceberg.connect.channel.KafkaClientFactory;
-import io.tabular.iceberg.connect.channel.KafkaUtils;
-import io.tabular.iceberg.connect.channel.NotRunningException;
-import io.tabular.iceberg.connect.channel.Worker;
-import io.tabular.iceberg.connect.data.IcebergWriterFactory;
+import io.tabular.iceberg.connect.channel.Task;
+import io.tabular.iceberg.connect.channel.TaskImpl;
 import io.tabular.iceberg.connect.data.Utilities;
 import java.util.Collection;
-import java.util.Comparator;
 import java.util.Map;
-import org.apache.iceberg.catalog.Catalog;
-import org.apache.iceberg.relocated.com.google.common.annotations.VisibleForTesting;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
-import org.apache.kafka.clients.admin.Admin;
-import org.apache.kafka.clients.admin.ConsumerGroupDescription;
-import org.apache.kafka.clients.admin.MemberDescription;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
-import org.apache.kafka.common.ConsumerGroupState;
 import org.apache.kafka.common.TopicPartition;
-import org.apache.kafka.connect.errors.ConnectException;
 import org.apache.kafka.connect.sink.SinkRecord;
 import org.apache.kafka.connect.sink.SinkTask;
 import org.slf4j.Logger;
@@ -49,21 +36,7 @@ public class IcebergSinkTask extends SinkTask {
   private static final Logger LOG = LoggerFactory.getLogger(IcebergSinkTask.class);
 
   private IcebergSinkConfig config;
-  private Catalog catalog;
-  private CoordinatorThread coordinatorThread;
-  private Worker worker;
-
-  static class TopicPartitionComparator implements Comparator<TopicPartition> {
-
-    @Override
-    public int compare(TopicPartition o1, TopicPartition o2) {
-      int result = o1.topic().compareTo(o2.topic());
-      if (result == 0) {
-        result = Integer.compare(o1.partition(), o2.partition());
-      }
-      return result;
-    }
-  }
+  private Task task;
 
   @Override
   public String version() {
@@ -79,43 +52,8 @@ public void start(Map<String, String> props) {
   public void open(Collection<TopicPartition> partitions) {
     // destroy any state if KC re-uses object
     clearObjectState();
-    catalog = Utilities.loadCatalog(config);
-    KafkaClientFactory clientFactory = new KafkaClientFactory(config.kafkaProps());
-
-    ConsumerGroupDescription groupDesc;
-    try (Admin admin = clientFactory.createAdmin()) {
-      groupDesc = KafkaUtils.consumerGroupDescription(config.connectGroupId(), admin);
-    }
 
-    if (groupDesc.state() == ConsumerGroupState.STABLE) {
-      Collection<MemberDescription> members = groupDesc.members();
-      if (isLeader(members, partitions)) {
-        LOG.info("Task elected leader, starting commit coordinator");
-        Coordinator coordinator = new Coordinator(catalog, config, members, clientFactory);
-        coordinatorThread = new CoordinatorThread(coordinator);
-        coordinatorThread.start();
-      }
-    }
-
-    LOG.info("Starting commit worker");
-    IcebergWriterFactory writerFactory = new IcebergWriterFactory(catalog, config);
-    worker = new Worker(config, clientFactory, writerFactory, context);
-    worker.syncCommitOffsets();
-    worker.start();
-  }
-
-  @VisibleForTesting
-  boolean isLeader(Collection<MemberDescription> members, Collection<TopicPartition> partitions) {
-    // there should only be one task assigned partition 0 of the first topic,
-    // so elect that one the leader
-    TopicPartition firstTopicPartition =
-        members.stream()
-            .flatMap(member -> member.assignment().topicPartitions().stream())
-            .min(new TopicPartitionComparator())
-            .orElseThrow(
-                () -> new ConnectException("No partitions assigned, cannot determine leader"));
-
-    return partitions.contains(firstTopicPartition);
+    task = new TaskImpl(context, config);
   }
 
   @Override
@@ -128,57 +66,21 @@ private void close() {
   }
 
   private void clearObjectState() {
-    if (worker != null) {
-      worker.stop();
-      worker = null;
-    }
-
-    if (coordinatorThread != null) {
-      coordinatorThread.terminate();
-      coordinatorThread = null;
-    }
-
-    if (catalog != null) {
-      if (catalog instanceof AutoCloseable) {
-        try {
-          ((AutoCloseable) catalog).close();
-        } catch (Exception e) {
-          LOG.warn("An error occurred closing catalog instance, ignoring...", e);
-        }
-      }
-      catalog = null;
-    }
+    Utilities.close(task);
+    task = null;
   }
 
   @Override
   public void put(Collection<SinkRecord> sinkRecords) {
-    if (sinkRecords != null && !sinkRecords.isEmpty() && worker != null) {
-      worker.save(sinkRecords);
-    }
-    processControlEvents();
-  }
-
-  @Override
-  public void flush(Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
-    processControlEvents();
-  }
-
-  private void processControlEvents() {
-    if (coordinatorThread != null && coordinatorThread.isTerminated()) {
-      throw new NotRunningException("Coordinator unexpectedly terminated");
-    }
-    if (worker != null) {
-      worker.process();
+    if (task != null) {
+      task.put(sinkRecords);
     }
   }
 
   @Override
   public Map<TopicPartition, OffsetAndMetadata> preCommit(
       Map<TopicPartition, OffsetAndMetadata> currentOffsets) {
-    if (worker == null) {
-      return ImmutableMap.of();
-    }
-    return worker.commitOffsets();
+    return ImmutableMap.of();
   }
 
   @Override

kafka-connect/src/main/java/io/tabular/iceberg/connect/channel/Channel.java

@@ -27,9 +27,11 @@
 import java.util.List;
 import java.util.Map;
 import java.util.UUID;
+import java.util.function.Function;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
 import org.apache.iceberg.relocated.com.google.common.collect.Maps;
+import org.apache.iceberg.util.Pair;
 import org.apache.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerGroupMetadata;
@@ -46,7 +48,6 @@ public abstract class Channel {
   private static final Logger LOG = LoggerFactory.getLogger(Channel.class);
 
   private final String controlTopic;
-  private final String controlGroupId;
   private final String groupId;
   private final Producer<String, byte[]> producer;
   private final Consumer<String, byte[]> consumer;
@@ -60,22 +61,26 @@ public Channel(
       IcebergSinkConfig config,
       KafkaClientFactory clientFactory) {
     this.controlTopic = config.controlTopic();
-    this.controlGroupId = config.controlGroupId();
     this.groupId = config.controlGroupId();
 
     String transactionalId = name + config.transactionalSuffix();
-    this.producer = clientFactory.createProducer(transactionalId);
+    Pair<UUID, Producer<String, byte[]>> pair = clientFactory.createProducer(transactionalId);
+    this.producer = pair.second();
     this.consumer = clientFactory.createConsumer(consumerGroupId);
+    consumer.subscribe(ImmutableList.of(controlTopic));
     this.admin = clientFactory.createAdmin();
 
-    this.producerId = UUID.randomUUID().toString();
+    this.producerId = pair.first().toString();
   }
 
   protected void send(Event event) {
-    send(ImmutableList.of(event), ImmutableMap.of());
+    send(ImmutableList.of(event), ImmutableMap.of(), null);
   }
 
-  protected void send(List<Event> events, Map<TopicPartition, Offset> sourceOffsets) {
+  protected void send(
+      List<Event> events,
+      Map<TopicPartition, Offset> sourceOffsets,
+      ConsumerGroupMetadata consumerGroupMetadata) {
     Map<TopicPartition, OffsetAndMetadata> offsetsToCommit = Maps.newHashMap();
     sourceOffsets.forEach((k, v) -> offsetsToCommit.put(k, new OffsetAndMetadata(v.offset())));
 
@@ -94,10 +99,9 @@ protected void send(List<Event> events, Map<TopicPartition, Offset> sourceOffset
       producer.beginTransaction();
       try {
         recordList.forEach(producer::send);
+        producer.flush();
         if (!sourceOffsets.isEmpty()) {
-          // TODO: this doesn't fence zombies
-          producer.sendOffsetsToTransaction(
-              offsetsToCommit, new ConsumerGroupMetadata(controlGroupId));
+          producer.sendOffsetsToTransaction(offsetsToCommit, consumerGroupMetadata);
         }
         producer.commitTransaction();
       } catch (Exception e) {
@@ -111,9 +115,7 @@ protected void send(List<Event> events, Map<TopicPartition, Offset> sourceOffset
     }
   }
 
-  protected abstract boolean receive(Envelope envelope);
-
-  protected void consumeAvailable(Duration pollDuration) {
+  protected void consumeAvailable(Duration pollDuration, Function<Envelope, Boolean> receiveFn) {
     ConsumerRecords<String, byte[]> records = consumer.poll(pollDuration);
     while (!records.isEmpty()) {
       records.forEach(
@@ -126,7 +128,7 @@ record -> {
 
             if (event.groupId().equals(groupId)) {
               LOG.debug("Received event of type: {}", event.type().name());
-              if (receive(new Envelope(event, record.partition(), record.offset()))) {
+              if (receiveFn.apply(new Envelope(event, record.partition(), record.offset()))) {
                 LOG.info("Handled event of type: {}", event.type().name());
               }
             }
@@ -152,13 +154,6 @@ protected Admin admin() {
     return admin;
   }
 
-  public void start() {
-    consumer.subscribe(ImmutableList.of(controlTopic));
-
-    // initial poll with longer duration so the consumer will initialize...
-    consumeAvailable(Duration.ofMillis(1000));
-  }
-
   public void stop() {
     LOG.info("Channel stopping");
     producer.close();

kafka-connect/src/main/java/io/tabular/iceberg/connect/channel/Committable.java

@@ -0,0 +1,47 @@
+/*
+ * 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 io.tabular.iceberg.connect.channel;
+
+import io.tabular.iceberg.connect.data.Offset;
+import io.tabular.iceberg.connect.data.WriterResult;
+import java.util.List;
+import java.util.Map;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.ImmutableMap;
+import org.apache.kafka.common.TopicPartition;
+
+class Committable {
+
+  private final ImmutableMap<TopicPartition, Offset> offsetsByTopicPartition;
+  private final ImmutableList<WriterResult> writerResults;
+
+  Committable(
+      Map<TopicPartition, Offset> offsetsByTopicPartition, List<WriterResult> writerResults) {
+    this.offsetsByTopicPartition = ImmutableMap.copyOf(offsetsByTopicPartition);
+    this.writerResults = ImmutableList.copyOf(writerResults);
+  }
+
+  public Map<TopicPartition, Offset> offsetsByTopicPartition() {
+    return offsetsByTopicPartition;
+  }
+
+  public List<WriterResult> writerResults() {
+    return writerResults;
+  }
+}

kafka-connect/src/main/java/io/tabular/iceberg/connect/channel/CommittableSupplier.java

@@ -0,0 +1,23 @@
+/*
+ * 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 io.tabular.iceberg.connect.channel;
+
+interface CommittableSupplier {
+  Committable committable();
+}

kafka-connect/src/main/java/io/tabular/iceberg/connect/channel/Committer.java

@@ -0,0 +1,23 @@
+/*
+ * 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 io.tabular.iceberg.connect.channel;
+
+interface Committer {
+  void commit(CommittableSupplier committableSupplier);
+}