fleetd/src/test/java/dev/ltms/fleet/msg/AmqpReplyInboxRecoveryRaceTest.java

@@ -8,7 +8,10 @@ import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
 
 import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Method;
 import java.lang.reflect.Proxy;
+import java.io.IOException;
+import java.util.Map;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.CountDownLatch;
@@ -20,6 +23,7 @@ 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;
+import static org.junit.jupiter.api.Assertions.assertEquals;
 
 /**
  * CB-528 follow-up: {@link AmqpReplyInbox#failPendingPublishesOnRecovery()} must not fail a publish
@@ -197,11 +201,78 @@ class AmqpReplyInboxRecoveryRaceTest {
                         + elapsedMillis.get() + "ms");
     }
 
+    @Test
+    void publishIOExceptionRemovesThePendingMessageId() throws Exception {
+        AtomicLong seqCounter = new AtomicLong();
+        Channel failing = fakeChannel(seqCounter, new CopyOnWriteArrayList<>(), new CopyOnWriteArrayList<>(),
+                new AtomicReference<>(), new AtomicReference<>(), true);
+        AmqpReplyInbox inbox = new AmqpReplyInbox(fakeConnection(failing, failing), AmqpReplyInbox.DEFAULT_PREFETCH);
+
+        try {
+            org.junit.jupiter.api.Assertions.assertThrows(IllegalStateException.class,
+                    () -> inbox.publish("worker", "catch", "body"));
+            assertEquals(0, pendingByMsgId(inbox).size(), "publish IOException must remove its msgId entry");
+        } finally {
+            inbox.close();
+        }
+    }
+
+    @Test
+    void interruptedPublishRemovesThePendingMessageIdInFinally() throws Exception {
+        InboxFixture fixture = new InboxFixture();
+        Thread publish = fixture.startPublish("finally");
+        fixture.awaitPublished("finally");
+        publish.interrupt();
+        publish.join(5_000);
+
+        assertEquals(0, pendingByMsgId(fixture.inbox).size(), "publish finally must remove its msgId entry");
+        fixture.inbox.close();
+    }
+
+    @Test
+    void confirmResolutionRemovesThePendingMessageId() throws Exception {
+        AmqpReplyInbox inbox = new InboxFixture().inbox;
+        try {
+            seedPending(inbox, 1, "confirm");
+            invoke(inbox, "resolveConfirm", new Class<?>[] {long.class, boolean.class, boolean.class}, 1L, false, true);
+            assertEquals(0, pendingByMsgId(inbox).size(), "confirm resolution must remove its msgId entry");
+        } finally {
+            inbox.close();
+        }
+    }
+
+    @Test
+    void recoverySweepRemovesThePendingMessageId() throws Exception {
+        AmqpReplyInbox inbox = new InboxFixture().inbox;
+        try {
+            seedPending(inbox, 1, "recovery");
+            inbox.failPendingPublishesOnRecovery();
+            assertEquals(0, pendingByMsgId(inbox).size(), "recovery sweep must remove its msgId entry");
+        } finally {
+            inbox.close();
+        }
+    }
+
+    @Test
+    void closeRemovesThePendingMessageId() throws Exception {
+        AmqpReplyInbox inbox = new InboxFixture().inbox;
+        seedPending(inbox, 1, "close");
+        inbox.close();
+
+        assertEquals(0, pendingByMsgId(inbox).size(), "close must remove its msgId entry");
+    }
+
     /** 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) {
+                                         AtomicReference<ConfirmCallback> ackCallback,
+                                         AtomicReference<ConfirmCallback> nackCallback) {
+        return fakeChannel(seqCounter, seqOrder, msgIdOrder, ackCallback, nackCallback, false);
+    }
+
+    private static Channel fakeChannel(AtomicLong seqCounter, List<Long> seqOrder, List<String> msgIdOrder,
+                                         AtomicReference<ConfirmCallback> ackCallback,
+                                         AtomicReference<ConfirmCallback> nackCallback, boolean failPublish) {
         InvocationHandler handler = (proxy, method, args) -> {
             String name = method.getName();
             if (name.equals("getNextPublishSeqNo")) {
@@ -210,6 +281,9 @@ class AmqpReplyInboxRecoveryRaceTest {
                 return value;
             }
             if (name.equals("basicPublish")) {
+                if (failPublish) {
+                    throw new IOException("test publish failure");
+                }
                 AMQP.BasicProperties props = (AMQP.BasicProperties) args[3];
                 msgIdOrder.add(props.getMessageId());
                 return null;
@@ -285,4 +359,60 @@ class AmqpReplyInboxRecoveryRaceTest {
         }
         return 0;
     }
+
+    @SuppressWarnings("unchecked")
+    private static Map<String, Object> pendingByMsgId(AmqpReplyInbox inbox) throws Exception {
+        var field = AmqpReplyInbox.class.getDeclaredField("pendingByMsgId");
+        field.setAccessible(true);
+        return (Map<String, Object>) field.get(inbox);
+    }
+
+    @SuppressWarnings("unchecked")
+    private static void seedPending(AmqpReplyInbox inbox, long seq, String msgId) throws Exception {
+        Class<?> pendingType = Class.forName(AmqpReplyInbox.class.getName() + "$Pending");
+        var constructor = pendingType.getDeclaredConstructor(String.class);
+        constructor.setAccessible(true);
+        Object pending = constructor.newInstance(msgId);
+        var seqField = AmqpReplyInbox.class.getDeclaredField("pendingBySeq");
+        seqField.setAccessible(true);
+        ((Map<Long, Object>) seqField.get(inbox)).put(seq, pending);
+        pendingByMsgId(inbox).put(msgId, pending);
+    }
+
+    private static void invoke(AmqpReplyInbox inbox, String name, Class<?>[] types, Object... args) throws Exception {
+        Method method = AmqpReplyInbox.class.getDeclaredMethod(name, types);
+        method.setAccessible(true);
+        method.invoke(inbox, args);
+    }
+
+    private static final class InboxFixture {
+        final AtomicLong seqCounter = new AtomicLong();
+        final List<Long> seqOrder = new CopyOnWriteArrayList<>();
+        final List<String> msgIdOrder = new CopyOnWriteArrayList<>();
+        final AtomicReference<ConfirmCallback> ackCallback = new AtomicReference<>();
+        final AtomicReference<ConfirmCallback> nackCallback = new AtomicReference<>();
+        final AmqpReplyInbox inbox = new AmqpReplyInbox(
+                fakeConnection(fakeChannel(seqCounter, seqOrder, msgIdOrder, ackCallback, nackCallback),
+                        fakeChannel(seqCounter, seqOrder, msgIdOrder, ackCallback, nackCallback)),
+                AmqpReplyInbox.DEFAULT_PREFETCH);
+
+        Thread startPublish(String msgId) {
+            Thread thread = Thread.ofVirtual().start(() -> {
+                try {
+                    inbox.publish("worker", msgId, "body");
+                } catch (IllegalStateException ignored) {
+                    // Interrupting the confirm wait is the path under test.
+                }
+            });
+            return thread;
+        }
+
+        void awaitPublished(String msgId) throws InterruptedException {
+            long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+            while (!msgIdOrder.contains(msgId) && System.nanoTime() < deadline) {
+                Thread.sleep(10);
+            }
+            assertTrue(msgIdOrder.contains(msgId), "publish did not register " + msgId);
+        }
+    }
 }

fleetd/src/test/java/dev/ltms/fleet/msg/LeadMailboxTest.java

@@ -12,7 +12,11 @@ import org.testcontainers.junit.jupiter.Testcontainers;
 import org.testcontainers.utility.DockerImageName;
 
 import java.io.IOException;
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.Method;
+import java.lang.reflect.Proxy;
 import java.util.List;
+import java.util.Map;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -373,6 +377,72 @@ class LeadMailboxTest {
                 () -> "expected AlreadyClosedException, got: " + thrown);
     }
 
+    @Test
+    void publishIOExceptionRemovesThePendingMessageId() throws Exception {
+        LeadMailbox mailbox = newMailbox(true);
+        try {
+            assertThrows(IllegalStateException.class,
+                    () -> mailbox.publish("target", new LeadMessage("catch", "from", "target", "body")));
+            assertEquals(0, pendingByMsgId(mailbox).size(), "publish IOException must remove its msgId entry");
+        } finally {
+            mailbox.close();
+        }
+    }
+
+    @Test
+    void interruptedPublishRemovesThePendingMessageIdInFinally() throws Exception {
+        LeadMailbox mailbox = newMailbox(false);
+        Thread publish = Thread.ofVirtual().start(() -> {
+            try {
+                mailbox.publish("target", new LeadMessage("finally", "from", "target", "body"));
+            } catch (IllegalStateException ignored) {
+                // Interrupting the confirm wait is the path under test.
+            }
+        });
+        awaitPending(mailbox, "finally");
+        publish.interrupt();
+        publish.join(5_000);
+
+        try {
+            assertEquals(0, pendingByMsgId(mailbox).size(), "publish finally must remove its msgId entry");
+        } finally {
+            mailbox.close();
+        }
+    }
+
+    @Test
+    void confirmResolutionRemovesThePendingMessageId() throws Exception {
+        LeadMailbox mailbox = newMailbox(false);
+        try {
+            seedPending(mailbox, 1, "confirm");
+            invoke(mailbox, "resolveConfirm", new Class<?>[] {long.class, boolean.class, boolean.class}, 1L, false, true);
+            assertEquals(0, pendingByMsgId(mailbox).size(), "confirm resolution must remove its msgId entry");
+        } finally {
+            mailbox.close();
+        }
+    }
+
+    @Test
+    void recoverySweepRemovesThePendingMessageId() throws Exception {
+        LeadMailbox mailbox = newMailbox(false);
+        try {
+            seedPending(mailbox, 1, "recovery");
+            mailbox.failPendingPublishesOnRecovery();
+            assertEquals(0, pendingByMsgId(mailbox).size(), "recovery sweep must remove its msgId entry");
+        } finally {
+            mailbox.close();
+        }
+    }
+
+    @Test
+    void closeRemovesThePendingMessageId() throws Exception {
+        LeadMailbox mailbox = newMailbox(false);
+        seedPending(mailbox, 1, "close");
+        mailbox.close();
+
+        assertEquals(0, pendingByMsgId(mailbox).size(), "close must remove its msgId entry");
+    }
+
     /** Poll peek until at least one message is held, or ~10s elapse (broker delivery is async). */
     @SuppressWarnings("BusyWait")
     private static List<LeadMessage> awaitPeek(LeadMailbox inbox) throws InterruptedException {
@@ -398,4 +468,108 @@ class LeadMailboxTest {
         }
         return state;
     }
