KAFKA-10826; Ensure raft io thread respects linger timeout

raft/src/main/java/org/apache/kafka/raft/KafkaRaftClient.java

@@ -1800,8 +1800,15 @@ public Long scheduleAppend(int epoch, List<T> records) {
             return Long.MAX_VALUE;
         }
 
+        boolean isFirstAppend = accumulator.isEmpty();
         Long offset = accumulator.append(epoch, records);
-        if (accumulator.needsDrain(time.milliseconds())) {
+
+        // Wakeup the network channel if either this is the first append
+        // or the accumulator is ready to drain now. Checking for the first
+        // append ensures that we give the IO thread a chance to observe
+        // the linger timeout so that it can schedule its own wakeup in case
+        // there are no additional appends.
+        if (isFirstAppend || accumulator.needsDrain(time.milliseconds())) {
             channel.wakeup();
         }
         return offset;

raft/src/main/java/org/apache/kafka/raft/internals/BatchAccumulator.java

@@ -141,6 +141,7 @@ private BatchBuilder<T> maybeAllocateBatch(int batchSize) {
             startNewBatch();
         } else if (!currentBatch.hasRoomFor(batchSize)) {
             completeCurrentBatch();
+            startNewBatch();
         }
         return currentBatch;
     }
@@ -271,22 +272,19 @@ private List<CompletedBatch<T>> drainCompleted() {
         }
     }
 
+    public boolean isEmpty() {
+        // The linger timer begins running when we have pending batches.
+        // We use this to infer when the accumulator is empty to avoid the
+        // need to acquire the append lock.
+        return !lingerTimer.isRunning();
+    }
+
     /**
-     * Get the number of batches including the one that is currently being
-     * written to (if it exists).
+     * Get the number of completed batches which are ready to be drained.
+     * This does not include the batch that is currently being filled.
      */
