Kafka Connect: make sink cleanup robust

kafka-connect/kafka-connect/src/main/java/org/apache/iceberg/connect/IcebergSinkTask.java

@@ -61,20 +61,23 @@ public void close(Collection<TopicPartition> partitions) {
   }
 
   private void close() {
-    if (committer != null) {
-      committer.close(List.of());
+    try {
+      if (committer != null) {
+        committer.close(List.of());
+      }
+    } finally {
       committer = 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);
+      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;
       }
-      catalog = null;
     }
   }
 

kafka-connect/kafka-connect/src/main/java/org/apache/iceberg/connect/channel/Channel.java

@@ -30,7 +30,6 @@
 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.kafka.clients.admin.Admin;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRecords;
 import org.apache.kafka.clients.consumer.OffsetAndMetadata;
@@ -50,7 +49,6 @@ abstract class Channel {
   private final Producer<String, byte[]> producer;
   private final Consumer<String, byte[]> consumer;
   private final SinkTaskContext context;
-  private final Admin admin;
   private final Map<Integer, Long> controlTopicOffsets = Maps.newHashMap();
   private final String producerId;
 
@@ -67,7 +65,6 @@ abstract class Channel {
     String transactionalId = config.transactionalPrefix() + name + config.transactionalSuffix();
     this.producer = clientFactory.createProducer(transactionalId);
     this.consumer = clientFactory.createConsumer(consumerGroupId);
-    this.admin = clientFactory.createAdmin();
 
     this.producerId = UUID.randomUUID().toString();
   }
@@ -160,8 +157,33 @@ void start() {
 
   void stop() {
     LOG.info("Channel stopping");
-    producer.close();
-    consumer.close();
-    admin.close();
+    RuntimeException failure = null;
+
+    try {
+      producer.close();
+    } catch (RuntimeException e) {
+      failure = appendFailure(failure, e);
+      LOG.warn("Error closing channel producer", e);
+    }
+
+    try {
+      consumer.close();
+    } catch (RuntimeException e) {
+      failure = appendFailure(failure, e);
+      LOG.warn("Error closing channel consumer", e);
+    }
+
+    if (failure != null) {
+      throw failure;
+    }
+  }
+
+  static RuntimeException appendFailure(RuntimeException failure, RuntimeException next) {
+    if (failure == null) {
+      return next;
+    }
+
+    failure.addSuppressed(next);
+    return failure;
   }
 }

kafka-connect/kafka-connect/src/main/java/org/apache/iceberg/connect/channel/CommitterImpl.java

@@ -159,31 +159,75 @@ public void stop() {
 
   @Override
   public void close(Collection<TopicPartition> closedPartitions) {
+    RuntimeException failure = null;
+
     // Always try to stop the worker to avoid duplicates.
-    stopWorker();
+    try {
+      stopWorker();
+    } catch (RuntimeException e) {
+      failure = Channel.appendFailure(failure, e);
+      LOG.warn("Committer {} failed to stop worker, continuing cleanup", taskId, e);
+    }
 
     // Defensive: close called without prior initialization (should not happen).
     if (!isInitialized.get()) {
       LOG.warn("Close unexpectedly called on committer {} without partition assignment", taskId);
+      if (failure != null) {
+        throw failure;
+      }
       return;
     }
 
-    // Empty partitions → task was stopped explicitly. Stop coordinator if running.
+    // Empty partitions: task was stopped explicitly. Stop coordinator if running.
     if (closedPartitions.isEmpty()) {
       LOG.info("Committer {} stopped. Closing coordinator.", taskId);
-      stopCoordinator();
+      try {
+        stopCoordinator();
+      } catch (RuntimeException e) {
+        failure = Channel.appendFailure(failure, e);
+        LOG.warn("Committer {} failed to stop coordinator", taskId, e);
+      }
+      if (failure != null) {
+        throw failure;
+      }
       return;
     }
 
     // Normal close: if leader partition is lost, stop coordinator.
-    if (hasLeaderPartition(closedPartitions)) {
+    boolean shouldStopCoordinator = false;
+    try {
+      shouldStopCoordinator = hasLeaderPartition(closedPartitions);
+    } catch (RuntimeException e) {
+      failure = Channel.appendFailure(failure, e);
+      shouldStopCoordinator = true;
+      LOG.warn(
+          "Committer {} could not determine whether leader partition was lost, stopping coordinator",
+          taskId,
+          e);
+    }
+
+    if (shouldStopCoordinator) {
       LOG.info("Committer {} lost leader partition. Stopping coordinator.", taskId);
-      stopCoordinator();
+      try {
+        stopCoordinator();
+      } catch (RuntimeException e) {
+        failure = Channel.appendFailure(failure, e);
+        LOG.warn("Committer {} failed to stop coordinator", taskId, e);
+      }
     }
 
     // Reset offsets to last committed to avoid data loss.
     LOG.info("Seeking to last committed offsets for worker {}.", taskId);
-    KafkaUtils.seekToLastCommittedOffsets(context);
+    try {
+      KafkaUtils.seekToLastCommittedOffsets(context);
+    } catch (RuntimeException e) {
+      failure = Channel.appendFailure(failure, e);
+      LOG.warn("Committer {} failed to seek to last committed offsets", taskId, e);
+    }
+
+    if (failure != null) {
+      throw failure;
+    }
   }
 
   @Override

kafka-connect/kafka-connect/src/main/java/org/apache/iceberg/connect/channel/CoordinatorThread.java

@@ -18,12 +18,14 @@
  */
 package org.apache.iceberg.connect.channel;
 
+import org.apache.kafka.connect.errors.ConnectException;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 class CoordinatorThread extends Thread {
   private static final Logger LOG = LoggerFactory.getLogger(CoordinatorThread.class);
   private static final String THREAD_NAME = "iceberg-coord";
+  private static final long TERMINATION_JOIN_TIMEOUT_MS = 60_000L;
 
   private final Coordinator coordinator;
   private volatile boolean terminated;
@@ -64,6 +66,32 @@ boolean isTerminated() {
 
   void terminate() {
     this.terminated = true;
-    coordinator.terminate();
+    RuntimeException failure = null;
+
+    try {
+      coordinator.terminate();
+    } catch (RuntimeException e) {
+      failure = Channel.appendFailure(failure, e);
+    }
+
+    if (Thread.currentThread() != this) {
+      try {
+        join(TERMINATION_JOIN_TIMEOUT_MS);
+        if (isAlive()) {
+          ConnectException timeout =
+              new ConnectException("Timed out waiting for coordinator thread shutdown");
+          failure = Channel.appendFailure(failure, timeout);
+        }
+      } catch (InterruptedException e) {
+        Thread.currentThread().interrupt();
+        ConnectException interrupted =
+            new ConnectException("Interrupted while waiting for coordinator thread shutdown", e);
+        failure = Channel.appendFailure(failure, interrupted);
+      }
+    }
+
+    if (failure != null) {
+      throw failure;
+    }
   }
 }

kafka-connect/kafka-connect/src/main/java/org/apache/iceberg/connect/channel/Worker.java

@@ -115,8 +115,23 @@ protected boolean receive(Envelope envelope) {
 
   @Override
   void stop() {
-    super.stop();
-    sinkWriter.close();
+    RuntimeException failure = null;
+
+    try {
+      super.stop();
+    } catch (RuntimeException e) {
+      failure = Channel.appendFailure(failure, e);
+    }
+
+    try {
+      sinkWriter.close();
+    } catch (RuntimeException e) {
+      failure = Channel.appendFailure(failure, e);
+    }
+
+    if (failure != null) {
+      throw failure;
+    }
   }
 
   void save(Collection<SinkRecord> sinkRecords) {

kafka-connect/kafka-connect/src/test/java/org/apache/iceberg/connect/channel/TestCommitterImpl.java

@@ -32,6 +32,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.iceberg.connect.IcebergSinkConfig;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableSet;
@@ -40,6 +41,8 @@
 import org.apache.kafka.clients.admin.MemberAssignment;
 import org.apache.kafka.clients.admin.MemberDescription;
 import org.apache.kafka.common.TopicPartition;
+import org.apache.kafka.connect.errors.ConnectException;
+import org.apache.kafka.connect.sink.SinkTaskContext;
 import org.junit.jupiter.api.Test;
 import org.mockito.MockedStatic;
 
@@ -165,4 +168,50 @@ public void testStartFailurePropagatesAsNotRunningException()
         .isInstanceOf(NotRunningException.class)
         .hasMessageContaining("Coordinator unexpectedly terminated");
   }
+
+  @Test
+  public void testCloseStopsCoordinatorWhenLeaderLookupFails()
+      throws NoSuchFieldException, IllegalAccessException {
+    CommitterImpl committer = new CommitterImpl();
+    setField(committer, "taskId", "demo1-foo-0");
+
+    Field initializedField = CommitterImpl.class.getDeclaredField("isInitialized");
+    initializedField.setAccessible(true);
+    ((AtomicBoolean) initializedField.get(committer)).set(true);
+
+    IcebergSinkConfig config = mock(IcebergSinkConfig.class);
+    when(config.connectGroupId()).thenReturn("demo1-foo-iceberg");
+    setField(committer, "config", config);
+
+    KafkaClientFactory clientFactory = mock(KafkaClientFactory.class);
+    when(clientFactory.createAdmin()).thenReturn(mock(Admin.class));
+    setField(committer, "clientFactory", clientFactory);
+
+    SinkTaskContext context = mock(SinkTaskContext.class);
+    setField(committer, "context", context);
+
+    CoordinatorThread coordinatorThread = mock(CoordinatorThread.class);
+    setField(committer, "coordinatorThread", coordinatorThread);
+
+    try (MockedStatic<KafkaUtils> mockKafkaUtils = mockStatic(KafkaUtils.class)) {
+      mockKafkaUtils
+          .when(() -> KafkaUtils.consumerGroupDescription(any(), any()))
+          .thenThrow(new ConnectException("Cannot retrieve members for consumer group"));
+
+      assertThatThrownBy(
+              () -> committer.close(ImmutableList.of(new TopicPartition("demo1-foo", 0))))
+          .isInstanceOf(ConnectException.class)
+          .hasMessageContaining("Cannot retrieve members for consumer group");
+
+      verify(coordinatorThread).terminate();
+      mockKafkaUtils.verify(() -> KafkaUtils.seekToLastCommittedOffsets(context));
+    }
+  }
+
+  private static void setField(Object target, String fieldName, Object value)
+      throws NoSuchFieldException, IllegalAccessException {
+    Field field = target.getClass().getDeclaredField(fieldName);
+    field.setAccessible(true);
+    field.set(target, value);
+  }
 }

kafka-connect/kafka-connect/src/test/java/org/apache/iceberg/connect/channel/TestCoordinatorThread.java

@@ -19,11 +19,14 @@
 package org.apache.iceberg.connect.channel;
 
 import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Mockito.doAnswer;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.timeout;
 import static org.mockito.Mockito.times;
 import static org.mockito.Mockito.verify;
 
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
 import org.junit.jupiter.api.Test;
 
 public class TestCoordinatorThread {
@@ -45,4 +48,36 @@ public void testRun() {
     verify(coordinator, timeout(1000)).stop();
     assertThat(coordinatorThread.isTerminated()).isTrue();
   }
+
+  @Test
+  public void testTerminateWaitsForCoordinatorStop() throws Exception {
+    Coordinator coordinator = mock(Coordinator.class);
+    CountDownLatch stopStarted = new CountDownLatch(1);
+    CountDownLatch stopCanFinish = new CountDownLatch(1);
+    doAnswer(
+            invocation -> {
+              stopStarted.countDown();
+              assertThat(stopCanFinish.await(5, TimeUnit.SECONDS)).isTrue();
+              return null;
+            })
+        .when(coordinator)
+        .stop();
+    CoordinatorThread coordinatorThread = new CoordinatorThread(coordinator);
+    coordinatorThread.start();
+
+    verify(coordinator, timeout(1000)).start();
+    verify(coordinator, timeout(1000).atLeast(1)).process();
+
+    Thread terminator = new Thread(coordinatorThread::terminate);
+    terminator.start();
+
+    assertThat(stopStarted.await(1, TimeUnit.SECONDS)).isTrue();
+    assertThat(terminator.isAlive()).isTrue();
+
+    stopCanFinish.countDown();
+    terminator.join(1000);
+
+    assertThat(terminator.isAlive()).isFalse();
+    assertThat(coordinatorThread.isAlive()).isFalse();
+  }
 }

kafka-connect/kafka-connect/src/test/java/org/apache/iceberg/connect/channel/TestWorker.java

@@ -22,6 +22,8 @@
 import static org.mockito.ArgumentMatchers.any;
 import static org.mockito.Mockito.mock;
 import static org.mockito.Mockito.mockStatic;
+import static org.mockito.Mockito.never;
+import static org.mockito.Mockito.verify;
 import static org.mockito.Mockito.when;
 
 import java.util.Map;
@@ -82,6 +84,7 @@ public void testSave() {
 
       Worker worker = new Worker(config, clientFactory, sinkWriter, context);
       worker.start();
+      verify(clientFactory, never()).createAdmin();
 
       // init consumer after subscribe()
       initConsumer();