Add optional message limit for receive command

This commit is contained in:
AsamK 2022-10-31 11:17:52 +01:00
parent 5ed9db4f08
commit de2bfc7f79
6 changed files with 59 additions and 43 deletions

View file

@ -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("Dont 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) {
m.receiveMessages(handler);
} else {
m.receiveMessages(Duration.ofMillis((long) (timeout * 1000)), handler);
}
final var duration = timeout < 0 ? null : Duration.ofMillis((long) (timeout * 1000));
final var maxMessages = maxMessagesRaw < 0 ? null : maxMessagesRaw;
m.receiveMessages(Optional.ofNullable(duration), Optional.ofNullable(maxMessages), handler);
} catch (IOException e) {
throw new IOErrorException("Error while receiving messages: " + e.getMessage(), e);
}

View file

@ -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,39 +498,46 @@ 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) {
try {
final var sleepTimeRemaining = timeout.toMillis() - (System.currentTimeMillis() - lastMessage.get());
if (sleepTimeRemaining < 0) {
break;
if (timeout.isPresent()) {
while (remainingMessages.get() != 0) {
try {
final var passedTime = System.currentTimeMillis() - lastMessage.get();
final var sleepTimeRemaining = timeout.get().toMillis() - passedTime;
if (sleepTimeRemaining < 0) {
break;
}
Thread.sleep(sleepTimeRemaining);
} catch (InterruptedException ignored) {
}
}
} else {
try {
synchronized (this) {
this.wait();
}
Thread.sleep(sleepTimeRemaining);
} catch (InterruptedException ignored) {
}
}
removeReceiveHandler(receiveHandler);
}