CB-528: close the recovery race in AmqpReplyInbox
failPendingPublishesOnRecovery walked and cleared pendingBySeq/pendingByMsgId without holding publishChannelLock, so a publish() that registered while the sweep was still iterating could be failed even though it published on the already-recovered channel — a successful publish reported as failed, and since MessageService.reply mints a fresh msgId per retry, dedup can't catch the resulting duplicate. Guard the sweep with publishChannelLock: publish() only holds it for the seq/map-put/basicPublish, so the sweep can only ever wait for an in-flight basicPublish to return, never a broker round trip. Also make close() fail in-flight publishes immediately with a clear message instead of leaving them to idle out the 10s confirm timeout, and record the (currently unreachable) msgId-uniqueness assumption pendingByMsgId relies on. AmqpReplyInboxRecoveryRaceTest drives the sweep and a real publish() against each other directly (no live broker reconnect) using Proxy-backed fake AMQP channels and a large in-flight backlog to make the race window observable; confirmed it fails without the guard (reply "fresh" wrongly failed as "connection recovered mid-publish") and passes with it.
This commit is contained in:
@@ -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<Long, Pending> 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<String, Pending> 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
|
||||
* <em>new</em> 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.
|
||||
*
|
||||
* <p>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) {
|
||||
|
||||
@@ -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).
|
||||
*
|
||||
* <p>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<Long> seqOrder = new CopyOnWriteArrayList<>();
|
||||
List<String> msgIdOrder = new CopyOnWriteArrayList<>();
|
||||
AtomicReference<ConfirmCallback> ackCallback = new AtomicReference<>();
|
||||
AtomicReference<ConfirmCallback> 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<Throwable> 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<Throwable> 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<Long> seqOrder = new CopyOnWriteArrayList<>();
|
||||
List<String> msgIdOrder = new CopyOnWriteArrayList<>();
|
||||
AtomicReference<ConfirmCallback> ackCallback = new AtomicReference<>();
|
||||
AtomicReference<ConfirmCallback> 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<Throwable> 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<Long> seqOrder, List<String> msgIdOrder,
|
||||
AtomicReference<ConfirmCallback> ackCallback,
|
||||
AtomicReference<ConfirmCallback> 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;
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user