Deliver the envelope in hand to the last receive handler before it is removed (#2153)

removeReceiveHandler removed the handler before it stopped the receive
thread. An envelope the thread was handling then was already cached and
acknowledged to the server; it reached no handler, and the cache dropped it
once handled, so it was lost. While removing the last handler stops the
thread, the handler stays in drainingHandlers: it still gets that envelope,
but doesn't keep the thread receiving.

Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

lib/src/main/java/org/asamk/signal/manager/internal/ManagerImpl.java

@@ -174,6 +174,9 @@ public class ManagerImpl implements Manager {
     private boolean isReceivingSynchronous;
     private final Set<ReceiveMessageHandler> weakHandlers = new HashSet<>();
     private final Set<ReceiveMessageHandler> messageHandlers = new HashSet<>();
+    // The last handler while its removal stops the receive thread: it still gets the envelope the thread is
+    // handling, but doesn't keep the thread receiving.
+    private final Set<ReceiveMessageHandler> drainingHandlers = new HashSet<>();
     private final List<Runnable> closedListeners = new ArrayList<>();
     private final List<Runnable> addressChangedListeners = new ArrayList<>();
     private final CompositeDisposable disposable = new CompositeDisposable();
@@ -1541,7 +1544,7 @@ public class ManagerImpl implements Manager {
 
     private void passReceivedMessageToHandlers(MessageEnvelope envelope, Throwable e) {
         synchronized (messageHandlers) {
-            Stream.concat(messageHandlers.stream(), weakHandlers.stream()).forEach(h -> {
+            Stream.of(messageHandlers, weakHandlers, drainingHandlers).flatMap(Set::stream).forEach(h -> {
                 try {
                     h.handleMessage(envelope, e);
                 } catch (Throwable ex) {
@@ -1556,14 +1559,26 @@ public class ManagerImpl implements Manager {
         final Thread thread;
         synchronized (messageHandlers) {
             weakHandlers.remove(handler);
-            messageHandlers.remove(handler);
+            final var removed = messageHandlers.remove(handler);
             if (!messageHandlers.isEmpty() || receiveThread == null || isReceivingSynchronous) {
                 return;
             }
             thread = receiveThread;
+            // The receive thread may be handling an envelope that is already cached and acknowledged to the server,
+            // and the cache drops it once it's handled. So the last handler keeps getting envelopes until the thread
+            // has stopped, otherwise that envelope would reach no handler and be lost.
+            if (removed) {
+                drainingHandlers.add(handler);
+            }
         }
 
-        stopReceiveThread(thread);
+        try {
+            stopReceiveThread(thread);
+        } finally {
+            synchronized (messageHandlers) {
+                drainingHandlers.remove(handler);
+            }
+        }
     }
 
     private void stopReceiveThread(final Thread thread) {

lib/src/test/java/org/asamk/signal/manager/internal/ReceiveThreadTest.java

@@ -23,8 +23,10 @@ import java.nio.file.Path;
 import java.security.Security;
 import java.time.Duration;
 import java.util.Optional;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.function.BooleanSupplier;
 
+import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
 @Timeout(60)
@@ -142,6 +144,28 @@ class ReceiveThreadTest {
         assertTrue(receiveHelper.await(() -> receiveHelper.startedThreads == 2 && receiveHelper.maxRunningLoops == 1));
     }
 
+    @Test
+    void lastHandlerStillGetsTheMessageInHandWhenItIsRemoved() throws Exception {
+        final var received = new AtomicInteger();
+        final var handler = new Manager.ReceiveMessageHandler() {
+            @Override
+            public void handleMessage(final MessageEnvelope envelope, final Throwable e) {
+                received.incrementAndGet();
+            }
+        };
+
+        manager.addReceiveHandler(handler);
+        assertTrue(receiveHelper.await(() -> receiveHelper.runningLoops == 1));
+
+        // Like an envelope that is already cached and acknowledged, and still being handled when the last
+        // handler goes: the loop passes it on before it stops
+        receiveHelper.setPassOnOneWhenStopped(true);
+        manager.removeReceiveHandler(handler);
+
+        assertEquals(1, received.get());
+        assertTrue(receiveHelper.await(() -> receiveHelper.runningLoops == 0 && receiveHelper.startedThreads == 1));
+    }
+
     private static Manager.ReceiveMessageHandler newHandler() {
         // An anonymous class instead of a lambda, so every handler is a separate instance
         return new Manager.ReceiveMessageHandler() {
@@ -161,6 +185,7 @@ class ReceiveThreadTest {
         private final Object lock = new Object();
         private boolean stopRequested;
         private boolean holdStoppingLoops;
+        private boolean passOnOneWhenStopped;
         private int stopRequests;
         private int startedThreads;
         private int startedLoops;
@@ -213,6 +238,9 @@ class ReceiveThreadTest {
                 maxRunningLoops = Math.max(maxRunningLoops, runningLoops);
                 lock.notifyAll();
                 if (await(() -> stopRequested)) {
+                    if (passOnOneWhenStopped) {
+                        handler.handleMessage(null, null);
+                    }
                     stoppingLoops++;
                     lock.notifyAll();
                     await(() -> !holdStoppingLoops);
@@ -223,6 +251,12 @@ class ReceiveThreadTest {
             }
         }
 
+        void setPassOnOneWhenStopped(final boolean passOn) {
+            synchronized (lock) {
+                passOnOneWhenStopped = passOn;
+            }
+        }
+
         void setHoldStoppingLoops(final boolean hold) {
             synchronized (lock) {
                 holdStoppingLoops = hold;