LoomKafkaProducer|Consumer let the background thread finish itself (#…

…4013)
knative-extensions · Sep 9, 2024 · 045f7e2 · 045f7e2
1 parent 8cd6f05
commit 045f7e2
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 18 deletions.
diff --git a/...oom/src/main/java/dev/knative/eventing/kafka/broker/dispatcherloom/LoomKafkaConsumer.java b/...oom/src/main/java/dev/knative/eventing/kafka/broker/dispatcherloom/LoomKafkaConsumer.java
@@ -27,6 +27,7 @@
 import java.util.Map;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.kafka.clients.consumer.Consumer;
 import org.apache.kafka.clients.consumer.ConsumerRebalanceListener;
@@ -44,15 +45,13 @@ public class LoomKafkaConsumer<K, V> implements ReactiveKafkaConsumer<K, V> {
     private final Consumer<K, V> consumer;
     private final BlockingQueue<Runnable> taskQueue;
     private final AtomicBoolean isClosed;
-    private final AtomicBoolean isFinished;
     private final Thread taskRunnerThread;
     private final Promise<Void> closePromise = Promise.promise();
 
     public LoomKafkaConsumer(Vertx vertx, Consumer<K, V> consumer) {
         this.consumer = consumer;
         this.taskQueue = new LinkedBlockingQueue<>();
         this.isClosed = new AtomicBoolean(false);
-        this.isFinished = new AtomicBoolean(false);
 
         if (Boolean.parseBoolean(System.getenv("ENABLE_VIRTUAL_THREADS"))) {
             this.taskRunnerThread = Thread.ofVirtual().start(this::processTaskQueue);
@@ -74,14 +73,15 @@ private void processTaskQueue() {
         // Process queue elements until this is closed and the tasks queue is empty
         while (!isClosed.get() || !taskQueue.isEmpty()) {
             try {
-                taskQueue.take().run();
+                Runnable task = taskQueue.poll(2000, TimeUnit.MILLISECONDS);
+                if (task != null) {
+                  task.run();
+                }
             } catch (InterruptedException e) {
                 logger.debug("Interrupted while waiting for task", e);
                 break;
             }
         }
-
-        isFinished.set(true);
     }
 
     @Override
@@ -126,16 +126,6 @@ public Future<Void> close() {
                 }
                 logger.debug("Queue is empty");
 
-                if (!isFinished.get()) {
-                    logger.debug("Background thread not finished yet, waiting for it to complete");
-                    Thread.sleep(2000L);
-
-                    if (!isFinished.get()) {
-                        logger.debug("Background thread still not finished yet, interrupting background thread");
-                        taskRunnerThread.interrupt();
-                    }
-                }
-
                 taskRunnerThread.join();
                 closePromise.tryComplete();
 

diff --git a/...-loom/src/main/java/dev/knative/eventing/kafka/broker/receiverloom/LoomKafkaProducer.java b/...-loom/src/main/java/dev/knative/eventing/kafka/broker/receiverloom/LoomKafkaProducer.java
@@ -26,6 +26,7 @@
 import java.util.Objects;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.LinkedBlockingQueue;
+import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicBoolean;
 import org.apache.kafka.clients.producer.Producer;
 import org.apache.kafka.clients.producer.ProducerRecord;
@@ -84,7 +85,11 @@ private void sendFromQueue() {
         // Process queue elements until this is closed and the tasks queue is empty
         while (!isClosed.get() || !eventQueue.isEmpty()) {
             try {
-                final var recordPromise = eventQueue.take();
+                final var recordPromise = eventQueue.poll(2000, TimeUnit.MILLISECONDS);
+                if (recordPromise == null) {
+                    continue;
+                }
+
                 final var startedSpan = this.tracer == null
                         ? null
                         : this.tracer.prepareSendMessage(recordPromise.getContext(), recordPromise.getRecord());
@@ -140,8 +145,6 @@ public Future<Void> close() {
                     logger.debug("Waiting for the eventQueue to become empty");
                     Thread.sleep(2000L);
                 }
-                logger.debug("Interrupting sendFromQueueThread thread");
-                sendFromQueueThread.interrupt();
                 logger.debug("Waiting for sendFromQueueThread thread to complete");
                 sendFromQueueThread.join();
                 logger.debug("Closing the producer");