Signal/Signal-Server
1
0
Fork 0
mirror of https://github.com/signalapp/Signal-Server synced 2026-04-21 12:38:04 +01:00
Code Issues Packages Projects Releases Wiki Activity

Split traffic between the Postgres and Dynamo pre-key stores.

Browse Source
This commit is contained in:
Jon Chambers
2021-01-26 11:32:15 -05:00
committed by Jon Chambers
parent d4d9403829
commit 3298db8683
3 changed files with 40 additions and 22 deletions
Download Patch File Download Diff File Expand all files Collapse all files

46
service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
View File

@@ -6,8 +6,6 @@ package org.whispersystems.textsecuregcm.controllers;
import com.codahale.metrics.annotation.Timed;
import io.dropwizard.auth.Auth;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.whispersystems.textsecuregcm.auth.AmbiguousIdentifier;
import org.whispersystems.textsecuregcm.auth.Anonymous;
import org.whispersystems.textsecuregcm.auth.DisabledPermittedAccount;
@@ -18,6 +16,7 @@ import org.whispersystems.textsecuregcm.entities.PreKeyResponse;
import org.whispersystems.textsecuregcm.entities.PreKeyResponseItem;
import org.whispersystems.textsecuregcm.entities.PreKeyState;
import org.whispersystems.textsecuregcm.entities.SignedPreKey;
import org.whispersystems.textsecuregcm.experiment.ExperimentEnrollmentManager;
import org.whispersystems.textsecuregcm.limits.RateLimiters;
import org.whispersystems.textsecuregcm.sqs.DirectoryQueue;
import org.whispersystems.textsecuregcm.storage.Account;
@@ -25,6 +24,8 @@ import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeyRecord;
import org.whispersystems.textsecuregcm.storage.Keys;
import org.whispersystems.textsecuregcm.storage.KeysDynamoDb;
import org.whispersystems.textsecuregcm.storage.PreKeyStore;
import javax.validation.Valid;
import javax.ws.rs.Consumes;
@@ -45,24 +46,29 @@ import java.util.Optional;
@Path("/v2/keys")
public class KeysController {
private static final Logger logger = LoggerFactory.getLogger(KeysController.class);
private final RateLimiters rateLimiters;
private final Keys keys;
private final KeysDynamoDb keysDynamoDb;
private final AccountsManager accounts;
private final DirectoryQueue directoryQueue;
private final ExperimentEnrollmentManager experimentEnrollmentManager;
private final RateLimiters rateLimiters;
private final Keys keys;
private final AccountsManager accounts;
private final DirectoryQueue directoryQueue;
private static final String DYNAMODB_CONSUMER_EXPERIMENT = "keys_dynamodb_consumer";
private static final String DYNAMODB_PRODUCER_EXPERIMENT = "keys_dynamodb_producer";
public KeysController(RateLimiters rateLimiters, Keys keys, AccountsManager accounts, DirectoryQueue directoryQueue) {
this.rateLimiters = rateLimiters;
this.keys = keys;
this.accounts = accounts;
this.directoryQueue = directoryQueue;
public KeysController(RateLimiters rateLimiters, Keys keys, KeysDynamoDb keysDynamoDb, AccountsManager accounts, DirectoryQueue directoryQueue, final ExperimentEnrollmentManager experimentEnrollmentManager) {
this.rateLimiters = rateLimiters;
this.keys = keys;
this.keysDynamoDb = keysDynamoDb;
this.accounts = accounts;
this.directoryQueue = directoryQueue;
this.experimentEnrollmentManager = experimentEnrollmentManager;
}
@GET
@Produces(MediaType.APPLICATION_JSON)
public PreKeyCount getStatus(@Auth Account account) {
int count = keys.getCount(account, account.getAuthenticatedDevice().get().getId());
int count = getPreKeyStoreForProducer(account).getCount(account, account.getAuthenticatedDevice().get().getId());
if (count > 0) {
count = count - 1;
@@ -98,7 +104,7 @@ public class KeysController {
}
}
keys.store(account, device.getId(), preKeys.getPreKeys());
getPreKeyStoreForProducer(account).store(account, device.getId(), preKeys.getPreKeys());
}
@Timed
@@ -179,14 +185,22 @@ public class KeysController {
private List<KeyRecord> getLocalKeys(Account destination, String deviceIdSelector) {
try {
if (deviceIdSelector.equals("*")) {
return keys.take(destination);
return getPreKeyStoreForConsumer(destination).take(destination);
}
long deviceId = Long.parseLong(deviceIdSelector);
return keys.take(destination, deviceId);
return getPreKeyStoreForConsumer(destination).take(destination, deviceId);
} catch (NumberFormatException e) {
throw new WebApplicationException(Response.status(422).build());
}
}
private PreKeyStore getPreKeyStoreForProducer(final Account account) {
return experimentEnrollmentManager.isEnrolled(account.getUuid(), DYNAMODB_PRODUCER_EXPERIMENT) ? keysDynamoDb : keys;
}
private PreKeyStore getPreKeyStoreForConsumer(final Account account) {
return experimentEnrollmentManager.isEnrolled(account.getUuid(), DYNAMODB_CONSUMER_EXPERIMENT) ? keysDynamoDb : keys;
}
}