Skip to content

Commit

Permalink
Add diagnostic dimensions to the "get keys" counter
Browse files Browse the repository at this point in the history
  • Loading branch information
@jon-signal
jon-signal committed Feb 18, 2024
1 parent 3e12a87 commit cd64390
Show file tree
Hide file tree
Showing 3 changed files with 33 additions and 15 deletions.
2 changes: 1 addition & 1 deletion service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -894,7 +894,7 @@ public void run(WhisperServerConfiguration config, Environment environment) thro
new DirectoryV2Controller(directoryV2CredentialsGenerator),
new DonationController(clock, zkReceiptOperations, redeemedReceiptsManager, accountsManager, config.getBadges(),
ReceiptCredentialPresentation::new),
new KeysController(rateLimiters, keysManager, accountsManager),
new KeysController(rateLimiters, keysManager, accountsManager, clientReleaseManager),
new MessageController(rateLimiters, messageByteLimitCardinalityEstimator, messageSender, receiptSender,
accountsManager, messagesManager, pushNotificationManager, reportMessageManager,
multiRecipientMessageExecutor, messageDeliveryScheduler, reportSpamTokenProvider, clientReleaseManager,
Expand Down
33 changes: 24 additions & 9 deletions service/src/main/java/org/whispersystems/textsecuregcm/controllers/KeysController.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,6 +58,7 @@
import org.whispersystems.textsecuregcm.metrics.UserAgentTagUtil;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.ClientReleaseManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeysManager;
import org.whispersystems.textsecuregcm.util.HeaderUtils;
Expand All @@ -72,6 +73,7 @@ public class KeysController {
private final RateLimiters rateLimiters;
private final KeysManager keysManager;
private final AccountsManager accounts;
private final ClientReleaseManager clientReleaseManager;

private static final String KEY_COUNT_DISTRIBUTION_NAME = MetricsUtil.name(KeysController.class, "getKeyCount");
private static final String GET_KEYS_COUNTER_NAME = MetricsUtil.name(KeysController.class, "getKeys");
Expand All @@ -82,10 +84,11 @@ public class KeysController {

private static final CompletableFuture<?>[] EMPTY_FUTURE_ARRAY = new CompletableFuture[0];

public KeysController(RateLimiters rateLimiters, KeysManager keysManager, AccountsManager accounts) {
public KeysController(RateLimiters rateLimiters, KeysManager keysManager, AccountsManager accounts, ClientReleaseManager clientReleaseManager) {
this.rateLimiters = rateLimiters;
this.keysManager = keysManager;
this.accounts = accounts;
this.clientReleaseManager = clientReleaseManager;
}

@GET
Expand Down Expand Up @@ -282,14 +285,26 @@ public PreKeyResponse getDeviceKeys(@Auth Optional<AuthenticatedAccount> auth,
final ECPreKey unsignedEcPreKey = unsignedEcPreKeyFuture.join().orElse(null);
final ECSignedPreKey signedEcPreKey = signedEcPreKeyFuture.join().orElse(null);

Metrics.counter(GET_KEYS_COUNTER_NAME, Tags.of(
Tag.of(PRIMARY_DEVICE_TAG_NAME, String.valueOf(device.isPrimary())),
UserAgentTagUtil.getPlatformTag(userAgent),
Tag.of("targetPlatform", getDevicePlatform(device).map(Enum::name).orElse("unknown")),
Tag.of(IDENTITY_TYPE_TAG_NAME, targetIdentifier.identityType().name()),
Tag.of("isStale", String.valueOf(isDeviceStale(device))),
Tag.of("oneTimeEcKeyAvailable", String.valueOf(unsignedEcPreKey != null))))
.increment();
Tags tags = Tags.of(
Tag.of(PRIMARY_DEVICE_TAG_NAME, String.valueOf(device.isPrimary())),
UserAgentTagUtil.getPlatformTag(userAgent),
Tag.of("targetPlatform", getDevicePlatform(device).map(Enum::name).orElse("unknown")),
Tag.of(IDENTITY_TYPE_TAG_NAME, targetIdentifier.identityType().name()),
Tag.of("isStale", String.valueOf(isDeviceStale(device))),
Tag.of("oneTimeEcKeyAvailable", String.valueOf(unsignedEcPreKey != null)),
Tag.of("authenticationType", auth.isPresent() ? "authenticated" : "anonymous"));

if (auth.isPresent()) {
tags = tags.and(Tag.of("targetIsSelf", String.valueOf(auth.get().getAccount().getUuid().equals(target.getUuid()))));
}

final Optional<Tag> maybeClientVersionTag = UserAgentTagUtil.getClientVersionTag(userAgent, clientReleaseManager);

if (maybeClientVersionTag.isPresent()) {
tags = tags.and(maybeClientVersionTag.get());
}

Metrics.counter(GET_KEYS_COUNTER_NAME, tags).increment();

if (signedEcPreKey != null || unsignedEcPreKey != null || pqPreKey != null) {
final int registrationId = switch (targetIdentifier.identityType()) {
Expand Down
13 changes: 8 additions & 5 deletions service/src/test/java/org/whispersystems/textsecuregcm/controllers/KeysControllerTest.java
View file
Original file line number Diff line number Diff line change
Expand Up @@ -32,7 +32,6 @@
import java.util.OptionalInt;
import java.util.UUID;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import javax.ws.rs.client.Entity;
import javax.ws.rs.core.MediaType;
import javax.ws.rs.core.Response;
Expand Down Expand Up @@ -64,6 +63,7 @@
import org.whispersystems.textsecuregcm.mappers.ServerRejectedExceptionMapper;
import org.whispersystems.textsecuregcm.storage.Account;
import org.whispersystems.textsecuregcm.storage.AccountsManager;
import org.whispersystems.textsecuregcm.storage.ClientReleaseManager;
import org.whispersystems.textsecuregcm.storage.Device;
import org.whispersystems.textsecuregcm.storage.KeysManager;
import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
Expand Down Expand Up @@ -122,9 +122,10 @@ class KeysControllerTest {
private final ECSignedPreKey VALID_DEVICE_SIGNED_KEY = KeysHelper.signedECPreKey(89898, IDENTITY_KEY_PAIR);
private final ECSignedPreKey VALID_DEVICE_PNI_SIGNED_KEY = KeysHelper.signedECPreKey(7777, PNI_IDENTITY_KEY_PAIR);

private final static KeysManager KEYS = mock(KeysManager.class );
private final static AccountsManager accounts = mock(AccountsManager.class );
private final static Account existsAccount = mock(Account.class );
private final static KeysManager KEYS = mock(KeysManager.class);
private final static AccountsManager accounts = mock(AccountsManager.class);
private final static Account existsAccount = mock(Account.class);
private final static ClientReleaseManager clientReleaseManager = mock(ClientReleaseManager.class);

private static final RateLimiters rateLimiters = mock(RateLimiters.class);
private static final RateLimiter rateLimiter = mock(RateLimiter.class );
Expand All @@ -136,7 +137,7 @@ class KeysControllerTest {
.addProvider(new AuthValueFactoryProvider.Binder<>(AuthenticatedAccount.class))
.setTestContainerFactory(new GrizzlyWebTestContainerFactory())
.addResource(new ServerRejectedExceptionMapper())
.addResource(new KeysController(rateLimiters, KEYS, accounts))
.addResource(new KeysController(rateLimiters, KEYS, accounts, clientReleaseManager))
.addResource(new RateLimitExceededExceptionMapper())
.build();

Expand Down Expand Up @@ -276,6 +277,8 @@ void setup() {

when(KEYS.getEcSignedPreKey(AuthHelper.VALID_PNI, AuthHelper.VALID_DEVICE.getId()))
.thenReturn(CompletableFuture.completedFuture(Optional.of(VALID_DEVICE_PNI_SIGNED_KEY)));

when(clientReleaseManager.isVersionActive(any(), any())).thenReturn(true);
}

@AfterEach
Expand Down

0 comments on commit cd64390

Please sign in to comment.