Add optional message limit for receive command

lib/src/main/java/org/asamk/signal/manager/Manager.java

@@ -202,12 +202,9 @@ public interface Manager extends Closeable {
     /**
      * Receive new messages from server, returns if no new message arrive in a timespan of timeout.
      */
-    void receiveMessages(Duration timeout, ReceiveMessageHandler handler) throws IOException;
+    public void receiveMessages(
-
+            Optional<Duration> timeout, Optional<Integer> maxMessages, ReceiveMessageHandler handler
-    /**
+    ) throws IOException;
-     * Receive new messages from server, returns only if the thread is interrupted.
-     */
-    void receiveMessages(ReceiveMessageHandler handler) throws IOException;
 
     void setReceiveConfig(ReceiveConfig receiveConfig);
 

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

@@ -961,17 +961,16 @@ class ManagerImpl implements Manager {
     }
 
     @Override
-    public void receiveMessages(Duration timeout, ReceiveMessageHandler handler) throws IOException {
+    public void receiveMessages(
-        receiveMessages(timeout, true, handler);
+            Optional<Duration> timeout,
-    }
+            Optional<Integer> maxMessages,
-
+            ReceiveMessageHandler handler
-    @Override
+    ) throws IOException {
-    public void receiveMessages(ReceiveMessageHandler handler) throws IOException {
+        receiveMessages(timeout.orElse(Duration.ofMinutes(1)), timeout.isPresent(), maxMessages.orElse(null), handler);
-        receiveMessages(Duration.ofMinutes(1), false, handler);
     }
 
     private void receiveMessages(
-            Duration timeout, boolean returnOnTimeout, ReceiveMessageHandler handler
+            Duration timeout, boolean returnOnTimeout, Integer maxMessages, ReceiveMessageHandler handler
     ) throws IOException {
         if (isReceiving()) {
             throw new IllegalStateException("Already receiving message.");
@@ -979,7 +978,7 @@ class ManagerImpl implements Manager {
         isReceivingSynchronous = true;
         receiveThread = Thread.currentThread();
         try {
-            context.getReceiveHelper().receiveMessages(timeout, returnOnTimeout, handler);
+            context.getReceiveHelper().receiveMessages(timeout, returnOnTimeout, maxMessages, handler);
         } finally {
             receiveThread = null;
             isReceivingSynchronous = false;

lib/src/main/java/org/asamk/signal/manager/helper/ReceiveHelper.java

@@ -80,7 +80,7 @@ public class ReceiveHelper {
     public void receiveMessagesContinuously(Manager.ReceiveMessageHandler handler) {
         while (!shouldStop) {
             try {
-                receiveMessages(Duration.ofMinutes(1), false, handler);
+                receiveMessages(Duration.ofMinutes(1), false, null, handler);
                 break;
             } catch (IOException e) {
                 logger.warn("Receiving messages failed, retrying", e);
@@ -89,7 +89,7 @@ public class ReceiveHelper {
     }
 
     public void receiveMessages(
-            Duration timeout, boolean returnOnTimeout, Manager.ReceiveMessageHandler handler
+            Duration timeout, boolean returnOnTimeout, Integer maxMessages, Manager.ReceiveMessageHandler handler
     ) throws IOException {
         needsToRetryFailedMessages = true;
         hasCaughtUpWithOldMessages = false;
@@ -107,7 +107,7 @@ public class ReceiveHelper {
         signalWebSocket.connect();
 
         try {
-            receiveMessagesInternal(signalWebSocket, timeout, returnOnTimeout, handler, queuedActions);
+            receiveMessagesInternal(signalWebSocket, timeout, returnOnTimeout, maxMessages, handler, queuedActions);
         } finally {
             hasCaughtUpWithOldMessages = false;
             handleQueuedActions(queuedActions.keySet());
@@ -122,13 +122,15 @@ public class ReceiveHelper {
             final SignalWebSocket signalWebSocket,
             Duration timeout,
             boolean returnOnTimeout,
+            Integer maxMessages,
             Manager.ReceiveMessageHandler handler,
             final Map<HandleAction, HandleAction> queuedActions
     ) throws IOException {
+        int remainingMessages = maxMessages == null ? -1 : maxMessages;
         var backOffCounter = 0;
         isWaitingForMessage = false;
 
-        while (!shouldStop) {
+        while (!shouldStop && remainingMessages != 0) {
             if (needsToRetryFailedMessages) {
                 retryFailedReceivedMessages(handler);
                 needsToRetryFailedMessages = false;
@@ -154,6 +156,9 @@ public class ReceiveHelper {
                 backOffCounter = 0;
 
                 if (result.isPresent()) {
+                    if (remainingMessages > 0) {
+                        remainingMessages -= 1;
+                    }
                     envelope = result.get();
                     logger.debug("New message received from server");
                 } else {

man/signal-cli.1.adoc

@@ -372,6 +372,9 @@ In json mode this is outputted as one json object per line.
 Number of seconds to wait for new messages (negative values disable timeout).
 Default is 5 seconds.
 
+*--max-messages*::
+Maximum number of messages to receive, before returning.
+
 *--ignore-attachments*::
 Don’t download attachments of received messages.
 

src/main/java/org/asamk/signal/commands/ReceiveCommand.java

@@ -20,6 +20,7 @@ import org.slf4j.LoggerFactory;
 import java.io.IOException;
 import java.time.Duration;
 import java.util.List;
+import java.util.Optional;
 
 public class ReceiveCommand implements LocalCommand {
 
@@ -37,6 +38,10 @@ public class ReceiveCommand implements LocalCommand {
                 .type(double.class)
                 .setDefault(3.0)
                 .help("Number of seconds to wait for new messages (negative values disable timeout)");
+        subparser.addArgument("--max-messages")
+                .type(int.class)
+                .setDefault(-1)
+                .help("Maximum number of messages to receive, before returning.");
         subparser.addArgument("--ignore-attachments")
                 .help("Don’t download attachments of received messages.")
                 .action(Arguments.storeTrue());
@@ -58,6 +63,7 @@ public class ReceiveCommand implements LocalCommand {
             final Namespace ns, final Manager m, final OutputWriter outputWriter
     ) throws CommandException {
         final var timeout = ns.getDouble("timeout");
+        final var maxMessagesRaw = ns.getInt("max-messages");
         final var ignoreAttachments = Boolean.TRUE.equals(ns.getBoolean("ignore-attachments"));
         final var ignoreStories = Boolean.TRUE.equals(ns.getBoolean("ignore-stories"));
         final var sendReadReceipts = Boolean.TRUE.equals(ns.getBoolean("send-read-receipts"));
@@ -65,11 +71,9 @@ public class ReceiveCommand implements LocalCommand {
         try {
             final var handler = outputWriter instanceof JsonWriter ? new JsonReceiveMessageHandler(m,
                     (JsonWriter) outputWriter) : new ReceiveMessageHandler(m, (PlainTextWriter) outputWriter);
-            if (timeout < 0) {
+            final var duration = timeout < 0 ? null : Duration.ofMillis((long) (timeout * 1000));
-                m.receiveMessages(handler);
+            final var maxMessages = maxMessagesRaw < 0 ? null : maxMessagesRaw;
-            } else {
+            m.receiveMessages(Optional.ofNullable(duration), Optional.ofNullable(maxMessages), handler);
-                m.receiveMessages(Duration.ofMillis((long) (timeout * 1000)), handler);
-            }
         } catch (IOException e) {
             throw new IOErrorException("Error while receiving messages: " + e.getMessage(), e);
         }

src/main/java/org/asamk/signal/dbus/DbusManagerImpl.java

@@ -58,6 +58,7 @@ import java.util.Map;
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
 import java.util.function.Supplier;
@@ -497,32 +498,30 @@ public class DbusManagerImpl implements Manager {
         }
     }
 
-    @Override
-    public void receiveMessages(final ReceiveMessageHandler handler) throws IOException {
-        addReceiveHandler(handler);
-        try {
-            synchronized (this) {
-                this.wait();
-            }
-        } catch (InterruptedException ignored) {
-        }
-        removeReceiveHandler(handler);
-    }
-
     @Override
     public void receiveMessages(
-            final Duration timeout, final ReceiveMessageHandler handler
+            Optional<Duration> timeout, Optional<Integer> maxMessages, ReceiveMessageHandler handler
     ) throws IOException {
+        final var remainingMessages = new AtomicInteger(maxMessages.orElse(-1));
         final var lastMessage = new AtomicLong(System.currentTimeMillis());
+        final var thread = Thread.currentThread();
 
         final ReceiveMessageHandler receiveHandler = (envelope, e) -> {
             lastMessage.set(System.currentTimeMillis());
             handler.handleMessage(envelope, e);
+            if (remainingMessages.get() > 0) {
+                if (remainingMessages.decrementAndGet() <= 0) {
+                    remainingMessages.set(0);
+                    thread.interrupt();
+                }
+            }
         };
         addReceiveHandler(receiveHandler);
-        while (true) {
+        if (timeout.isPresent()) {
+            while (remainingMessages.get() != 0) {
                 try {
-                final var sleepTimeRemaining = timeout.toMillis() - (System.currentTimeMillis() - lastMessage.get());
+                    final var passedTime = System.currentTimeMillis() - lastMessage.get();
+                    final var sleepTimeRemaining = timeout.get().toMillis() - passedTime;
                     if (sleepTimeRemaining < 0) {
                         break;
                     }
@@ -530,6 +529,15 @@ public class DbusManagerImpl implements Manager {
                 } catch (InterruptedException ignored) {
                 }
             }
+        } else {
+            try {
+                synchronized (this) {
+                    this.wait();
+                }
+            } catch (InterruptedException ignored) {
+            }
+        }
+
         removeReceiveHandler(receiveHandler);
     }