josefk31 commented on code in PR #22505:
URL: https://github.com/apache/kafka/pull/22505#discussion_r3422386916
##########
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) {
Review Comment:
Nit: could we explicitly name this `pollMs`? I think that would make it more
clear to users that the polling must be milliseconds. We could also just pass
through the same arguments to `queue.poll` as well. I think making mistakes
with units can be quite common and is something to be actively avoided.
##########
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:
Hmm; not in the PR but is this handling correct? I read through Java docs
and it seems interupt works in the following way:
1. Once we reach the catch block, the interrupted status is cleared - so the
thread state is not interrupted.
2. In this thread I assume we just want to keep the thread interrupted state
and allow the threads owner to figure out what to do?
3. In that case - according to the docs: "Code that throws another type of
exception with the InterruptedException as
[cause](https://docs.oracle.com/en/java/javase/26/docs/api/java.base/java/lang/Throwable.html#getCause()),
or the InterruptedException as a [suppressed
exception](https://docs.oracle.com/en/java/javase/26/docs/api/java.base/java/lang/Throwable.html#addSuppressed(java.lang.Throwable)),
should also restore the interrupted status before throwing the exception."
4. This implies to me that we must rethrow the exception or call
`Thread.currentThread().interrupt()`
Source:
https://docs.oracle.com/en/java/javase/26/docs/api/java.base/java/lang/Thread.html#thread-interruption
##########
raft/src/main/java/org/apache/kafka/raft/KafkaNetworkChannel.java:
##########
@@ -116,29 +118,35 @@ public int newCorrelationId() {
}
@Override
- public void send(RaftRequest.Outbound request) {
+ public CompletionStage<RaftResponse.Inbound> send(RaftRequest.Outbound
request) {
Node node = request.destination();
if (node != null) {
- requestThread.sendRequest(new RequestAndCompletionHandler(
- request.createdTimeMs(),
- node,
- buildRequest(request.data()),
- response -> sendOnComplete(request, response)
- ));
- } else
- sendCompleteFuture(request, errorResponse(request.data(),
Errors.BROKER_NOT_AVAILABLE));
- }
-
- private void sendCompleteFuture(RaftRequest.Outbound request, ApiMessage
message) {
- RaftResponse.Inbound response = new RaftResponse.Inbound(
- request.correlationId(),
- message,
- request.destination()
- );
- request.completion.complete(response);
+ var future = new CompletableFuture<RaftResponse.Inbound>();
+ requestThread.sendRequest(
+ new RequestAndCompletionHandler(
+ request.createdTimeMs(),
+ node,
+ buildRequest(request.data()),
+ response -> sendOnComplete(request, response, future)
+ )
+ );
+ return future;
+ } else {
+ return CompletableFuture.completedFuture(
+ new RaftResponse.Inbound(
+ request.correlationId(),
+ errorResponse(request.data(), Errors.BROKER_NOT_AVAILABLE),
+ request.destination()
+ )
+ );
+ }
}
- private void sendOnComplete(RaftRequest.Outbound request, ClientResponse
clientResponse) {
+ private void sendOnComplete(
Review Comment:
Very small nitpick: it may be clearer to read this as `onSendCompleted`
which emphasizes that this is an action which is performed when we complete the
method.
##########
raft/src/test/java/org/apache/kafka/raft/internals/BlockingMessageQueueTest.java:
##########
@@ -21,40 +21,90 @@
import org.junit.jupiter.api.Test;
import org.mockito.Mockito;
+import java.util.Optional;
+
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
-import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
class BlockingMessageQueueTest {
@Test
public void testOfferAndPoll() {
- BlockingMessageQueue queue = new BlockingMessageQueue();
+ var queue = new BlockingMessageQueue();
assertTrue(queue.isEmpty());
- assertNull(queue.poll(0));
+ assertEquals(Optional.empty(), queue.poll(0));
- RaftMessage message1 = Mockito.mock(RaftMessage.class);
- queue.add(message1);
+ var mockMessage1 = Mockito.mock(RaftMessage.class);
+ queue.add(mockMessage1);
assertFalse(queue.isEmpty());
- assertEquals(message1, queue.poll(0));
+ var entry1 = queue.poll(0);
+ assertTrue(entry1.isPresent());
+ assertEquals(mockMessage1, entry1.get().message());
assertTrue(queue.isEmpty());
- RaftMessage message2 = Mockito.mock(RaftMessage.class);
- RaftMessage message3 = Mockito.mock(RaftMessage.class);
- queue.add(message2);
- queue.add(message3);
+ var mockMessage2 = Mockito.mock(RaftMessage.class);
+ var mockMessage3 = Mockito.mock(RaftMessage.class);
+ queue.add(mockMessage2);
+ queue.add(mockMessage3);
assertFalse(queue.isEmpty());
- assertEquals(message2, queue.poll(0));
- assertEquals(message3, queue.poll(0));
+ var entry2 = queue.poll(0);
+ var entry3 = queue.poll(0);
+ assertTrue(entry2.isPresent());
+ assertTrue(entry3.isPresent());
+ assertEquals(mockMessage2, entry2.get().message());
+ assertEquals(mockMessage3, entry3.get().message());
}
@Test
public void testWakeupFromPoll() {
- BlockingMessageQueue queue = new BlockingMessageQueue();
+ var queue = new BlockingMessageQueue();
queue.wakeup();
- assertNull(queue.poll(Long.MAX_VALUE));
+ assertEquals(Optional.empty(), queue.poll(Long.MAX_VALUE));
}
-}
\ No newline at end of file
+ @Test
+ public void testWakeupsAreTransparentToIsEmptyAndDrainedOnPoll() {
+ var queue = new BlockingMessageQueue();
+
+ // Wakeups alone should not affect isEmpty
+ queue.wakeup();
+ queue.wakeup();
+ assertTrue(queue.isEmpty());
+
+ // Adding a real message makes the queue non-empty
+ var mockMessage = Mockito.mock(RaftMessage.class);
+ queue.add(mockMessage);
+ assertFalse(queue.isEmpty());
+
+ // Poll should drain all wakeups and return the message in one call
+ var entry = queue.poll(0);
+ assertTrue(entry.isPresent());
+ assertEquals(mockMessage, entry.get().message());
+ assertTrue(queue.isEmpty());
+ }
+
+ @Test
+ public void testAddRejectsNullMessage() {
+ var queue = new BlockingMessageQueue();
+
+ // Null message should be rejected
+ IllegalArgumentException exception = assertThrows(
Review Comment:
Not really a comment but more a question - here we handle the `null`
validation before scheduling the `CompletionStage`. We could also do so within
the `CompletionStage` computation. Is there a good way to decide when to favour
the former or the latter? I think in this case it makes sense since "we should
not schedule something obviously wrong" but it also creates a bit of doubt when
something might crash and who is responsible for handling the resulting
problems.
--
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]