-    public int count() {
-        appendLock.lock();
-        try {
-            int count = completed.size();
-            if (currentBatch != null) {
-                return count + 1;
-            } else {
-                return count;
-            }
-        } finally {
-            appendLock.unlock();
-        }
+    public int numCompletedBatches() {
+        return completed.size();
     }
 
     @Override

raft/src/test/java/org/apache/kafka/raft/KafkaRaftClientTest.java

@@ -477,6 +477,74 @@ public void testAccumulatorClearedAfterBecomingUnattached() throws Exception {
         Mockito.verify(memoryPool).release(buffer);
     }
 
+    @Test
+    public void testChannelWokenUpIfLingerTimeoutReachedWithoutAppend() throws Exception {
+        // This test verifies that the client will set its poll timeout accounting
+        // for the lingerMs of a pending append
+
+        int localId = 0;
+        int otherNodeId = 1;
+        int lingerMs = 50;
+        Set<Integer> voters = Utils.mkSet(localId, otherNodeId);
+
+        RaftClientTestContext context = new RaftClientTestContext.Builder(localId, voters)
+            .withAppendLingerMs(lingerMs)
+            .build();
+
+        context.becomeLeader();
+        assertEquals(OptionalInt.of(localId), context.currentLeader());
+        assertEquals(1L, context.log.endOffset().offset);
+
+        int epoch = context.currentEpoch();
+        assertEquals(1L, context.client.scheduleAppend(epoch, singletonList("a")));
+        assertTrue(context.channel.wakeupRequested());
+
+        context.client.poll();
+        assertEquals(OptionalLong.of(lingerMs), context.channel.lastReceiveTimeout());
+
+        context.time.sleep(20);
+        context.client.poll();
+        assertEquals(OptionalLong.of(30), context.channel.lastReceiveTimeout());
+
+        context.time.sleep(30);
+        context.client.poll();
+        assertEquals(2L, context.log.endOffset().offset);
+    }
+
+    @Test
+    public void testChannelWokenUpIfLingerTimeoutReachedDuringAppend() throws Exception {
+        // This test verifies that the client will get woken up immediately
+        // if the linger timeout has expired during an append
+
+        int localId = 0;
+        int otherNodeId = 1;
+        int lingerMs = 50;
+        Set<Integer> voters = Utils.mkSet(localId, otherNodeId);
+
+        RaftClientTestContext context = new RaftClientTestContext.Builder(localId, voters)
+            .withAppendLingerMs(lingerMs)
+            .build();
+
+        context.becomeLeader();
+        assertEquals(OptionalInt.of(localId), context.currentLeader());
+        assertEquals(1L, context.log.endOffset().offset);
+
+        int epoch = context.currentEpoch();
+        assertEquals(1L, context.client.scheduleAppend(epoch, singletonList("a")));
+        assertTrue(context.channel.wakeupRequested());
+
+        context.client.poll();
+        assertFalse(context.channel.wakeupRequested());
+        assertEquals(OptionalLong.of(lingerMs), context.channel.lastReceiveTimeout());
+
+        context.time.sleep(lingerMs);
+        assertEquals(2L, context.client.scheduleAppend(epoch, singletonList("b")));
+        assertTrue(context.channel.wakeupRequested());
+
+        context.client.poll();
+        assertEquals(3L, context.log.endOffset().offset);
+    }
+
     @Test
     public void testHandleEndQuorumRequest() throws Exception {
         int localId = 0;

raft/src/test/java/org/apache/kafka/raft/MockNetworkChannel.java

@@ -24,10 +24,16 @@
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
+import java.util.OptionalLong;
+import java.util.concurrent.atomic.AtomicBoolean;
 import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
 
 public class MockNetworkChannel implements NetworkChannel {
     private final AtomicInteger requestIdCounter;
+    private final AtomicBoolean wakeupRequested = new AtomicBoolean(false);
+    private final AtomicLong lastReceiveTimeout = new AtomicLong(-1);
+
     private List<RaftMessage> sendQueue = new ArrayList<>();
     private List<RaftMessage> receiveQueue = new ArrayList<>();
     private Map<Integer, InetSocketAddress> addressCache = new HashMap<>();
@@ -59,50 +65,34 @@ public void send(RaftMessage message) {
 
     @Override
     public List<RaftMessage> receive(long timeoutMs) {
+        wakeupRequested.set(false);
+        lastReceiveTimeout.set(timeoutMs);
         List<RaftMessage> messages = receiveQueue;
         receiveQueue = new ArrayList<>();
         return messages;
     }
 
-    @Override
-    public void wakeup() {}
-
-    @Override
-    public void updateEndpoint(int id, InetSocketAddress address) {
-        addressCache.put(id, address);
+    OptionalLong lastReceiveTimeout() {
+        long timeout = lastReceiveTimeout.get();
+        if (timeout < 0) {
+            return OptionalLong.empty();
+        } else {
+            return OptionalLong.of(timeout);
+        }
     }
 
-    public RaftRequest.Outbound drainNextRequest(int destinationId) {
-        Iterator<RaftMessage> iterator = sendQueue.iterator();
-        while (iterator.hasNext()) {
-            RaftMessage message = iterator.next();
-            if (message instanceof RaftRequest.Outbound) {
-                RaftRequest.Outbound request = (RaftRequest.Outbound) message;
-                if (request.destinationId() == destinationId) {
-                    iterator.remove();
-                    return request;
-                }
-            }
-        }
-        return null;
+    boolean wakeupRequested() {
+        return wakeupRequested.get();
     }
 
-    public RaftResponse.Outbound drainNextResponse(int correlationId) {
-        Iterator<RaftMessage> iterator = sendQueue.iterator();
-        while (iterator.hasNext()) {
-            RaftMessage message = iterator.next();
-            if (message.correlationId() == correlationId) {
-                if (!(message instanceof RaftResponse.Outbound)) {
-                    throw new IllegalStateException("Message " + message +
-                        " is not an outbound response as we expected");
-                }
+    @Override
+    public void wakeup() {
+        wakeupRequested.set(true);
+    }
 
-                RaftResponse.Outbound response = (RaftResponse.Outbound) message;
-                iterator.remove();
-                return response;
-            }
-        }
-        return null;
+    @Override
+    public void updateEndpoint(int id, InetSocketAddress address) {
+        addressCache.put(id, address);
     }
 
     public List<RaftMessage> drainSendQueue() {
@@ -148,9 +138,4 @@ public void mockReceive(RaftMessage message) {
         receiveQueue.add(message);
     }
 
-    void clear() {
-        sendQueue.clear();
-        receiveQueue.clear();
-        requestIdCounter.set(0);
-    }
 }

raft/src/test/java/org/apache/kafka/raft/internals/BatchAccumulatorTest.java

@@ -72,6 +72,7 @@ public void testLingerIgnoredIfAccumulatorEmpty() {
             maxBatchSize
         );
 
+        assertTrue(acc.isEmpty());
         assertFalse(acc.needsDrain(time.milliseconds()));
         assertEquals(Long.MAX_VALUE - time.milliseconds(), acc.timeUntilDrain(time.milliseconds()));
     }
@@ -96,13 +97,16 @@ public void testLingerBeginsOnFirstWrite() {
         time.sleep(15);
         assertEquals(baseOffset, acc.append(leaderEpoch, singletonList("a")));
         assertEquals(lingerMs, acc.timeUntilDrain(time.milliseconds()));
+        assertFalse(acc.isEmpty());
 
         time.sleep(lingerMs / 2);
         assertEquals(lingerMs / 2, acc.timeUntilDrain(time.milliseconds()));
+        assertFalse(acc.isEmpty());
 
         time.sleep(lingerMs / 2);
         assertEquals(0, acc.timeUntilDrain(time.milliseconds()));
         assertTrue(acc.needsDrain(time.milliseconds()));
+        assertFalse(acc.isEmpty());
     }
 
     @Test
@@ -211,12 +215,13 @@ public void testMultipleBatchAccumulation() {
             maxBatchSize
         );
 
-        while (acc.count() < 3) {
+        // Append entries until we have 4 batches to drain (3 completed, 1 building)
+        while (acc.numCompletedBatches() < 3) {
             acc.append(leaderEpoch, singletonList("foo"));
         }
 
         List<BatchAccumulator.CompletedBatch<String>> batches = acc.drain();
-        assertEquals(3, batches.size());
+        assertEquals(4, batches.size());
         assertTrue(batches.stream().allMatch(batch -> batch.data.sizeInBytes() <= maxBatchSize));
     }