FlowControlledReplyChannelWriter.stream() can lose the QueryComplete when the buffer closes while another thread holds the flow gate #523

Basic information

• Axon Server Connector Java version: 2026.0.0 and 2026.1.0 (FlowControlledReplyChannelWriter is byte-identical in both)
• JDK version: 21 (Temurin 21.0.11)
• Complete executable reproducer if available (e.g. GitHub Repo): LostCompletionRepro.java below, a single class run against the published connector jar

Context

We found this while chasing a different bug: AF5 point-to-point query responses that were never sent, which turned out to be in FlowControlledResponseSender in axon-server-connector (AxonIQ/axoniq-framework#666). While going through the reply path for that one, we came across this check-then-act race one layer down. We haven't tied it to a production hang, and we can't say how often it happens unforced. We're reporting it because the source shows the gap and a forced interleaving reproduces it.

Steps to reproduce

The race, in io.axoniq.axonserver.connector.impl.FlowControlledReplyChannelWriter:

private void stream() {                                       // line 94
    do {
        if (!flowGate.compareAndSet(false, true)) {
            return;                                           // a concurrent caller just leaves
        }
        try {
            long currentPermits = requestedRef.get();
            long sent = send(currentPermits);                 // the only path to completeIfBufferIsClosedAndEmpty()
            requestedRef.getAndAccumulate(sent, (current, s) -> current - s);
        } finally {
            flowGate.set(false);
        }
    } while (requestedRef.get() > 0 && !buffer.isEmpty());   // line 106: no check for a pending completion
}

The completion is only sent from send() → completeIfBufferIsClosedAndEmpty() (lines 110, 116 and 123–133), which only runs while the caller holds flowGate.

1. Thread A handles a flow-control request (QueryChannelImpl.handleFlowControlRequest → writer.request() → stream()) and holds flowGate. It drains the buffer, finds it empty, and sees buffer.closed() is still false.
2. Thread B finishes the handler and closes the buffer. BlockingCloseableBuffer.close() sets closed and runs the onAvailable callback, which is the writer's stream(), synchronously on B. B's CAS on flowGate fails, so it returns.
3. Thread A releases the gate. Line 106 sees an empty buffer and exits.

Nobody sends QueryComplete.

With Axoniq Framework 5 on the responding side, A probably holds the gate for longer than with AF4: FlowControlledResponseSender.request(n) runs synchronously inside A's stream() and passes through several message-stream locks, where AF4 handed it to an executor. That is our reading of the code, not something we measured.

To reproduce, compile and run the class below against the connector jar and its dependencies. A hook on the buffer closes it from a second thread at exactly the moment the writer reads closed(), then returns the value from before the close, which is what thread A would have seen.

import io.axoniq.axonserver.connector.ErrorCategory;
import io.axoniq.axonserver.connector.ReplyChannel;
import io.axoniq.axonserver.connector.impl.FlowControlledReplyChannelWriter;
import io.axoniq.axonserver.connector.impl.buffer.BlockingCloseableBuffer;
import io.axoniq.axonserver.connector.impl.DisposableReadonlyBuffer;
import io.axoniq.axonserver.grpc.ErrorMessage;

import java.util.Optional;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

/**
 * Deterministic reproducer: the producer closes the response buffer while the gRPC thread is inside
 * FlowControlledReplyChannelWriter.stream(), after its last closed() check but before it releases flowGate. The
 * producer's onAvailable -> stream() loses the CAS and returns; the gRPC thread's do/while only re-checks
 * "permits > 0 && !buffer.isEmpty()", so complete() (QueryComplete) is never sent.
 */
public class LostCompletionRepro {
    public static void main(String[] args) throws Exception {
        BlockingCloseableBuffer<String> real = new BlockingCloseableBuffer<>();
        AtomicBoolean drained = new AtomicBoolean();
        AtomicBoolean injected = new AtomicBoolean();
        DisposableReadonlyBuffer<String> hooked = new DisposableReadonlyBuffer<>() {
            public Optional<String> poll() { Optional<String> r = real.poll(); if (r.isPresent()) drained.set(true); return r; }
            public boolean isEmpty() { return real.isEmpty(); }
            public int capacity() { return real.capacity(); }
            public void onAvailable(Runnable r) { real.onAvailable(r); }
            public void dispose() { }
            public Optional<ErrorMessage> error() { return real.error(); }
            public boolean closed() {
                boolean c = real.closed();
                if (drained.get() && !c && injected.compareAndSet(false, true)) {
                    // the producer thread completes the stream right now, while this (gRPC) thread holds flowGate
                    Thread producer = new Thread(real::close, "producer");
                    producer.start();
                    try { producer.join(); } catch (InterruptedException e) { throw new RuntimeException(e); }
                }
                return c;
            }
        };
        AtomicInteger sent = new AtomicInteger();
        AtomicBoolean completed = new AtomicBoolean();
        ReplyChannel<String> toAxonServer = new ReplyChannel<>() {
            public void send(String m) { sent.incrementAndGet(); }
            public void complete() { completed.set(true); }
            public void completeWithError(ErrorMessage e) { completed.set(true); }
            public void completeWithError(ErrorCategory c, String m) { completed.set(true); }
        };
        FlowControlledReplyChannelWriter<String> writer = new FlowControlledReplyChannelWriter<>(hooked, toAxonServer);
        real.put("response");          // handler result already buffered
        writer.request(32);            // flow-control request from Axon Server, on the gRPC thread
        System.out.printf("responses sent=%d, QueryComplete sent=%s, buffer closed=%s empty=%s%n",
                          sent.get(), completed.get(), real.closed(), real.isEmpty());
        System.out.println(completed.get() ? "NOT reproduced" : "REPRODUCED: completion lost, nothing left to trigger it");
    }
}

Expected behaviour

Once the buffer is closed and drained, the writer sends QueryComplete, whichever thread closed it.

Actual behaviour

The responses are sent, but QueryComplete is not, and nothing is left to trigger it until the next permit grant or a cancel. Nothing fails and nothing is logged. On both 2026.0.0 and 2026.1.0, with the stock writer:

responses sent=1, QueryComplete sent=false, buffer closed=true empty=true
REPRODUCED: completion lost, nothing left to trigger it

The hook forces the interleaving, so this shows the race exists, not how likely it is.

For a single-result query the requester usually already has its response, so it does not hang (see AxonIQ/axoniq-framework#666). It does hang where it waits for the completion: queryMany, a handler returning null or empty (AF5 then sends only an empty-payload marker followed by the completion), or a streaming query. Axon Server does not time out streaming queries, so the query stays active there as well.

Suggested fix

Also loop while a completion is pending, so whichever thread last holds the gate sends it:

    } while ((requestedRef.get() > 0 && !buffer.isEmpty())
            || (buffer.closed() && buffer.isEmpty() && !completed.get()));

This covers both orderings. If the buffer closes before the gate holder evaluates the condition, the holder loops again and completes. If it closes after the holder released the gate, the closing thread's own CAS succeeds and it completes. With this condition applied to the 2026.0.0 and 2026.1.0 sources, the same reproducer gives QueryComplete sent=true on both. It can serve as the regression test.