When persisting messages fails due to a full queue in DynamoDB, automatically unlink one device to free up room.
Co-authored-by: Chris Eager <79161849+eager-signal@users.noreply.github.com>
This commit is contained in:
Jonathan Klabunde Tomer
GitHub
parent
ce60f13320
commit
8f7bae54fe
|
|
@ -16,17 +16,30 @@ import com.codahale.metrics.Timer;
|
|||
import com.google.common.annotations.VisibleForTesting;
|
||||
import io.dropwizard.lifecycle.Managed;
|
||||
import io.micrometer.core.instrument.Counter;
|
||||
import reactor.core.publisher.Flux;
|
||||
import reactor.util.function.Tuple2;
|
||||
import reactor.util.function.Tuples;
|
||||
import software.amazon.awssdk.services.dynamodb.model.ItemCollectionSizeLimitExceededException;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.time.Instant;
|
||||
import java.util.Comparator;
|
||||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.TimeUnit;
|
||||
import java.util.function.Supplier;
|
||||
import java.util.stream.Collectors;
|
||||
|
||||
import org.apache.commons.lang3.StringUtils;
|
||||
import org.slf4j.Logger;
|
||||
import org.slf4j.LoggerFactory;
|
||||
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
|
||||
import org.whispersystems.textsecuregcm.entities.MessageProtos;
|
||||
import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
|
||||
import org.whispersystems.textsecuregcm.util.Constants;
|
||||
import org.whispersystems.textsecuregcm.util.Util;
|
||||
|
||||
|
|
@ -35,6 +48,8 @@ public class MessagePersister implements Managed {
|
|||
private final MessagesCache messagesCache;
|
||||
private final MessagesManager messagesManager;
|
||||
private final AccountsManager accountsManager;
|
||||
private final ClientPresenceManager clientPresenceManager;
|
||||
private final KeysManager keysManager;
|
||||
|
||||
private final Duration persistDelay;
|
||||
|
||||
|
|
@ -50,27 +65,35 @@ public class MessagePersister implements Managed {
|
|||
private final Counter oversizedQueueCounter = counter(name(MessagePersister.class, "persistQueueOversized"));
|
||||
private final Histogram queueCountHistogram = metricRegistry.histogram(name(MessagePersister.class, "queueCount"));
|
||||
private final Histogram queueSizeHistogram = metricRegistry.histogram(name(MessagePersister.class, "queueSize"));
|
||||
private final ExecutorService executor;
|
||||
|
||||
static final int QUEUE_BATCH_LIMIT = 100;
|
||||
static final int MESSAGE_BATCH_LIMIT = 100;
|
||||
|
||||
private static final long EXCEPTION_PAUSE_MILLIS = Duration.ofSeconds(3).toMillis();
|
||||
public static final Duration UNLINK_TIMEOUT = Duration.ofHours(1);
|
||||
|
||||
private static final int CONSECUTIVE_EMPTY_CACHE_REMOVAL_LIMIT = 3;
|
||||
|
||||
private static final Logger logger = LoggerFactory.getLogger(MessagePersister.class);
|
||||
|
||||
public MessagePersister(final MessagesCache messagesCache, final MessagesManager messagesManager,
|
||||
final AccountsManager accountsManager,
|
||||
final AccountsManager accountsManager, final ClientPresenceManager clientPresenceManager,
|
||||
final KeysManager keysManager,
|
||||
final DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
|
||||
final Duration persistDelay,
|
||||
final int dedicatedProcessWorkerThreadCount) {
|
||||
final int dedicatedProcessWorkerThreadCount,
|
||||
final ExecutorService executor
|
||||
) {
|
||||
this.messagesCache = messagesCache;
|
||||
this.messagesManager = messagesManager;
|
||||
this.accountsManager = accountsManager;
|
||||
this.clientPresenceManager = clientPresenceManager;
|
||||
this.keysManager = keysManager;
|
||||
this.persistDelay = persistDelay;
|
||||
this.workerThreads = new Thread[dedicatedProcessWorkerThreadCount];
|
||||
this.dedicatedProcess = true;
|
||||
this.executor = executor;
|
||||
|
||||
for (int i = 0; i < workerThreads.length; i++) {
|
||||
workerThreads[i] = new Thread(() -> {
|
||||
|
|
@ -139,12 +162,14 @@ public class MessagePersister implements Managed {
|
|||
final UUID accountUuid = MessagesCache.getAccountUuidFromQueueName(queue);
|
||||
final byte deviceId = MessagesCache.getDeviceIdFromQueueName(queue);
|
||||
|
||||
final Optional<Account> maybeAccount = accountsManager.getByAccountIdentifier(accountUuid);
|
||||
if (maybeAccount.isEmpty()) {
|
||||
logger.error("No account record found for account {}", accountUuid);
|
||||
continue;
|
||||
}
|
||||
try {
|
||||
persistQueue(accountUuid, deviceId);
|
||||
persistQueue(maybeAccount.get(), deviceId);
|
||||
} catch (final Exception e) {
|
||||
if (e instanceof ItemCollectionSizeLimitExceededException) {
|
||||
oversizedQueueCounter.increment();
|
||||
}
|
||||
persistQueueExceptionMeter.mark();
|
||||
logger.warn("Failed to persist queue {}::{}; will schedule for retry", accountUuid, deviceId, e);
|
||||
|
||||
|
|
@ -161,14 +186,8 @@ public class MessagePersister implements Managed {
|
|||
}
|
||||
|
||||
@VisibleForTesting
|
||||
void persistQueue(final UUID accountUuid, final byte deviceId) throws MessagePersistenceException {
|
||||
final Optional<Account> maybeAccount = accountsManager.getByAccountIdentifier(accountUuid);
|
||||
|
||||
if (maybeAccount.isEmpty()) {
|
||||
logger.error("No account record found for account {}", accountUuid);
|
||||
return;
|
||||
}
|
||||
|
||||
void persistQueue(final Account account, final byte deviceId) throws MessagePersistenceException {
|
||||
final UUID accountUuid = account.getUuid();
|
||||
try (final Timer.Context ignored = persistQueueTimer.time()) {
|
||||
messagesCache.lockQueueForPersistence(accountUuid, deviceId);
|
||||
|
||||
|
|
@ -197,9 +216,73 @@ public class MessagePersister implements Managed {
|
|||
} while (!messages.isEmpty());
|
||||
|
||||
queueSizeHistogram.update(messageCount);
|
||||
} catch (ItemCollectionSizeLimitExceededException e) {
|
||||
oversizedQueueCounter.increment();
|
||||
unlinkLeastActiveDevice(account, deviceId); // this will either do a deferred reschedule for retry or throw
|
||||
} finally {
|
||||
messagesCache.unlockQueueForPersistence(accountUuid, deviceId);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@VisibleForTesting
|
||||
void unlinkLeastActiveDevice(final Account account, byte destinationDeviceId) throws MessagePersistenceException {
|
||||
if (!messagesCache.lockAccountForMessagePersisterCleanup(account.getUuid())) {
|
||||
// don't try to unlink an account multiple times in parallel; just fail this persist attempt
|
||||
// and we'll try again, hopefully deletions in progress will have made room
|
||||
throw new MessagePersistenceException("account has a full queue and another device-unlinking attempt is in progress");
|
||||
}
|
||||
|
||||
// Selection logic:
|
||||
|
||||
// The primary device is never unlinked
|
||||
// The least-recently-seen inactive device gets unlinked
|
||||
// If there are none, the device with the oldest queued message (not necessarily the
|
||||
// least-recently-seen; a device could be connecting frequently but have some problem fetching
|
||||
// its messages) is unlinked
|
||||
final Device deviceToDelete = account.getDevices()
|
||||
.stream()
|
||||
.filter(d -> !d.isPrimary() && !d.isEnabled())
|
||||
.min(Comparator.comparing(Device::getLastSeen))
|
||||
.or(() ->
|
||||
Flux.fromIterable(account.getDevices())
|
||||
.filter(d -> !d.isPrimary())
|
||||
.flatMap(d ->
|
||||
messagesManager
|
||||
.getEarliestUndeliveredTimestampForDevice(account.getUuid(), d.getId())
|
||||
.map(t -> Tuples.of(d, t)))
|
||||
.sort(Comparator.comparing(Tuple2::getT2))
|
||||
.map(Tuple2::getT1)
|
||||
.next()
|
||||
.blockOptional())
|
||||
.orElseThrow(() -> new MessagePersistenceException("account has a full queue and no unlinkable devices"));
|
||||
|
||||
logger.warn("Failed to persist queue {}::{} due to overfull queue; will unlink device {}{}",
|
||||
account.getUuid(), destinationDeviceId, deviceToDelete.getId(), deviceToDelete.getId() == destinationDeviceId ? "" : " and schedule for retry");
|
||||
executor.execute(
|
||||
() -> {
|
||||
try {
|
||||
// Lock the device's auth token first to prevent it from connecting while we're
|
||||
// destroying its queue, but we don't want to completely remove the device from the
|
||||
// account until we're sure the messages have been cleared because otherwise we won't
|
||||
// be able to find it later to try again, in the event of a failure or timeout
|
||||
final Account updatedAccount = accountsManager.update(
|
||||
account,
|
||||
a -> a.getDevice(deviceToDelete.getId()).ifPresent(Device::lockAuthTokenHash));
|
||||
clientPresenceManager.disconnectPresence(account.getUuid(), deviceToDelete.getId());
|
||||
CompletableFuture
|
||||
.allOf(
|
||||
keysManager.delete(account.getUuid(), deviceToDelete.getId()),
|
||||
messagesManager.clear(account.getUuid(), deviceToDelete.getId()))
|
||||
.orTimeout((UNLINK_TIMEOUT.toSeconds() * 3) / 4, TimeUnit.SECONDS)
|
||||
.join();
|
||||
accountsManager.update(updatedAccount, a -> a.removeDevice(deviceToDelete.getId()));
|
||||
} finally {
|
||||
messagesCache.unlockAccountForMessagePersisterCleanup(account.getUuid());
|
||||
if (deviceToDelete.getId() != destinationDeviceId) { // no point in persisting a device we just purged
|
||||
messagesCache.addQueueToPersist(account.getUuid(), destinationDeviceId);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
|
|
|
|||
|
|
@ -12,6 +12,7 @@ import com.google.protobuf.InvalidProtocolBufferException;
|
|||
import io.dropwizard.lifecycle.Managed;
|
||||
import io.lettuce.core.ScoredValue;
|
||||
import io.lettuce.core.ScriptOutputType;
|
||||
import io.lettuce.core.SetArgs;
|
||||
import io.lettuce.core.ZAddArgs;
|
||||
import io.lettuce.core.cluster.SlotHash;
|
||||
import io.lettuce.core.cluster.models.partitions.RedisClusterNode;
|
||||
|
|
@ -382,6 +383,20 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
|
|||
connection -> connection.sync().del(getPersistInProgressKey(accountUuid, deviceId)));
|
||||
}
|
||||
|
||||
boolean lockAccountForMessagePersisterCleanup(final UUID accountUuid) {
|
||||
return readDeleteCluster.withBinaryCluster(
|
||||
connection -> "OK".equals(
|
||||
connection.sync().set(
|
||||
getUnlinkInProgressKey(accountUuid),
|
||||
LOCK_VALUE,
|
||||
new SetArgs().ex(MessagePersister.UNLINK_TIMEOUT.toSeconds()).nx())));
|
||||
}
|
||||
|
||||
void unlockAccountForMessagePersisterCleanup(final UUID accountUuid) {
|
||||
readDeleteCluster.useBinaryCluster(
|
||||
connection -> connection.sync().del(getUnlinkInProgressKey(accountUuid)));
|
||||
}
|
||||
|
||||
public void addMessageAvailabilityListener(final UUID destinationUuid, final byte deviceId,
|
||||
final MessageAvailabilityListener listener) {
|
||||
final String queueName = getQueueName(destinationUuid, deviceId);
|
||||
|
|
@ -531,6 +546,10 @@ public class MessagesCache extends RedisClusterPubSubAdapter<String, String> imp
|
|||
return ("user_queue_persisting::{" + accountUuid + "::" + deviceId + "}").getBytes(StandardCharsets.UTF_8);
|
||||
}
|
||||
|
||||
private static byte[] getUnlinkInProgressKey(final UUID accountUuid) {
|
||||
return ("user_account_unlinking::{" + accountUuid + "}").getBytes(StandardCharsets.UTF_8);
|
||||
}
|
||||
|
||||
static UUID getAccountUuidFromQueueName(final String queueName) {
|
||||
final int startOfHashTag = queueName.indexOf('{');
|
||||
|
||||
|
|
|
|||
|
|
@ -102,6 +102,10 @@ public class MessagesManager {
|
|||
.tap(Micrometer.metrics(Metrics.globalRegistry));
|
||||
}
|
||||
|
||||
public Mono<Long> getEarliestUndeliveredTimestampForDevice(UUID destinationUuid, byte destinationDevice) {
|
||||
return Mono.from(messagesDynamoDb.load(destinationUuid, destinationDevice, 1)).map(Envelope::getServerTimestamp);
|
||||
}
|
||||
|
||||
public CompletableFuture<Void> clear(UUID destinationUuid) {
|
||||
return CompletableFuture.allOf(
|
||||
messagesCache.clear(destinationUuid),
|
||||
|
|
|
|||
|
|
@ -62,9 +62,14 @@ public class MessagePersisterServiceCommand extends ServerCommand<WhisperServerC
|
|||
|
||||
final MessagePersister messagePersister = new MessagePersister(deps.messagesCache(), deps.messagesManager(),
|
||||
deps.accountsManager(),
|
||||
deps.clientPresenceManager(),
|
||||
deps.keysManager(),
|
||||
dynamicConfigurationManager,
|
||||
Duration.ofMinutes(configuration.getMessageCacheConfiguration().getPersistDelayMinutes()),
|
||||
namespace.getInt(WORKER_COUNT));
|
||||
namespace.getInt(WORKER_COUNT),
|
||||
environment.lifecycle().executorService("messagePersisterUnlinkDeviceExecutor-%d")
|
||||
.maxThreads(2)
|
||||
.build());
|
||||
|
||||
environment.lifecycle().manage(deps.messagesCache());
|
||||
environment.lifecycle().manage(messagePersister);
|
||||
|
|
|
|||
|
|
@ -32,7 +32,9 @@ import org.junit.jupiter.api.Test;
|
|||
import org.junit.jupiter.api.extension.RegisterExtension;
|
||||
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
|
||||
import org.whispersystems.textsecuregcm.entities.MessageProtos;
|
||||
import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
|
||||
import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
|
||||
import org.whispersystems.textsecuregcm.storage.KeysManager;
|
||||
import org.whispersystems.textsecuregcm.storage.DynamoDbExtensionSchema.Tables;
|
||||
import reactor.core.scheduler.Scheduler;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
|
|
@ -83,7 +85,8 @@ class MessagePersisterIntegrationTest {
|
|||
messagesManager = new MessagesManager(messagesDynamoDb, messagesCache, mock(ReportMessageManager.class),
|
||||
messageDeletionExecutorService);
|
||||
messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager,
|
||||
dynamicConfigurationManager, PERSIST_DELAY, 1);
|
||||
mock(ClientPresenceManager.class), mock(KeysManager.class), dynamicConfigurationManager, PERSIST_DELAY, 1,
|
||||
Executors.newSingleThreadExecutor());
|
||||
|
||||
account = mock(Account.class);
|
||||
|
||||
|
|
|
|||
|
|
@ -6,6 +6,7 @@
|
|||
package org.whispersystems.textsecuregcm.storage;
|
||||
|
||||
import static org.junit.jupiter.api.Assertions.assertEquals;
|
||||
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
|
||||
import static org.junit.jupiter.api.Assertions.assertThrows;
|
||||
import static org.junit.jupiter.api.Assertions.assertTimeoutPreemptively;
|
||||
import static org.mockito.ArgumentMatchers.any;
|
||||
|
|
@ -16,9 +17,13 @@ import static org.mockito.Mockito.atLeastOnce;
|
|||
import static org.mockito.Mockito.doAnswer;
|
||||
import static org.mockito.Mockito.mock;
|
||||
import static org.mockito.Mockito.never;
|
||||
import static org.mockito.Mockito.times;
|
||||
import static org.mockito.Mockito.reset;
|
||||
import static org.mockito.Mockito.verify;
|
||||
import static org.mockito.Mockito.when;
|
||||
import static org.whispersystems.textsecuregcm.util.MockUtils.exactly;
|
||||
|
||||
import com.google.common.util.concurrent.MoreExecutors;
|
||||
import com.google.protobuf.ByteString;
|
||||
import io.lettuce.core.cluster.SlotHash;
|
||||
import java.nio.charset.StandardCharsets;
|
||||
|
|
@ -28,6 +33,7 @@ import java.time.Instant;
|
|||
import java.util.List;
|
||||
import java.util.Optional;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
import java.util.concurrent.ExecutorService;
|
||||
import java.util.concurrent.Executors;
|
||||
import java.util.concurrent.ScheduledExecutorService;
|
||||
|
|
@ -41,9 +47,14 @@ import org.mockito.ArgumentCaptor;
|
|||
import org.mockito.stubbing.Answer;
|
||||
import org.whispersystems.textsecuregcm.configuration.dynamic.DynamicConfiguration;
|
||||
import org.whispersystems.textsecuregcm.entities.MessageProtos;
|
||||
import org.whispersystems.textsecuregcm.push.ClientPresenceManager;
|
||||
import org.whispersystems.textsecuregcm.redis.RedisClusterExtension;
|
||||
import org.whispersystems.textsecuregcm.storage.KeysManager;
|
||||
|
||||
import reactor.core.publisher.Mono;
|
||||
import reactor.core.scheduler.Scheduler;
|
||||
import reactor.core.scheduler.Schedulers;
|
||||
import software.amazon.awssdk.services.dynamodb.model.ItemCollectionSizeLimitExceededException;
|
||||
|
||||
class MessagePersisterTest {
|
||||
|
||||
|
|
@ -57,7 +68,10 @@ class MessagePersisterTest {
|
|||
private MessagesDynamoDb messagesDynamoDb;
|
||||
private MessagePersister messagePersister;
|
||||
private AccountsManager accountsManager;
|
||||
private ClientPresenceManager clientPresenceManager;
|
||||
private KeysManager keysManager;
|
||||
private MessagesManager messagesManager;
|
||||
private Account destinationAccount;
|
||||
|
||||
private static final UUID DESTINATION_ACCOUNT_UUID = UUID.randomUUID();
|
||||
private static final String DESTINATION_ACCOUNT_NUMBER = "+18005551234";
|
||||
|
|
@ -74,11 +88,13 @@ class MessagePersisterTest {
|
|||
|
||||
messagesDynamoDb = mock(MessagesDynamoDb.class);
|
||||
accountsManager = mock(AccountsManager.class);
|
||||
clientPresenceManager = mock(ClientPresenceManager.class);
|
||||
keysManager = mock(KeysManager.class);
|
||||
destinationAccount = mock(Account.class);;
|
||||
|
||||
final Account account = mock(Account.class);
|
||||
|
||||
when(accountsManager.getByAccountIdentifier(DESTINATION_ACCOUNT_UUID)).thenReturn(Optional.of(account));
|
||||
when(account.getNumber()).thenReturn(DESTINATION_ACCOUNT_NUMBER);
|
||||
when(accountsManager.getByAccountIdentifier(DESTINATION_ACCOUNT_UUID)).thenReturn(Optional.of(destinationAccount));
|
||||
when(destinationAccount.getUuid()).thenReturn(DESTINATION_ACCOUNT_UUID);
|
||||
when(destinationAccount.getNumber()).thenReturn(DESTINATION_ACCOUNT_NUMBER);
|
||||
when(dynamicConfigurationManager.getConfiguration()).thenReturn(new DynamicConfiguration());
|
||||
|
||||
sharedExecutorService = Executors.newSingleThreadExecutor();
|
||||
|
|
@ -87,8 +103,8 @@ class MessagePersisterTest {
|
|||
messagesCache = new MessagesCache(REDIS_CLUSTER_EXTENSION.getRedisCluster(),
|
||||
REDIS_CLUSTER_EXTENSION.getRedisCluster(), sharedExecutorService, messageDeliveryScheduler,
|
||||
sharedExecutorService, Clock.systemUTC());
|
||||
messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager,
|
||||
dynamicConfigurationManager, PERSIST_DELAY, 1);
|
||||
messagePersister = new MessagePersister(messagesCache, messagesManager, accountsManager, clientPresenceManager,
|
||||
keysManager, dynamicConfigurationManager, PERSIST_DELAY, 1, MoreExecutors.newDirectExecutorService());
|
||||
|
||||
when(messagesManager.persistMessages(any(UUID.class), anyByte(), any())).thenAnswer(invocation -> {
|
||||
final UUID destinationUuid = invocation.getArgument(0);
|
||||
|
|
@ -172,6 +188,7 @@ class MessagePersisterTest {
|
|||
final Account account = mock(Account.class);
|
||||
|
||||
when(accountsManager.getByAccountIdentifier(accountUuid)).thenReturn(Optional.of(account));
|
||||
when(account.getUuid()).thenReturn(accountUuid);
|
||||
when(account.getNumber()).thenReturn(accountNumber);
|
||||
|
||||
insertMessages(accountUuid, deviceId, messagesPerQueue, now);
|
||||
|
|
@ -223,7 +240,150 @@ class MessagePersisterTest {
|
|||
|
||||
assertTimeoutPreemptively(Duration.ofSeconds(1), () ->
|
||||
assertThrows(MessagePersistenceException.class,
|
||||
() -> messagePersister.persistQueue(DESTINATION_ACCOUNT_UUID, DESTINATION_DEVICE_ID)));
|
||||
() -> messagePersister.persistQueue(destinationAccount, DESTINATION_DEVICE_ID)));
|
||||
}
|
||||
|
||||
@Test
|
||||
void testUnlinkFirstInactiveDeviceOnFullQueue() {
|
||||
final String queueName = new String(
|
||||
MessagesCache.getMessageQueueKey(DESTINATION_ACCOUNT_UUID, DESTINATION_DEVICE_ID), StandardCharsets.UTF_8);
|
||||
final int messageCount = 1;
|
||||
final Instant now = Instant.now();
|
||||
|
||||
insertMessages(DESTINATION_ACCOUNT_UUID, DESTINATION_DEVICE_ID, messageCount, now);
|
||||
setNextSlotToPersist(SlotHash.getSlot(queueName));
|
||||
|
||||
final Device primary = mock(Device.class);
|
||||
when(primary.getId()).thenReturn((byte) 1);
|
||||
when(primary.isPrimary()).thenReturn(true);
|
||||
when(primary.isEnabled()).thenReturn(true);
|
||||
final Device activeA = mock(Device.class);
|
||||
when(activeA.getId()).thenReturn((byte) 2);
|
||||
when(activeA.isEnabled()).thenReturn(true);
|
||||
final Device inactiveB = mock(Device.class);
|
||||
final byte inactiveId = 3;
|
||||
when(inactiveB.getId()).thenReturn(inactiveId);
|
||||
when(inactiveB.isEnabled()).thenReturn(false);
|
||||
final Device inactiveC = mock(Device.class);
|
||||
when(inactiveC.getId()).thenReturn((byte) 4);
|
||||
when(inactiveC.isEnabled()).thenReturn(false);
|
||||
final Device activeD = mock(Device.class);
|
||||
when(activeD.getId()).thenReturn((byte) 5);
|
||||
when(activeD.isEnabled()).thenReturn(true);
|
||||
final Device destination = mock(Device.class);
|
||||
when(destination.getId()).thenReturn(DESTINATION_DEVICE_ID);
|
||||
when(destination.isEnabled()).thenReturn(true);
|
||||
|
||||
when(destinationAccount.getDevices()).thenReturn(List.of(primary, activeA, inactiveB, inactiveC, activeD, destination));
|
||||
|
||||
when(messagesManager.persistMessages(any(UUID.class), anyByte(), anyList())).thenThrow(ItemCollectionSizeLimitExceededException.builder().build());
|
||||
when(messagesManager.clear(any(UUID.class), anyByte())).thenReturn(CompletableFuture.completedFuture(null));
|
||||
when(keysManager.delete(any(), eq(inactiveId))).thenReturn(CompletableFuture.completedFuture(null));
|
||||
|
||||
assertTimeoutPreemptively(Duration.ofSeconds(1), () ->
|
||||
messagePersister.persistQueue(destinationAccount, DESTINATION_DEVICE_ID));
|
||||
|
||||
verify(messagesManager, exactly()).clear(DESTINATION_ACCOUNT_UUID, inactiveId);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testUnlinkActiveDeviceWithOldestMessageOnFullQueueWithNoInactiveDevices() {
|
||||
final String queueName = new String(
|
||||
MessagesCache.getMessageQueueKey(DESTINATION_ACCOUNT_UUID, DESTINATION_DEVICE_ID), StandardCharsets.UTF_8);
|
||||
final int messageCount = 1;
|
||||
final Instant now = Instant.now();
|
||||
|
||||
insertMessages(DESTINATION_ACCOUNT_UUID, DESTINATION_DEVICE_ID, messageCount, now);
|
||||
setNextSlotToPersist(SlotHash.getSlot(queueName));
|
||||
|
||||
final Device primary = mock(Device.class);
|
||||
final byte primaryId = 1;
|
||||
when(primary.getId()).thenReturn(primaryId);
|
||||
when(primary.isPrimary()).thenReturn(true);
|
||||
when(primary.isEnabled()).thenReturn(true);
|
||||
when(messagesManager.getEarliestUndeliveredTimestampForDevice(any(), eq(primaryId)))
|
||||
.thenReturn(Mono.just(4L));
|
||||
|
||||
final Device deviceA = mock(Device.class);
|
||||
final byte deviceIdA = 2;
|
||||
when(deviceA.getId()).thenReturn(deviceIdA);
|
||||
when(deviceA.isEnabled()).thenReturn(true);
|
||||
when(messagesManager.getEarliestUndeliveredTimestampForDevice(any(), eq(deviceIdA)))
|
||||
.thenReturn(Mono.empty());
|
||||
|
||||
final Device deviceB = mock(Device.class);
|
||||
final byte deviceIdB = 3;
|
||||
when(deviceB.getId()).thenReturn(deviceIdB);
|
||||
when(deviceB.isEnabled()).thenReturn(true);
|
||||
when(messagesManager.getEarliestUndeliveredTimestampForDevice(any(), eq(deviceIdB)))
|
||||
.thenReturn(Mono.just(2L));
|
||||
|
||||
final Device destination = mock(Device.class);
|
||||
when(destination.getId()).thenReturn(DESTINATION_DEVICE_ID);
|
||||
when(destination.isEnabled()).thenReturn(true);
|
||||
when(messagesManager.getEarliestUndeliveredTimestampForDevice(any(), eq(DESTINATION_DEVICE_ID)))
|
||||
.thenReturn(Mono.just(5L));
|
||||
|
||||
when(destinationAccount.getDevices()).thenReturn(List.of(primary, deviceA, deviceB, destination));
|
||||
|
||||
when(messagesManager.persistMessages(any(UUID.class), anyByte(), anyList())).thenThrow(ItemCollectionSizeLimitExceededException.builder().build());
|
||||
when(messagesManager.clear(any(UUID.class), anyByte())).thenReturn(CompletableFuture.completedFuture(null));
|
||||
when(keysManager.delete(any(), eq(deviceIdB))).thenReturn(CompletableFuture.completedFuture(null));
|
||||
|
||||
assertTimeoutPreemptively(Duration.ofSeconds(1), () ->
|
||||
messagePersister.persistQueue(destinationAccount, DESTINATION_DEVICE_ID));
|
||||
|
||||
verify(messagesManager, exactly()).clear(DESTINATION_ACCOUNT_UUID, deviceIdB);
|
||||
}
|
||||
|
||||
@Test
|
||||
void testUnlinkDestinationDevice() {
|
||||
final String queueName = new String(
|
||||
MessagesCache.getMessageQueueKey(DESTINATION_ACCOUNT_UUID, DESTINATION_DEVICE_ID), StandardCharsets.UTF_8);
|
||||
final int messageCount = 1;
|
||||
final Instant now = Instant.now();
|
||||
|
||||
insertMessages(DESTINATION_ACCOUNT_UUID, DESTINATION_DEVICE_ID, messageCount, now);
|
||||
setNextSlotToPersist(SlotHash.getSlot(queueName));
|
||||
|
||||
final Device primary = mock(Device.class);
|
||||
final byte primaryId = 1;
|
||||
when(primary.getId()).thenReturn(primaryId);
|
||||
when(primary.isPrimary()).thenReturn(true);
|
||||
when(primary.isEnabled()).thenReturn(true);
|
||||
when(messagesManager.getEarliestUndeliveredTimestampForDevice(any(), eq(primaryId)))
|
||||
.thenReturn(Mono.just(1L));
|
||||
|
||||
final Device deviceA = mock(Device.class);
|
||||
final byte deviceIdA = 2;
|
||||
when(deviceA.getId()).thenReturn(deviceIdA);
|
||||
when(deviceA.isEnabled()).thenReturn(true);
|
||||
when(messagesManager.getEarliestUndeliveredTimestampForDevice(any(), eq(deviceIdA)))
|
||||
.thenReturn(Mono.just(3L));
|
||||
|
||||
final Device deviceB = mock(Device.class);
|
||||
final byte deviceIdB = 2;
|
||||
when(deviceB.getId()).thenReturn(deviceIdB);
|
||||
when(deviceB.isEnabled()).thenReturn(true);
|
||||
when(messagesManager.getEarliestUndeliveredTimestampForDevice(any(), eq(deviceIdB)))
|
||||
.thenReturn(Mono.empty());
|
||||
|
||||
final Device destination = mock(Device.class);
|
||||
when(destination.getId()).thenReturn(DESTINATION_DEVICE_ID);
|
||||
when(destination.isEnabled()).thenReturn(true);
|
||||
when(messagesManager.getEarliestUndeliveredTimestampForDevice(any(), eq(DESTINATION_DEVICE_ID)))
|
||||
.thenReturn(Mono.just(2L));
|
||||
|
||||
when(destinationAccount.getDevices()).thenReturn(List.of(primary, deviceA, deviceB, destination));
|
||||
|
||||
when(messagesManager.persistMessages(any(UUID.class), anyByte(), anyList())).thenThrow(ItemCollectionSizeLimitExceededException.builder().build());
|
||||
when(messagesManager.clear(any(UUID.class), anyByte())).thenReturn(CompletableFuture.completedFuture(null));
|
||||
when(keysManager.delete(any(), anyByte())).thenReturn(CompletableFuture.completedFuture(null));
|
||||
|
||||
assertTimeoutPreemptively(Duration.ofSeconds(1), () ->
|
||||
messagePersister.persistQueue(destinationAccount, DESTINATION_DEVICE_ID));
|
||||
|
||||
verify(messagesManager, exactly()).clear(DESTINATION_ACCOUNT_UUID, DESTINATION_DEVICE_ID);
|
||||
}
|
||||
|
||||
@SuppressWarnings("SameParameterValue")
|
||||
|
|
@ -265,5 +425,4 @@ class MessagePersisterTest {
|
|||
REDIS_CLUSTER_EXTENSION.getRedisCluster().useCluster(
|
||||
connection -> connection.sync().set(MessagesCache.NEXT_SLOT_TO_PERSIST_KEY, String.valueOf(nextSlot - 1)));
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
|
|
@ -9,12 +9,23 @@ import static org.mockito.ArgumentMatchers.eq;
|
|||
import static org.mockito.Mockito.doNothing;
|
||||
import static org.mockito.Mockito.doReturn;
|
||||
import static org.mockito.Mockito.doThrow;
|
||||
import static org.mockito.internal.exceptions.Reporter.noMoreInteractionsWanted;
|
||||
import static org.mockito.internal.exceptions.Reporter.wantedButNotInvoked;
|
||||
import static org.mockito.internal.invocation.InvocationMarker.markVerified;
|
||||
import static org.mockito.internal.invocation.InvocationsFinder.findFirstUnverified;
|
||||
import static org.mockito.internal.invocation.InvocationsFinder.findInvocations;
|
||||
|
||||
import java.time.Duration;
|
||||
import java.util.List;
|
||||
import java.util.UUID;
|
||||
import java.util.concurrent.CompletableFuture;
|
||||
import java.util.function.Predicate;
|
||||
|
||||
import org.apache.commons.lang3.RandomUtils;
|
||||
import org.mockito.Mockito;
|
||||
import org.mockito.invocation.Invocation;
|
||||
import org.mockito.invocation.MatchableInvocation;
|
||||
import org.mockito.verification.VerificationMode;
|
||||
import org.whispersystems.textsecuregcm.configuration.secrets.SecretBytes;
|
||||
import org.whispersystems.textsecuregcm.controllers.RateLimitExceededException;
|
||||
import org.whispersystems.textsecuregcm.limits.RateLimiter;
|
||||
|
|
@ -154,4 +165,30 @@ public final class MockUtils {
|
|||
}
|
||||
return new SecretBytes(bytes);
|
||||
}
|
||||
|
||||
/**
|
||||
* modeled after {@link org.mockito.Mockito#only()}, verifies that the matched invocation is the only invocation of
|
||||
* this method
|
||||
*/
|
||||
public static VerificationMode exactly() {
|
||||
return data -> {
|
||||
MatchableInvocation target = data.getTarget();
|
||||
final List<Invocation> allInvocations = data.getAllInvocations();
|
||||
List<Invocation> chunk = findInvocations(allInvocations, target);
|
||||
List<Invocation> otherInvocations = allInvocations.stream()
|
||||
.filter(target::hasSameMethod)
|
||||
.filter(Predicate.not(target::matches))
|
||||
.toList();
|
||||
|
||||
if (!otherInvocations.isEmpty()) {
|
||||
Invocation unverified = findFirstUnverified(otherInvocations);
|
||||
throw noMoreInteractionsWanted(unverified, (List) allInvocations);
|
||||
}
|
||||
if (chunk.isEmpty()) {
|
||||
throw wantedButNotInvoked(target);
|
||||
}
|
||||
markVerified(chunk.get(0), target);
|
||||
};
|
||||
}
|
||||
|
||||
}
|
||||
|
|
|
|||
Loading…
Reference in New Issue