+
+    private static LeadMailbox newMailbox(boolean failPublish) {
+        AtomicLong sequence = new AtomicLong();
+        Channel consume = fakeChannel(sequence, false);
+        Channel publish = fakeChannel(sequence, failPublish);
+        return new LeadMailbox(fakeConnection(consume, publish), "self");
+    }
+
+    private static Channel fakeChannel(AtomicLong sequence, boolean failPublish) {
+        InvocationHandler handler = (proxy, method, args) -> {
+            if (method.getName().equals("getNextPublishSeqNo")) {
+                return sequence.incrementAndGet();
+            }
+            if (method.getName().equals("basicPublish") && failPublish) {
+                throw new IOException("test publish failure");
+            }
+            if (method.getName().equals("equals")) {
+                return proxy == args[0];
+            }
+            if (method.getName().equals("hashCode")) {
+                return System.identityHashCode(proxy);
+            }
+            return defaultValue(method.getReturnType());
+        };
+        return (Channel) Proxy.newProxyInstance(LeadMailboxTest.class.getClassLoader(), new Class<?>[] {Channel.class}, handler);
+    }
+
+    private static Connection fakeConnection(Channel first, Channel second) {
+        AtomicLong calls = new AtomicLong();
+        InvocationHandler handler = (proxy, method, args) -> {
+            if (method.getName().equals("createChannel") && (args == null || args.length == 0)) {
+                return calls.getAndIncrement() == 0 ? first : second;
+            }
+            if (method.getName().equals("equals")) {
+                return proxy == args[0];
+            }
+            if (method.getName().equals("hashCode")) {
+                return System.identityHashCode(proxy);
+            }
+            return defaultValue(method.getReturnType());
+        };
+        return (Connection) Proxy.newProxyInstance(LeadMailboxTest.class.getClassLoader(), new Class<?>[] {Connection.class}, handler);
+    }
+
+    @SuppressWarnings("unchecked")
+    private static Map<String, Object> pendingByMsgId(LeadMailbox mailbox) throws Exception {
+        var field = LeadMailbox.class.getDeclaredField("pendingByMsgId");
+        field.setAccessible(true);
+        return (Map<String, Object>) field.get(mailbox);
+    }
+
+    @SuppressWarnings("unchecked")
+    private static void seedPending(LeadMailbox mailbox, long seq, String msgId) throws Exception {
+        Class<?> pendingType = Class.forName(LeadMailbox.class.getName() + "$Pending");
+        var constructor = pendingType.getDeclaredConstructor(String.class);
+        constructor.setAccessible(true);
+        Object pending = constructor.newInstance(msgId);
+        var seqField = LeadMailbox.class.getDeclaredField("pendingBySeq");
+        seqField.setAccessible(true);
+        ((Map<Long, Object>) seqField.get(mailbox)).put(seq, pending);
+        pendingByMsgId(mailbox).put(msgId, pending);
+    }
+
+    private static void invoke(LeadMailbox mailbox, String name, Class<?>[] types, Object... args) throws Exception {
+        Method method = LeadMailbox.class.getDeclaredMethod(name, types);
+        method.setAccessible(true);
+        method.invoke(mailbox, args);
+    }
+
+    private static void awaitPending(LeadMailbox mailbox, String msgId) throws Exception {
+        long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(5);
+        while (!pendingByMsgId(mailbox).containsKey(msgId) && System.nanoTime() < deadline) {
+            Thread.sleep(10);
+        }
+        assertTrue(pendingByMsgId(mailbox).containsKey(msgId), "publish did not register " + msgId);
+    }
+
+    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;
+    }
 }