This is an automated email from the ASF dual-hosted git repository.
wenjin272 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/flink-agents.git
The following commit(s) were added to refs/heads/main by this push:
new e5734203 [hotfix][runtime] Close the Kafka consumer when the producer
fails to close (#948)
e5734203 is described below
commit e5734203721eeacee51f98c6f4e0ef843fb229ba
Author: Weiqing Yang <[email protected]>
AuthorDate: Wed Aug 5 03:51:58 2026 -0700
[hotfix][runtime] Close the Kafka consumer when the producer fails to close
(#948)
---
.../runtime/actionstate/KafkaActionStateStore.java | 20 ++++-
.../actionstate/KafkaActionStateStoreTest.java | 89 ++++++++++++++++++++++
2 files changed, 107 insertions(+), 2 deletions(-)
diff --git
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
index f09e5cd2..ea7ea146 100644
---
a/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
+++
b/runtime/src/main/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStore.java
@@ -331,11 +331,27 @@ public class KafkaActionStateStore implements
ActionStateStore {
@Override
public void close() throws Exception {
+ Exception firstException = null;
if (producer != null) {
- producer.close();
+ try {
+ producer.close();
+ } catch (Exception e) {
+ firstException = e;
+ }
}
if (consumer != null) {
- consumer.close();
+ try {
+ consumer.close();
+ } catch (Exception e) {
+ if (firstException == null) {
+ firstException = e;
+ } else {
+ firstException.addSuppressed(e);
+ }
+ }
+ }
+ if (firstException != null) {
+ throw firstException;
}
}
diff --git
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
index 7cf829c0..c0a01bd0 100644
---
a/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
+++
b/runtime/src/test/java/org/apache/flink/agents/runtime/actionstate/KafkaActionStateStoreTest.java
@@ -21,9 +21,11 @@ import org.apache.flink.agents.api.Event;
import org.apache.flink.agents.api.InputEvent;
import org.apache.flink.agents.plan.AgentConfiguration;
import org.apache.flink.agents.plan.actions.Action;
+import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.MockConsumer;
import org.apache.kafka.clients.producer.MockProducer;
+import org.apache.kafka.clients.producer.Producer;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
@@ -38,6 +40,9 @@ import java.util.Map;
import static
org.apache.kafka.clients.consumer.internals.AutoOffsetResetStrategy.EARLIEST;
import static org.assertj.core.api.Assertions.assertThat;
import static org.junit.jupiter.api.Assertions.*;
+import static org.mockito.Mockito.doThrow;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.verify;
/** Unit tests for {@link KafkaActionStateStore}. */
public class KafkaActionStateStoreTest {
@@ -254,4 +259,88 @@ public class KafkaActionStateStoreTest {
ActionStateUtil.generateKey(TEST_KEY, 3L,
testAction, testEvent)))
.isEqualTo(thirdState);
}
+
+ /** Contract: the consumer is closed even when closing the producer
throws. */
+ @Test
+ @SuppressWarnings("unchecked")
+ void testCloseClosesConsumerWhenProducerCloseFails() {
+ Producer<String, ActionState> failingProducer = mock(Producer.class);
+ Consumer<String, ActionState> consumer = mock(Consumer.class);
+ doThrow(new RuntimeException("producer close
failed")).when(failingProducer).close();
+
+ KafkaActionStateStore store =
+ new KafkaActionStateStore(
+ actionStates,
+ new AgentConfiguration(),
+ failingProducer,
+ consumer,
+ TEST_TOPIC);
+
+ assertThrows(RuntimeException.class, store::close);
+
+ verify(consumer).close();
+ }
+
+ /**
+ * Contract: when both closes fail, the producer's exception is the one
thrown and the
+ * consumer's is attached to it as a suppressed exception, so neither
failure is lost.
+ */
+ @Test
+ @SuppressWarnings("unchecked")
+ void testCloseKeepsProducerFailureWhenBothCloseFail() {
+ Producer<String, ActionState> failingProducer = mock(Producer.class);
+ Consumer<String, ActionState> failingConsumer = mock(Consumer.class);
+ RuntimeException producerFailure = new RuntimeException("producer
close failed");
+ RuntimeException consumerFailure = new RuntimeException("consumer
close failed");
+ doThrow(producerFailure).when(failingProducer).close();
+ doThrow(consumerFailure).when(failingConsumer).close();
+
+ KafkaActionStateStore store =
+ new KafkaActionStateStore(
+ actionStates,
+ new AgentConfiguration(),
+ failingProducer,
+ failingConsumer,
+ TEST_TOPIC);
+
+ RuntimeException thrown = assertThrows(RuntimeException.class,
store::close);
+
+ assertThat(thrown).isSameAs(producerFailure);
+ assertThat(thrown.getSuppressed()).containsExactly(consumerFailure);
+ }
+
+ /**
+ * Contract: when only the consumer close fails, its exception reaches the
caller unchanged,
+ * with nothing attached as suppressed.
+ */
+ @Test
+ @SuppressWarnings("unchecked")
+ void testCloseThrowsConsumerFailureWhenOnlyConsumerCloseFails() {
+ Producer<String, ActionState> producer = mock(Producer.class);
+ Consumer<String, ActionState> failingConsumer = mock(Consumer.class);
+ RuntimeException consumerFailure = new RuntimeException("consumer
close failed");
+ doThrow(consumerFailure).when(failingConsumer).close();
+
+ KafkaActionStateStore store =
+ new KafkaActionStateStore(
+ actionStates,
+ new AgentConfiguration(),
+ producer,
+ failingConsumer,
+ TEST_TOPIC);
+
+ RuntimeException thrown = assertThrows(RuntimeException.class,
store::close);
+
+ assertThat(thrown).isSameAs(consumerFailure);
+ assertThat(thrown.getSuppressed()).isEmpty();
+ }
+
+ /** Contract: both the producer and the consumer are closed when neither
close fails. */
+ @Test
+ void testCloseClosesProducerAndConsumer() throws Exception {
+ actionStateStore.close();
+
+ assertThat(mockProducer.closed()).isTrue();
+ assertThat(mockConsumer.closed()).isTrue();
+ }
}
