diff --git a/src/main/java/io/axoniq/axonserver/connector/impl/FlowControlledReplyChannelWriter.java b/src/main/java/io/axoniq/axonserver/connector/impl/FlowControlledReplyChannelWriter.java index 9a7f3f64..ba24dc74 100644 --- a/src/main/java/io/axoniq/axonserver/connector/impl/FlowControlledReplyChannelWriter.java +++ b/src/main/java/io/axoniq/axonserver/connector/impl/FlowControlledReplyChannelWriter.java @@ -103,7 +103,10 @@ private void stream() { } finally { flowGate.set(false); } - } while (requestedRef.get() > 0 && !buffer.isEmpty()); + // Also loop while a completion is pending: a concurrent close() loses the flowGate CAS and returns, + // so the current gate holder is the only one left to send the completion. + } while ((requestedRef.get() > 0 && !buffer.isEmpty()) + || (buffer.closed() && buffer.isEmpty() && !completed.get())); } private long send(long maxMessages) { diff --git a/src/test/java/io/axoniq/axonserver/connector/impl/FlowControlledReplyChannelWriterTest.java b/src/test/java/io/axoniq/axonserver/connector/impl/FlowControlledReplyChannelWriterTest.java index 6e1f6455..12a4c5bf 100644 --- a/src/test/java/io/axoniq/axonserver/connector/impl/FlowControlledReplyChannelWriterTest.java +++ b/src/test/java/io/axoniq/axonserver/connector/impl/FlowControlledReplyChannelWriterTest.java @@ -29,6 +29,7 @@ import java.util.concurrent.ScheduledFuture; import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; +import java.util.concurrent.atomic.AtomicBoolean; import java.util.concurrent.atomic.AtomicReference; import java.util.stream.LongStream; @@ -160,9 +161,101 @@ void testRequestWhenSourceCompleted() { verify(destination).complete(); } + @Test + void completionIsSentWhenBufferClosesWhileAnotherThreadHoldsTheFlowGate() { + BlockingCloseableBuffer real = new BlockingCloseableBuffer<>(); + DisposableReadonlyBuffer hooked = closingDuringClosedCheck(real, real::close); + writer = new FlowControlledReplyChannelWriter<>(hooked, destination); + real.put("response"); + + writer.request(32); + + verify(destination, times(1)).send("response"); + verify(destination, times(1)).complete(); + verify(destination, never()).completeWithError(any(ErrorMessage.class)); + } + + @Test + void errorCompletionIsSentWhenBufferClosesExceptionallyWhileAnotherThreadHoldsTheFlowGate() { + BlockingCloseableBuffer real = new BlockingCloseableBuffer<>(); + ErrorMessage error = ErrorMessage.newBuilder().setMessage("failure").build(); + DisposableReadonlyBuffer hooked = + closingDuringClosedCheck(real, () -> real.closeExceptionally(error)); + writer = new FlowControlledReplyChannelWriter<>(hooked, destination); + real.put("response"); + + writer.request(32); + + verify(destination, times(1)).send("response"); + verify(destination, times(1)).completeWithError(error); + verify(destination, never()).complete(); + } + @Test void testCancelPropagation() { writer.cancel(); verify(source).dispose(); } + + /** + * Wraps the given {@code real} buffer so that the first {@code closed()} check after the buffer has been drained + * triggers the given {@code closer} on a second thread, whilst still returning the value from before the close. + * This forces the interleaving in which the closing thread loses the flow gate CAS in the writer. + */ + private static DisposableReadonlyBuffer closingDuringClosedCheck(BlockingCloseableBuffer real, + Runnable closer) { + AtomicBoolean drained = new AtomicBoolean(); + AtomicBoolean injected = new AtomicBoolean(); + return new DisposableReadonlyBuffer() { + @Override + public Optional poll() { + Optional result = real.poll(); + if (result.isPresent()) { + drained.set(true); + } + return result; + } + + @Override + public boolean isEmpty() { + return real.isEmpty(); + } + + @Override + public int capacity() { + return real.capacity(); + } + + @Override + public void onAvailable(Runnable onAvailable) { + real.onAvailable(onAvailable); + } + + @Override + public void dispose() { + // noop + } + + @Override + public Optional error() { + return real.error(); + } + + @Override + public boolean closed() { + boolean closed = real.closed(); + if (drained.get() && !closed && injected.compareAndSet(false, true)) { + Thread producer = new Thread(closer, "producer"); + producer.start(); + try { + producer.join(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + throw new IllegalStateException(e); + } + } + return closed; + } + }; + } }