jsancio commented on code in PR #22505:
URL: https://github.com/apache/kafka/pull/22505#discussion_r3431008086


##########
raft/src/main/java/org/apache/kafka/raft/internals/BlockingMessageQueue.java:
##########
@@ -17,60 +17,104 @@
 package org.apache.kafka.raft.internals;
 
 import org.apache.kafka.common.errors.InterruptException;
-import org.apache.kafka.common.protocol.ApiMessage;
 import org.apache.kafka.raft.RaftMessage;
 import org.apache.kafka.raft.RaftMessageQueue;
 
+import java.util.Optional;
 import java.util.concurrent.BlockingQueue;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.CompletionStage;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
 public class BlockingMessageQueue implements RaftMessageQueue {
-    private static final RaftMessage WAKEUP_MESSAGE = new RaftMessage() {
+    private final BlockingQueue<InternalQueueEntry> queue = new 
LinkedBlockingQueue<>();
+    private final AtomicInteger messageCount = new AtomicInteger(0);
+
+    /**
+     * Internal queue entry type used to discriminate between messages and 
wakeup signals.
+     *
+     * This sealed interface ensures type safety when polling the queue.
+     */
+    private sealed interface InternalQueueEntry { }
+
+    /**
+     * Marker entry used to unblock threads waiting on {@link #poll(long)} 
without delivering a message.
+     *
+     * Wakeup entries are drained during polling and do not contribute to the 
message count.
+     */
+    private record WakeupMarker() implements InternalQueueEntry { }
+
+    private static final WakeupMarker WAKEUP = new WakeupMarker();
+
+    /**
+     * A queue entry that contains a message and its associated future.
+     */
+    private static final class MessageEntry implements QueueEntry, 
InternalQueueEntry {
+        private final CompletableFuture<RaftMessage> future = new 
CompletableFuture<>();
+        private final RaftMessage message;
+
+        MessageEntry(RaftMessage message) {
+            this.message = message;
+        }
+
         @Override
-        public int correlationId() {
-            return 0;
+        public RaftMessage message() {
+            return message;
         }
 
         @Override
-        public ApiMessage data() {
-            return null;
+        public CompletableFuture<RaftMessage> future() {
+            return future;
         }
-    };
 
-    private final BlockingQueue<RaftMessage> queue = new 
LinkedBlockingQueue<>();
-    private final AtomicInteger size = new AtomicInteger(0);
+        @Override
+        public String toString() {
+            return String.format(
+                "MessageEntry(message=%s, future.isDone=%s)",
+                message,
+                future.isDone()
+            );
+        }
+    }
 
     @Override
-    public RaftMessage poll(long timeoutMs) {
+    public Optional<QueueEntry> poll(long timeoutMs) {
         try {
-            RaftMessage message = queue.poll(timeoutMs, TimeUnit.MILLISECONDS);
-            if (message == null || message == WAKEUP_MESSAGE) {
-                return null;
-            } else {
-                size.decrementAndGet();
-                return message;
+            InternalQueueEntry entry = queue.poll(timeoutMs, 
TimeUnit.MILLISECONDS);
+            // Drain all wakeup markers until we find a message or the queue 
is empty
+            while (entry instanceof WakeupMarker) {
+                entry = queue.poll();
             }
+            if (entry instanceof MessageEntry messageEntry) {
+                messageCount.decrementAndGet();
+                return Optional.of(messageEntry);
+            }
+            return Optional.empty();
         } catch (InterruptedException e) {
             throw new InterruptException(e);

Review Comment:
   The constructor for Kafka's unchecked InterruptException sets the interrupt 
flag.



-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]

Reply via email to