Add optional message limit for receive command

[signal-cli] / lib / src / main / java / org / asamk / signal / manager / helper / ReceiveHelper.java

diff --git a/lib/src/main/java/org/asamk/signal/manager/helper/ReceiveHelper.java b/lib/src/main/java/org/asamk/signal/manager/helper/ReceiveHelper.java
index 9fe1bf54587d78435d8db902161bb772f77d1c7f..c15f4f94cf89d6993a282a52fc89ad70f45199b7 100644 (file)
--- a/lib/src/main/java/org/asamk/signal/manager/helper/ReceiveHelper.java
+++ b/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 {