diff --git a/bridged/src/main/java/dev/ltms/bridged/msg/AmqpReplyInbox.java b/bridged/src/main/java/dev/ltms/bridged/msg/AmqpReplyInbox.java index e586372..5f5846d 100644 --- a/bridged/src/main/java/dev/ltms/bridged/msg/AmqpReplyInbox.java +++ b/bridged/src/main/java/dev/ltms/bridged/msg/AmqpReplyInbox.java @@ -105,7 +105,13 @@ public final class AmqpReplyInbox implements ReplyInbox, AutoCloseable { private final Object publishChannelLock = new Object(); /** In-flight publishes awaiting their confirm, keyed by the publish channel's sequence number. */ private final ConcurrentSkipListMap pendingBySeq = new ConcurrentSkipListMap<>(); - /** The same in-flight publishes, keyed by {@code msgId} — a broker {@code Return} carries no delivery tag. */ + /** + * The same in-flight publishes, keyed by {@code msgId} — a broker {@code Return} carries no delivery + * tag. Assumes {@code msgId} is unique per in-flight publish: a second {@link #publish} for a + * {@code msgId} still awaiting its confirm would overwrite this entry and misdirect + * {@link #onReturn}'s lookup. Not reachable today — {@code MessageService.reply} generates a fresh + * {@code UUID} per call — so no guard is added for it. + */ private final ConcurrentHashMap pendingByMsgId = new ConcurrentHashMap<>(); /** A message pulled off the broker but not yet acked: its delivery-tag plus the port payload. */ @@ -374,15 +380,48 @@ public final class AmqpReplyInbox implements ReplyInbox, AutoCloseable { } } - /** Fail every publish still awaiting its confirm — their sequence numbers are stale after recovery. */ - private void failPendingPublishesOnRecovery() { - for (var it = pendingBySeq.entrySet().iterator(); it.hasNext(); ) { - Pending pending = it.next().getValue(); - it.remove(); - pendingByMsgId.remove(pending.msgId, pending); - pending.confirmed.completeExceptionally(new IllegalStateException( - "AMQP connection recovered mid-publish; confirm status of reply " + pending.msgId - + " is unknown")); + /** + * Fail every publish still awaiting its confirm — their sequence numbers are stale after recovery. + * Guarded by {@link #publishChannelLock}, the same lock {@link #publish} holds while it takes its + * sequence number and registers its {@link Pending}: without it, a {@link #publish} that starts + * after the connection has already recovered (so it publishes — and will be confirmed — on the + * new channel) can register between this sweep's iteration and its clear, and this sweep + * then fails a publish that actually succeeded. {@link #publish} only holds the lock for the + * seq/map-put/{@code basicPublish} — it awaits the confirm outside it — so this sweep can only ever + * wait for an in-flight {@code basicPublish} call to return, never for a broker round trip. No + * deadlock. + * + *

Package-private (rather than {@code private}) only so the unit test can drive it directly + * against a concurrent {@link #publish} without a live broker reconnect. + */ + void failPendingPublishesOnRecovery() { + synchronized (publishChannelLock) { + for (var it = pendingBySeq.entrySet().iterator(); it.hasNext(); ) { + Pending pending = it.next().getValue(); + it.remove(); + pendingByMsgId.remove(pending.msgId, pending); + pending.confirmed.completeExceptionally(new IllegalStateException( + "AMQP connection recovered mid-publish; confirm status of reply " + pending.msgId + + " is unknown")); + } + } + } + + /** + * Fail every publish still awaiting its confirm with a clear, immediate error instead of leaving it + * to time out after {@link #CONFIRM_TIMEOUT_MS} once the channels are closed underneath it. Guarded + * by {@link #publishChannelLock} for the same reason as {@link #failPendingPublishesOnRecovery}. + */ + private void failPendingPublishesOnClose() { + synchronized (publishChannelLock) { + for (var it = pendingBySeq.entrySet().iterator(); it.hasNext(); ) { + Pending pending = it.next().getValue(); + it.remove(); + pendingByMsgId.remove(pending.msgId, pending); + pending.confirmed.completeExceptionally(new IllegalStateException( + "AMQP reply inbox closed while publish of reply " + pending.msgId + + " was still awaiting its confirm")); + } } } @@ -392,6 +431,7 @@ public final class AmqpReplyInbox implements ReplyInbox, AutoCloseable { @Override public void close() { + failPendingPublishesOnClose(); try { channel.close(); } catch (Exception e) { diff --git a/bridged/src/test/java/dev/ltms/bridged/msg/AmqpReplyInboxRecoveryRaceTest.java b/bridged/src/test/java/dev/ltms/bridged/msg/AmqpReplyInboxRecoveryRaceTest.java new file mode 100644 index 0000000..4031ea6 --- /dev/null +++ b/bridged/src/test/java/dev/ltms/bridged/msg/AmqpReplyInboxRecoveryRaceTest.java @@ -0,0 +1,264 @@ +package dev.ltms.bridged.msg; + +import com.rabbitmq.client.AMQP; +import com.rabbitmq.client.Channel; +import com.rabbitmq.client.ConfirmCallback; +import com.rabbitmq.client.Connection; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.Timeout; + +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.Proxy; +import java.util.List; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.concurrent.atomic.AtomicReference; + +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; + +/** + * CB-528 follow-up: {@link AmqpReplyInbox#failPendingPublishesOnRecovery()} must not fail a publish + * that registers concurrently with (and was not yet visible when) the recovery sweep began — that + * would report a publish that actually succeeded as failed, and {@code MessageService.reply} retries + * with a fresh {@code msgId}, so the reply is delivered twice. A live broker reconnect cannot be + * forced reliably, so this drives {@link AmqpReplyInbox#failPendingPublishesOnRecovery} and a real + * {@link AmqpReplyInbox#publish} against each other directly, against fake AMQP channels built with + * {@link Proxy} (no mocking library is on the classpath). + * + *

Also covers Finding 2 (CB-528 follow-up): {@link AmqpReplyInbox#close()} must fail an in-flight + * publish promptly instead of leaving it to idle out the 10s confirm timeout. + */ +class AmqpReplyInboxRecoveryRaceTest { + + /** Large enough that the (unfixed) unsynchronized sweep's iteration is a real, observable window + * a concurrently-started publish can land in — not just a best case, single-entry sprint. */ + private static final int STALE_PUBLISHES = 100_000; + + @Test + @Timeout(30) + void recoverySweepDoesNotFailAPublishThatRegistersWhileItIsRunning() throws Exception { + AtomicLong seqCounter = new AtomicLong(); + List seqOrder = new CopyOnWriteArrayList<>(); + List msgIdOrder = new CopyOnWriteArrayList<>(); + AtomicReference ackCallback = new AtomicReference<>(); + AtomicReference nackCallback = new AtomicReference<>(); + + Channel publishChannel = fakeChannel(seqCounter, seqOrder, msgIdOrder, ackCallback, nackCallback); + Channel consumeChannel = fakeChannel(seqCounter, seqOrder, msgIdOrder, ackCallback, nackCallback); + Connection connection = fakeConnection(consumeChannel, publishChannel); + + AmqpReplyInbox inbox = new AmqpReplyInbox(connection, AmqpReplyInbox.DEFAULT_PREFETCH); + + // STALE_PUBLISHES in-flight publishes that never get confirmed — they sit in pendingBySeq / + // pendingByMsgId exactly like publishes whose confirm never arrived before a connection drop. + // Virtual threads make this many concurrent blocking publish() calls cheap. + CountDownLatch staleStarted = new CountDownLatch(STALE_PUBLISHES); + for (int i = 0; i < STALE_PUBLISHES; i++) { + String msgId = "stale-" + i; + Thread.ofVirtual().start(() -> { + staleStarted.countDown(); + try { + inbox.publish("worker-stale", msgId, "x"); + } catch (IllegalStateException expected) { + // resolved (failed by the sweep) — that is exactly what this thread is here for + } + }); + } + staleStarted.await(); + // Let the registrations (the synchronized put into pendingBySeq/pendingByMsgId) actually land + // for all of them before the sweep starts, so the sweep begins with a large, real backlog. + Thread.sleep(300); + + AtomicReference sweepError = new AtomicReference<>(); + Thread sweepThread = new Thread(() -> { + try { + inbox.failPendingPublishesOnRecovery(); + } catch (Throwable t) { + sweepError.set(t); + } + }, "recovery-sweep"); + sweepThread.start(); + // A short, deliberate head start: with STALE_PUBLISHES this large, the (unfixed) sweep's own + // iteration takes several milliseconds, so this guarantees the sweep has already begun — + // and, once guarded, is already holding publishChannelLock — before "fresh" attempts to + // register. Without this head start, "fresh" sometimes wins the race for the lock and + // registers before the sweep even starts, which is the accepted "already in flight when + // recovery fires" case (correctly failed either way) rather than the bug under test. + Thread.sleep(5); + + // This is the exact interleaving CB-528's follow-up describes: "the still-running recovery + // sweep" racing a publish that registers while it is mid-flight. + AtomicReference publishError = new AtomicReference<>(); + Thread freshThread = new Thread(() -> { + try { + inbox.publish("worker-fresh", "fresh", "hello"); + } catch (Throwable t) { + publishError.set(t); + } + }, "fresh-publish"); + freshThread.start(); + + sweepThread.join(20_000); + assertNull(sweepError.get(), "sweep threw: " + sweepError.get()); + + // Simulate the broker's real confirm for "fresh" now that the sweep is done, so a correct + // implementation's publish() returns normally instead of idling out CONFIRM_TIMEOUT_MS. Poll + // for the registration rather than checking once: freshThread may still be contending for + // publishChannelLock (behind the 20,000 stale threads' own lock acquisitions) even though the + // sweep itself has already finished. + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(9); + int idx = -1; + while (idx < 0 && System.nanoTime() < deadline) { + idx = msgIdOrder.indexOf("fresh"); + if (idx < 0) { + Thread.sleep(20); + } + } + if (idx >= 0 && ackCallback.get() != null) { + ackCallback.get().handle(seqOrder.get(idx), false); + } + freshThread.join(15_000); + + assertNull(publishError.get(), + "a publish that registered while the recovery sweep was running must not be failed by " + + "it, but got: " + publishError.get()); + } + + @Test + @Timeout(15) + void closeFailsInFlightPublishPromptlyInsteadOfWaitingOutTheConfirmTimeout() throws Exception { + AtomicLong seqCounter = new AtomicLong(); + List seqOrder = new CopyOnWriteArrayList<>(); + List msgIdOrder = new CopyOnWriteArrayList<>(); + AtomicReference ackCallback = new AtomicReference<>(); + AtomicReference nackCallback = new AtomicReference<>(); + + Channel publishChannel = fakeChannel(seqCounter, seqOrder, msgIdOrder, ackCallback, nackCallback); + Channel consumeChannel = fakeChannel(seqCounter, seqOrder, msgIdOrder, ackCallback, nackCallback); + Connection connection = fakeConnection(consumeChannel, publishChannel); + + AmqpReplyInbox inbox = new AmqpReplyInbox(connection, AmqpReplyInbox.DEFAULT_PREFETCH); + + CountDownLatch publishReturned = new CountDownLatch(1); + AtomicReference publishError = new AtomicReference<>(); + AtomicLong elapsedMillis = new AtomicLong(); + Thread publishThread = new Thread(() -> { + long start = System.nanoTime(); + try { + inbox.publish("worker-close", "never-confirmed", "x"); + } catch (Throwable t) { + publishError.set(t); + } finally { + elapsedMillis.set((System.nanoTime() - start) / 1_000_000); + publishReturned.countDown(); + } + }, "publish-during-close"); + publishThread.start(); + + Thread.sleep(200); // let publish() register before close() runs + inbox.close(); + + assertTrue(publishReturned.await(5, TimeUnit.SECONDS), "publish() did not return after close()"); + assertNotNull(publishError.get(), "a publish in flight when close() runs must fail, not hang"); + assertTrue(publishError.get().getMessage() != null + && publishError.get().getMessage().toLowerCase().contains("closed"), + "expected a clear closed-inbox message, got: " + publishError.get()); + assertTrue(elapsedMillis.get() < 5_000, + "close() should fail the in-flight publish promptly, not wait out the confirm timeout — took " + + elapsedMillis.get() + "ms"); + } + + /** A {@link Proxy}-backed {@link Channel}: only the calls {@link AmqpReplyInbox} actually makes + * are meaningfully implemented; everything else returns a harmless default. */ + private static Channel fakeChannel(AtomicLong seqCounter, List seqOrder, List msgIdOrder, + AtomicReference ackCallback, + AtomicReference nackCallback) { + InvocationHandler handler = (proxy, method, args) -> { + String name = method.getName(); + if (name.equals("getNextPublishSeqNo")) { + long value = seqCounter.incrementAndGet(); + seqOrder.add(value); + return value; + } + if (name.equals("basicPublish")) { + AMQP.BasicProperties props = (AMQP.BasicProperties) args[3]; + msgIdOrder.add(props.getMessageId()); + return null; + } + if (name.equals("addConfirmListener")) { + ackCallback.set((ConfirmCallback) args[0]); + nackCallback.set((ConfirmCallback) args[1]); + return null; + } + if (name.equals("equals")) { + return proxy == args[0]; + } + if (name.equals("hashCode")) { + return System.identityHashCode(proxy); + } + if (name.equals("toString")) { + return "FakeChannel"; + } + return defaultValue(method.getReturnType()); + }; + return (Channel) Proxy.newProxyInstance(AmqpReplyInboxRecoveryRaceTest.class.getClassLoader(), + new Class[] {Channel.class}, handler); + } + + /** A {@link Proxy}-backed {@link Connection} handing out {@code first} then {@code second} from + * successive {@code createChannel()} calls, matching {@link AmqpReplyInbox}'s constructor. */ + private static Connection fakeConnection(Channel first, Channel second) { + AtomicInteger calls = new AtomicInteger(); + InvocationHandler handler = (proxy, method, args) -> { + String name = method.getName(); + if (name.equals("createChannel") && (args == null || args.length == 0)) { + return calls.getAndIncrement() == 0 ? first : second; + } + if (name.equals("equals")) { + return proxy == args[0]; + } + if (name.equals("hashCode")) { + return System.identityHashCode(proxy); + } + if (name.equals("toString")) { + return "FakeConnection"; + } + return defaultValue(method.getReturnType()); + }; + return (Connection) Proxy.newProxyInstance(AmqpReplyInboxRecoveryRaceTest.class.getClassLoader(), + new Class[] {Connection.class}, handler); + } + + private static Object defaultValue(Class type) { + if (!type.isPrimitive() || type == void.class) { + return null; + } + if (type == boolean.class) { + return Boolean.FALSE; + } + if (type == long.class) { + return 0L; + } + if (type == short.class) { + return (short) 0; + } + if (type == byte.class) { + return (byte) 0; + } + if (type == char.class) { + return (char) 0; + } + if (type == double.class) { + return 0.0d; + } + if (type == float.class) { + return 0.0f; + } + return 0; + } +}