MirrorMirror
/
Signal-Server
mirror of https://github.com/signalapp/Signal-Server.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

Send directory updates in batches.

This commit is contained in:
Jon Chambers 2021-05-03 12:33:31 -04:00 committed by Jon Chambers
parent 30c9968928
commit 8fdbcbef44
5 changed files with 121 additions and 44 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
service/pom.xml
View File

@ -145,6 +145,12 @@
<version>3.1.0</version>
</dependency>
<dependency>
<groupId>com.google.guava</groupId>
<artifactId>guava</artifactId>
<version>30.1.1-jre</version>
</dependency>
<dependency>
<groupId>com.googlecode.libphonenumber</groupId>
<artifactId>libphonenumber</artifactId>

88
service/src/main/java/org/whispersystems/textsecuregcm/sqs/DirectoryQueue.java
View File

@ -4,6 +4,8 @@
*/
package org.whispersystems.textsecuregcm.sqs;
import static com.codahale.metrics.MetricRegistry.name;
import com.amazonaws.AmazonClientException;
import com.amazonaws.AmazonServiceException;
import com.amazonaws.auth.AWSCredentials;
@ -12,33 +14,33 @@ import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.services.sqs.AmazonSQS;
import com.amazonaws.services.sqs.AmazonSQSClientBuilder;
import com.amazonaws.services.sqs.model.MessageAttributeValue;
import com.amazonaws.services.sqs.model.SendMessageRequest;
import com.amazonaws.services.sqs.model.SendMessageBatchRequest;
import com.amazonaws.services.sqs.model.SendMessageBatchRequestEntry;
import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import com.codahale.metrics.SharedMetricRegistries;
import com.codahale.metrics.Timer;
import com.google.common.annotations.VisibleForTesting;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.stream.Collectors;
import com.google.common.collect.Iterables;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.configuration.SqsConfiguration;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.util.Constants;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import static com.codahale.metrics.MetricRegistry.name;
import org.whispersystems.textsecuregcm.util.Pair;
public class DirectoryQueue {
private static final Logger logger = LoggerFactory.getLogger(DirectoryQueue.class);
private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
private final Meter serviceErrorMeter = metricRegistry.meter(name(DirectoryQueue.class, "serviceError"));
private final Meter clientErrorMeter = metricRegistry.meter(name(DirectoryQueue.class, "clientError"));
private final Timer sendMessageTimer = metricRegistry.timer(name(DirectoryQueue.class, "sendMessage"));
private final MetricRegistry metricRegistry = SharedMetricRegistries.getOrCreate(Constants.METRICS_NAME);
private final Meter serviceErrorMeter = metricRegistry.meter(name(DirectoryQueue.class, "serviceError"));
private final Meter clientErrorMeter = metricRegistry.meter(name(DirectoryQueue.class, "clientError"));
private final Timer sendMessageBatchTimer = metricRegistry.timer(name(DirectoryQueue.class, "sendMessageBatch"));
private final List<String> queueUrls;
private final AmazonSQS sqs;
@ -58,36 +60,54 @@ public class DirectoryQueue {
}
public void refreshRegisteredUser(final Account account) {
sendMessage(account.isEnabled() && account.isDiscoverableByPhoneNumber() ? "add" : "delete", account.getUuid(), account.getNumber());
refreshRegisteredUsers(List.of(account));
}
public void refreshRegisteredUsers(final List<Account> accounts) {
final List<Pair<Account, String>> accountsAndActions = accounts.stream()
.map(account -> new Pair<>(account, account.isEnabled() && account.isDiscoverableByPhoneNumber() ? "add" : "delete"))
.collect(Collectors.toList());
sendUpdateMessages(accountsAndActions);
}
public void deleteAccount(final Account account) {
sendMessage("delete", account.getUuid(), account.getNumber());
sendUpdateMessages(List.of(new Pair<>(account, "delete")));
}
private void sendMessage(String action, UUID uuid, String number) {
final Map<String, MessageAttributeValue> messageAttributes = new HashMap<>();
messageAttributes.put("id", new MessageAttributeValue().withDataType("String").withStringValue(number));
messageAttributes.put("uuid", new MessageAttributeValue().withDataType("String").withStringValue(uuid.toString()));
messageAttributes.put("action", new MessageAttributeValue().withDataType("String").withStringValue(action));
private void sendUpdateMessages(final List<Pair<Account, String>> accountsAndActions) {
for (final String queueUrl : queueUrls) {
final SendMessageRequest sendMessageRequest = new SendMessageRequest()
.withQueueUrl(queueUrl)
for (final List<Pair<Account, String>> partition : Iterables.partition(accountsAndActions, 10)) {
final List<SendMessageBatchRequestEntry> entries = partition.stream().map(pair -> {
final Account account = pair.first();
final String action = pair.second();
return new SendMessageBatchRequestEntry()
.withMessageBody("-")
.withMessageDeduplicationId(UUID.randomUUID().toString())
.withMessageGroupId(number)
.withMessageAttributes(messageAttributes);
try (final Timer.Context ignored = sendMessageTimer.time()) {
sqs.sendMessage(sendMessageRequest);
} catch (AmazonServiceException ex) {
serviceErrorMeter.mark();
logger.warn("sqs service error: ", ex);
} catch (AmazonClientException ex) {
clientErrorMeter.mark();
logger.warn("sqs client error: ", ex);
} catch (Throwable t) {
logger.warn("sqs unexpected error: ", t);
.withMessageGroupId(account.getNumber())
.withMessageAttributes(Map.of(
"id", new MessageAttributeValue().withDataType("String").withStringValue(account.getNumber()),
"uuid", new MessageAttributeValue().withDataType("String").withStringValue(account.getUuid().toString()),
"action", new MessageAttributeValue().withDataType("String").withStringValue(action)
));
}).collect(Collectors.toList());
final SendMessageBatchRequest sendMessageBatchRequest = new SendMessageBatchRequest()
.withQueueUrl(queueUrl)
.withEntries(entries);
try (final Timer.Context ignored = sendMessageBatchTimer.time()) {
sqs.sendMessageBatch(sendMessageBatchRequest);
} catch (AmazonServiceException ex) {
serviceErrorMeter.mark();
logger.warn("sqs service error: ", ex);
} catch (AmazonClientException ex) {
clientErrorMeter.mark();
logger.warn("sqs client error: ", ex);
} catch (Throwable t) {
logger.warn("sqs unexpected error: ", t);
}
}
}
}

9
service/src/main/java/org/whispersystems/textsecuregcm/storage/PushFeedbackProcessor.java
View File

@ -12,6 +12,7 @@ import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
import org.whispersystems.textsecuregcm.util.Constants;
import org.whispersystems.textsecuregcm.util.Util;
import java.util.ArrayList;
import java.util.List;
import java.util.Optional;
import java.util.UUID;
@ -41,6 +42,8 @@ public class PushFeedbackProcessor extends AccountDatabaseCrawlerListener {
@Override
protected void onCrawlChunk(Optional<UUID> fromUuid, List<Account> chunkAccounts) {
final List<Account> directoryUpdateAccounts = new ArrayList<>();
for (Account account : chunkAccounts) {
boolean update = false;
@ -74,8 +77,12 @@ public class PushFeedbackProcessor extends AccountDatabaseCrawlerListener {
if (update) {
accountsManager.update(account);
directoryQueue.refreshRegisteredUser(account);
directoryUpdateAccounts.add(account);
}
}
if (!directoryUpdateAccounts.isEmpty()) {
directoryQueue.refreshRegisteredUsers(directoryUpdateAccounts);
}
}
}

54
service/src/test/java/org/whispersystems/textsecuregcm/sqs/DirectoryQueueTest.java
View File

@ -7,6 +7,7 @@ package org.whispersystems.textsecuregcm.sqs;
import com.amazonaws.services.sqs.AmazonSQS;
import com.amazonaws.services.sqs.model.MessageAttributeValue;
import com.amazonaws.services.sqs.model.SendMessageBatchRequest;
import com.amazonaws.services.sqs.model.SendMessageRequest;
import junitparams.JUnitParamsRunner;
import junitparams.Parameters;
@ -43,13 +44,50 @@ public class DirectoryQueueTest {
directoryQueue.refreshRegisteredUser(account);
final ArgumentCaptor<SendMessageRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
verify(sqs).sendMessage(requestCaptor.capture());
final ArgumentCaptor<SendMessageBatchRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageBatchRequest.class);
verify(sqs).sendMessageBatch(requestCaptor.capture());
final Map<String, MessageAttributeValue> messageAttributes = requestCaptor.getValue().getMessageAttributes();
assertEquals(1, requestCaptor.getValue().getEntries().size());
final Map<String, MessageAttributeValue> messageAttributes = requestCaptor.getValue().getEntries().get(0).getMessageAttributes();
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(expectedAction), messageAttributes.get("action"));
}
@Test
public void testRefreshBatch() {
final AmazonSQS sqs = mock(AmazonSQS.class);
final DirectoryQueue directoryQueue = new DirectoryQueue(List.of("sqs://test"), sqs);
final Account discoverableAccount = mock(Account.class);
when(discoverableAccount.getNumber()).thenReturn("+18005556543");
when(discoverableAccount.getUuid()).thenReturn(UUID.randomUUID());
when(discoverableAccount.isEnabled()).thenReturn(true);
when(discoverableAccount.isDiscoverableByPhoneNumber()).thenReturn(true);
final Account undiscoverableAccount = mock(Account.class);
when(undiscoverableAccount.getNumber()).thenReturn("+18005550987");
when(undiscoverableAccount.getUuid()).thenReturn(UUID.randomUUID());
when(undiscoverableAccount.isEnabled()).thenReturn(true);
when(undiscoverableAccount.isDiscoverableByPhoneNumber()).thenReturn(false);
directoryQueue.refreshRegisteredUsers(List.of(discoverableAccount, undiscoverableAccount));
final ArgumentCaptor<SendMessageBatchRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageBatchRequest.class);
verify(sqs).sendMessageBatch(requestCaptor.capture());
assertEquals(2, requestCaptor.getValue().getEntries().size());
final Map<String, MessageAttributeValue> discoverableAccountAttributes = requestCaptor.getValue().getEntries().get(0).getMessageAttributes();
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(discoverableAccount.getNumber()), discoverableAccountAttributes.get("id"));
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(discoverableAccount.getUuid().toString()), discoverableAccountAttributes.get("uuid"));
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue("add"), discoverableAccountAttributes.get("action"));
final Map<String, MessageAttributeValue> undiscoverableAccountAttributes = requestCaptor.getValue().getEntries().get(1).getMessageAttributes();
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(undiscoverableAccount.getNumber()), undiscoverableAccountAttributes.get("id"));
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue(undiscoverableAccount.getUuid().toString()), undiscoverableAccountAttributes.get("uuid"));
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue("delete"), undiscoverableAccountAttributes.get("action"));
}
@Test
public void testSendMessageMultipleQueues() {
final AmazonSQS sqs = mock(AmazonSQS.class);
@ -63,11 +101,13 @@ public class DirectoryQueueTest {
directoryQueue.refreshRegisteredUser(account);
final ArgumentCaptor<SendMessageRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageRequest.class);
verify(sqs, times(2)).sendMessage(requestCaptor.capture());
final ArgumentCaptor<SendMessageBatchRequest> requestCaptor = ArgumentCaptor.forClass(SendMessageBatchRequest.class);
verify(sqs, times(2)).sendMessageBatch(requestCaptor.capture());
for (final SendMessageRequest sendMessageRequest : requestCaptor.getAllValues()) {
final Map<String, MessageAttributeValue> messageAttributes = sendMessageRequest.getMessageAttributes();
for (final SendMessageBatchRequest sendMessageBatchRequest : requestCaptor.getAllValues()) {
assertEquals(1, requestCaptor.getValue().getEntries().size());
final Map<String, MessageAttributeValue> messageAttributes = sendMessageBatchRequest.getEntries().get(0).getMessageAttributes();
assertEquals(new MessageAttributeValue().withDataType("String").withStringValue("add"), messageAttributes.get("action"));
}
}

8
service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/PushFeedbackProcessorTest.java
View File

@ -7,6 +7,7 @@ package org.whispersystems.textsecuregcm.tests.storage;
import org.junit.Before;
import org.junit.Test;
import org.mockito.ArgumentCaptor;
import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountDatabaseCrawlerRestartException;
@ -22,6 +23,7 @@ import java.util.Set;
import java.util.UUID;
import java.util.concurrent.TimeUnit;
import static org.junit.Assert.assertTrue;
import static org.mockito.Mockito.*;
public class PushFeedbackProcessorTest {
@ -131,8 +133,10 @@ public class PushFeedbackProcessorTest {
verify(accountsManager).update(eq(stillActiveAccount));
verify(directoryQueue).refreshRegisteredUser(undiscoverableAccount);
verify(directoryQueue).refreshRegisteredUser(uninstalledAccount);
final ArgumentCaptor<List<Account>> refreshedAccountArgumentCaptor = ArgumentCaptor.forClass(List.class);
verify(directoryQueue).refreshRegisteredUsers(refreshedAccountArgumentCaptor.capture());
assertTrue(refreshedAccountArgumentCaptor.getValue().containsAll(List.of(undiscoverableAccount, uninstalledAccount)));
}