From bf78e7bbf9f25c57493a776137e7ef8443dac1e2 Mon Sep 17 00:00:00 2001
From: Ashwin Somasundara <ashwin.somasundara@gmail.com>
Date: Sun, 1 Nov 2026 13:03:37 -0500
Subject: [PATCH] Release Kafka sink lock on early return and failure

Signed-off-by: Ashwin Somasundara <ashwin.somasundara@gmail.com>
---
 .../plugins/kafka/sink/KafkaSink.java         |  5 ++--
 .../plugins/kafka/sink/KafkaSinkTest.java     | 27 +++++++++++++++++++
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSink.java b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSink.java
index 97c93d22de..2de1cd47c7 100644
--- a/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSink.java
+++ b/data-prepper-plugins/kafka-plugins/src/main/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSink.java
@@ -116,10 +116,10 @@ private void doInitializeInternal() {
 
     @Override
     public void doOutput(Collection<Record<Event>> records) {
-        reentrantLock.lock();
         if (records.isEmpty()) {
             return;
         }
+        reentrantLock.lock();
         try {
             // TODO: Looks like this call to prepareTopicAndSchema is unnecessary as it is 
             // done in createProducer().
@@ -133,8 +133,9 @@ public void doOutput(Collection<Record<Event>> records) {
         } catch (Exception e) {
             LOG.error("Failed to setup the Kafka sink Plugin.", e);
             throw new RuntimeException(e.getMessage());
+        } finally {
+            reentrantLock.unlock();
         }
-        reentrantLock.unlock();
     }
 
     private void prepareTopicAndSchema() {
diff --git a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkTest.java b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkTest.java
index 6d57c7803a..20193f49ac 100644
--- a/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkTest.java
+++ b/data-prepper-plugins/kafka-plugins/src/test/java/org/opensearch/dataprepper/plugins/kafka/sink/KafkaSinkTest.java
@@ -47,7 +47,10 @@
 import java.util.concurrent.Executors;
 import java.util.concurrent.FutureTask;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.locks.ReentrantLock;
 
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.equalTo;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertThrows;
 import static org.mockito.ArgumentMatchers.any;
@@ -173,6 +176,30 @@ public void doOutputEmptyRecordsTest() {
 
     }
 
+    @Test
+    public void doOutputReleasesLockWhenRecordsAreEmpty() {
+        final KafkaSink objectUnderTest = createObjectUnderTest();
+
+        objectUnderTest.doOutput(Arrays.asList());
+
+        assertThat(getLock(objectUnderTest).isLocked(), equalTo(false));
+    }
+
+    @Test
+    public void doOutputReleasesLockWhenSubmitFails() {
+        final Collection records = Arrays.asList(new Record(event));
+        when(executorService.submit(any(ProducerWorker.class))).thenThrow(new RuntimeException());
+        final KafkaSink objectUnderTest = createObjectUnderTest();
+
+        assertThrows(RuntimeException.class, () -> objectUnderTest.doOutput(records));
+
+        assertThat(getLock(objectUnderTest).isLocked(), equalTo(false));
+    }
+
+    private ReentrantLock getLock(final KafkaSink kafkaSink) {
+        return (ReentrantLock) ReflectionTestUtils.getField(kafkaSink, "reentrantLock");
+    }
+
     @Test
     public void shutdownTest() {
         final KafkaSink objectUnderTest = createObjectUnderTest();