Skip to content

Commit 282bcf6

Browse files
ravi-signalravi-signal
committed
Add persistent timer utility backed by redis
1 parent 1446d1a commit 282bcf6

File tree

5 files changed

+257
-47
lines changed

5 files changed

+257
-47
lines changed

‎service/src/main/java/org/whispersystems/textsecuregcm/WhisperServerService.java

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,6 +227,7 @@
227227
import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
228228
import org.whispersystems.textsecuregcm.storage.MessagesManager;
229229
import org.whispersystems.textsecuregcm.storage.OneTimeDonationsManager;
230+
import org.whispersystems.textsecuregcm.storage.PersistentTimer;
230231
import org.whispersystems.textsecuregcm.storage.PhoneNumberIdentifiers;
231232
import org.whispersystems.textsecuregcm.storage.Profiles;
232233
import org.whispersystems.textsecuregcm.storage.ProfilesManager;
@@ -1097,6 +1098,7 @@ protected void configureServer(final ServerBuilder<?> serverBuilder) {
10971098
log.info("Registered spam filter: {}", filter.getClass().getName());
10981099
});
10991100

1101+
final PersistentTimer persistentTimer = new PersistentTimer(rateLimitersCluster, clock);
11001102

11011103
final PhoneVerificationTokenManager phoneVerificationTokenManager = new PhoneVerificationTokenManager(
11021104
phoneNumberIdentifiers, registrationServiceClient, registrationRecoveryPasswordsManager, registrationRecoveryChecker);
@@ -1115,7 +1117,7 @@ protected void configureServer(final ServerBuilder<?> serverBuilder) {
11151117
config.getDeliveryCertificate().ecPrivateKey(), config.getDeliveryCertificate().expiresDays()),
11161118
zkAuthOperations, callingGenericZkSecretParams, clock),
11171119
new ChallengeController(rateLimitChallengeManager, challengeConstraintChecker),
1118-
new DeviceController(accountsManager, clientPublicKeysManager, rateLimiters, config.getMaxDevices()),
1120+
new DeviceController(accountsManager, clientPublicKeysManager, rateLimiters, persistentTimer, config.getMaxDevices()),
11191121
new DeviceCheckController(clock, backupAuthManager, appleDeviceCheckManager, rateLimiters,
11201122
config.getDeviceCheck().backupRedemptionLevel(),
11211123
config.getDeviceCheck().backupRedemptionDuration()),

‎service/src/main/java/org/whispersystems/textsecuregcm/controllers/DeviceController.java

Lines changed: 44 additions & 45 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,6 @@
77
import com.google.common.annotations.VisibleForTesting;
88
import com.google.common.net.HttpHeaders;
99
import io.dropwizard.auth.Auth;
10-
import io.lettuce.core.RedisException;
1110
import io.micrometer.core.instrument.Metrics;
1211
import io.micrometer.core.instrument.Tags;
1312
import io.micrometer.core.instrument.Timer;
@@ -81,6 +80,7 @@
8180
import org.whispersystems.textsecuregcm.storage.DeviceCapability;
8281
import org.whispersystems.textsecuregcm.storage.DeviceSpec;
8382
import org.whispersystems.textsecuregcm.storage.LinkDeviceTokenAlreadyUsedException;
83+
import org.whispersystems.textsecuregcm.storage.PersistentTimer;
8484
import org.whispersystems.textsecuregcm.util.DeviceCapabilityAdapter;
8585
import org.whispersystems.textsecuregcm.util.EnumMapUtil;
8686
import org.whispersystems.textsecuregcm.util.ExceptionUtils;
@@ -100,6 +100,7 @@ public class DeviceController {
100100
private final AccountsManager accounts;
101101
private final ClientPublicKeysManager clientPublicKeysManager;
102102
private final RateLimiters rateLimiters;
103+
private final PersistentTimer persistentTimer;
103104
private final Map<String, Integer> maxDeviceConfiguration;
104105

105106
private final EnumMap<ClientPlatform, AtomicInteger> linkedDeviceListenersByPlatform;
@@ -108,9 +109,11 @@ public class DeviceController {
108109
private static final String LINKED_DEVICE_LISTENER_GAUGE_NAME =
109110
MetricsUtil.name(DeviceController.class, "linkedDeviceListeners");
110111

112+
private static final String WAIT_FOR_LINKED_DEVICE_TIMER_NAMESPACE = "wait_for_linked_device";
111113
private static final String WAIT_FOR_LINKED_DEVICE_TIMER_NAME =
112114
MetricsUtil.name(DeviceController.class, "waitForLinkedDeviceDuration");
113115

116+
private static final String WAIT_FOR_TRANSFER_ARCHIVE_TIMER_NAMESPACE = "wait_for_transfer_archive";
114117
private static final String WAIT_FOR_TRANSFER_ARCHIVE_TIMER_NAME =
115118
MetricsUtil.name(DeviceController.class, "waitForTransferArchiveDuration");
116119

@@ -124,11 +127,13 @@ public class DeviceController {
124127
public DeviceController(final AccountsManager accounts,
125128
final ClientPublicKeysManager clientPublicKeysManager,
126129
final RateLimiters rateLimiters,
130+
final PersistentTimer persistentTimer,
127131
final Map<String, Integer> maxDeviceConfiguration) {
128132

129133
this.accounts = accounts;
130134
this.clientPublicKeysManager = clientPublicKeysManager;
131135
this.rateLimiters = rateLimiters;
136+
this.persistentTimer = persistentTimer;
132137
this.maxDeviceConfiguration = maxDeviceConfiguration;
133138

134139
linkedDeviceListenersByPlatform =
@@ -366,32 +371,30 @@ The amount of time (in seconds) to wait for a response. If the expected device i
366371
@HeaderParam(HttpHeaders.USER_AGENT) String userAgent) {
367372
final AtomicInteger linkedDeviceListenerCounter = getCounterForLinkedDeviceListeners(userAgent);
368373
linkedDeviceListenerCounter.incrementAndGet();
369-
final Timer.Sample sample = Timer.start();
370374

371375
return rateLimiters.getWaitForLinkedDeviceLimiter()
372376
.validateAsync(authenticatedDevice.getAccount().getIdentifier(IdentityType.ACI))
373-
.thenCompose(ignored -> accounts.waitForNewLinkedDevice(
374-
authenticatedDevice.getAccount().getUuid(),
375-
authenticatedDevice.getAuthenticatedDevice(),
376-
tokenIdentifier,
377-
Duration.ofSeconds(timeoutSeconds)))
378-
.thenApply(maybeDeviceInfo -> maybeDeviceInfo
379-
.map(deviceInfo -> Response.status(Response.Status.OK).entity(deviceInfo).build())
380-
.orElseGet(() -> Response.status(Response.Status.NO_CONTENT).build()))
381-
.exceptionally(ExceptionUtils.exceptionallyHandler(IllegalArgumentException.class,
382-
e -> Response.status(Response.Status.BAD_REQUEST).build()))
383-
.whenComplete((response, throwable) -> {
384-
linkedDeviceListenerCounter.decrementAndGet();
385-
386-
if (response != null) {
387-
sample.stop(Timer.builder(WAIT_FOR_LINKED_DEVICE_TIMER_NAME)
388-
.publishPercentileHistogram(true)
389-
.tags(Tags.of(UserAgentTagUtil.getPlatformTag(userAgent),
390-
io.micrometer.core.instrument.Tag.of("deviceFound",
391-
String.valueOf(response.getStatus() == Response.Status.OK.getStatusCode()))))
392-
.register(Metrics.globalRegistry));
393-
}
394-
});
377+
.thenCompose(ignored -> persistentTimer.start(WAIT_FOR_LINKED_DEVICE_TIMER_NAMESPACE, tokenIdentifier))
378+
.thenCompose(sample -> accounts.waitForNewLinkedDevice(
379+
authenticatedDevice.getAccount().getUuid(),
380+
authenticatedDevice.getAuthenticatedDevice(),
381+
tokenIdentifier,
382+
Duration.ofSeconds(timeoutSeconds))
383+
.thenApply(maybeDeviceInfo -> maybeDeviceInfo
384+
.map(deviceInfo -> Response.status(Response.Status.OK).entity(deviceInfo).build())
385+
.orElseGet(() -> Response.status(Response.Status.NO_CONTENT).build()))
386+
.exceptionally(ExceptionUtils.exceptionallyHandler(IllegalArgumentException.class,
387+
e -> Response.status(Response.Status.BAD_REQUEST).build()))
388+
.whenComplete((response, throwable) -> {
389+
linkedDeviceListenerCounter.decrementAndGet();
390+
391+
if (response != null && response.getStatus() == Response.Status.OK.getStatusCode()) {
392+
sample.stop(Timer.builder(WAIT_FOR_LINKED_DEVICE_TIMER_NAME)
393+
.publishPercentileHistogram(true)
394+
.tags(Tags.of(UserAgentTagUtil.getPlatformTag(userAgent)))
395+
.register(Metrics.globalRegistry));
396+
}
397+
}));
395398
}
396399

397400
private AtomicInteger getCounterForLinkedDeviceListeners(final String userAgent) {
@@ -529,7 +532,8 @@ The amount of time (in seconds) to wait for a response. If a transfer archive fo
529532
public CompletionStage<Void> recordTransferArchiveUploaded(@ReadOnly @Auth final AuthenticatedDevice authenticatedDevice,
530533
@NotNull @Valid final TransferArchiveUploadedRequest transferArchiveUploadedRequest) {
531534

532-
return rateLimiters.getUploadTransferArchiveLimiter().validateAsync(authenticatedDevice.getAccount().getIdentifier(IdentityType.ACI))
535+
return rateLimiters.getUploadTransferArchiveLimiter()
536+
.validateAsync(authenticatedDevice.getAccount().getIdentifier(IdentityType.ACI))
533537
.thenCompose(ignored -> accounts.recordTransferArchiveUpload(authenticatedDevice.getAccount(),
534538
transferArchiveUploadedRequest.destinationDeviceId(),
535539
Instant.ofEpochMilli(transferArchiveUploadedRequest.destinationDeviceCreated()),
@@ -568,30 +572,25 @@ The amount of time (in seconds) to wait for a response. If a transfer archive fo
568572

569573
@HeaderParam(HttpHeaders.USER_AGENT) @Nullable String userAgent) {
570574

571-
final Timer.Sample sample = Timer.start();
572575

573576
final String rateLimiterKey = authenticatedDevice.getAccount().getIdentifier(IdentityType.ACI) +
574577
":" + authenticatedDevice.getAuthenticatedDevice().getId();
575578

576579
return rateLimiters.getWaitForTransferArchiveLimiter().validateAsync(rateLimiterKey)
577-
.thenCompose(ignored -> accounts.waitForTransferArchive(authenticatedDevice.getAccount(),
578-
authenticatedDevice.getAuthenticatedDevice(),
579-
Duration.ofSeconds(timeoutSeconds)))
580-
.thenApply(maybeTransferArchive -> maybeTransferArchive
581-
.map(transferArchive -> Response.status(Response.Status.OK).entity(transferArchive).build())
582-
.orElseGet(() -> Response.status(Response.Status.NO_CONTENT).build()))
583-
.whenComplete((response, throwable) -> {
584-
if (response == null) {
585-
return;
586-
}
587-
sample.stop(Timer.builder(WAIT_FOR_TRANSFER_ARCHIVE_TIMER_NAME)
588-
.publishPercentileHistogram(true)
589-
.tags(Tags.of(
590-
UserAgentTagUtil.getPlatformTag(userAgent),
591-
io.micrometer.core.instrument.Tag.of(
592-
"archiveUploaded",
593-
String.valueOf(response.getStatus() == Response.Status.OK.getStatusCode()))))
594-
.register(Metrics.globalRegistry));
595-
});
580+
.thenCompose(ignored -> persistentTimer.start(WAIT_FOR_TRANSFER_ARCHIVE_TIMER_NAMESPACE, rateLimiterKey))
581+
.thenCompose(sample -> accounts.waitForTransferArchive(authenticatedDevice.getAccount(),
582+
authenticatedDevice.getAuthenticatedDevice(),
583+
Duration.ofSeconds(timeoutSeconds))
584+
.thenApply(maybeTransferArchive -> maybeTransferArchive
585+
.map(transferArchive -> Response.status(Response.Status.OK).entity(transferArchive).build())
586+
.orElseGet(() -> Response.status(Response.Status.NO_CONTENT).build()))
587+
.whenComplete((response, throwable) -> {
588+
if (response != null && response.getStatus() == Response.Status.OK.getStatusCode()) {
589+
sample.stop(Timer.builder(WAIT_FOR_TRANSFER_ARCHIVE_TIMER_NAME)
590+
.publishPercentileHistogram(true)
591+
.tags(Tags.of(UserAgentTagUtil.getPlatformTag(userAgent)))
592+
.register(Metrics.globalRegistry));
593+
}
594+
}));
596595
}
597596
}

‎service/src/main/java/org/whispersystems/textsecuregcm/storage/PersistentTimer.java

Lines changed: 104 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,104 @@
1+
/*
2+
* Copyright 2025 Signal Messenger, LLC
3+
* SPDX-License-Identifier: AGPL-3.0-only
4+
*/
5+
6+
package org.whispersystems.textsecuregcm.storage;
7+
8+
import com.google.common.annotations.VisibleForTesting;
9+
import io.lettuce.core.SetArgs;
10+
import io.micrometer.core.instrument.Timer;
11+
import java.time.Clock;
12+
import java.time.Duration;
13+
import java.time.Instant;
14+
import java.util.Optional;
15+
import java.util.concurrent.CompletableFuture;
16+
import javax.annotation.Nullable;
17+
import org.slf4j.Logger;
18+
import org.slf4j.LoggerFactory;
19+
import org.whispersystems.textsecuregcm.redis.FaultTolerantRedisClusterClient;
20+
import org.whispersystems.textsecuregcm.util.Util;
21+
22+
/**
23+
* Timers for operations that may span machines or requests and require a persistently stored timer start itme
24+
*/
25+
public class PersistentTimer {
26+
27+
private static final Logger logger = LoggerFactory.getLogger(PersistentTimer.class);
28+
29+
private static String TIMER_NAMESPACE = "persistent_timer";
30+
@VisibleForTesting
31+
static final Duration TIMER_TTL = Duration.ofHours(1);
32+
33+
private final FaultTolerantRedisClusterClient redisClient;
34+
private final Clock clock;
35+
36+
37+
public PersistentTimer(final FaultTolerantRedisClusterClient redisClient, final Clock clock) {
38+
this.redisClient = redisClient;
39+
this.clock = clock;
40+
}
41+
42+
public class Sample {
43+
44+
private final Instant start;
45+
private final String redisKey;
46+
47+
public Sample(final Instant start, final String redisKey) {
48+
this.start = start;
49+
this.redisKey = redisKey;
50+
}
51+
52+
/**
53+
* Stop the timer, recording the duration between now and the first call to start. This deletes the persistent timer.
54+
*
55+
* @param timer The micrometer timer to record the duration to
56+
* @return A future that completes when the resources associated with the persistent timer have been destroyed
57+
*/
58+
public CompletableFuture<Void> stop(Timer timer) {
59+
Duration duration = Duration.between(start, clock.instant());
60+
timer.record(duration);
61+
return redisClient.withCluster(connection -> connection.async().del(redisKey))
62+
.toCompletableFuture()
63+
.thenRun(Util.NOOP);
64+
}
65+
}
66+
67+
/**
68+
* Start the timer if a timer with the provided namespaced key has not already been started, otherwise return the
69+
* existing sample.
70+
*
71+
* @param namespace A namespace prefix to use for the timer
72+
* @param key The unique key within the namespace that identifies the timer
73+
* @return A future that completes with a {@link Sample} that can later be used to record the final duration.
74+
*/
75+
public CompletableFuture<Sample> start(final String namespace, final String key) {
76+
final Instant now = clock.instant();
77+
final String redisKey = redisKey(namespace, key);
78+
79+
return redisClient.withCluster(connection ->
80+
connection.async().setGet(redisKey, String.valueOf(now.getEpochSecond()), SetArgs.Builder.nx().ex(TIMER_TTL)))
81+
.toCompletableFuture()
82+
.thenApply(serialized -> new Sample(parseStoredTimestamp(serialized).orElse(now), redisKey));
83+
}
84+
85+
@VisibleForTesting
86+
String redisKey(final String namespace, final String key) {
87+
return String.format("%s::%s::%s", TIMER_NAMESPACE, namespace, key);
88+
}
89+
90+
private static Optional<Instant> parseStoredTimestamp(final @Nullable String serialized) {
91+
return Optional
92+
.ofNullable(serialized)
93+
.flatMap(s -> {
94+
try {
95+
return Optional.of(Long.parseLong(s));
96+
} catch (NumberFormatException e) {
97+
logger.warn("Failed to parse stored timestamp {}", s, e);
98+
return Optional.empty();
99+
}
100+
})
101+
.map(Instant::ofEpochSecond);
102+
}
103+
104+
}

‎service/src/test/java/org/whispersystems/textsecuregcm/controllers/DeviceControllerTest.java

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,7 +12,6 @@
1212
import static org.mockito.ArgumentMatchers.anyByte;
1313
import static org.mockito.Mockito.anyString;
1414
import static org.mockito.Mockito.clearInvocations;
15-
import static org.mockito.Mockito.doThrow;
1615
import static org.mockito.Mockito.eq;
1716
import static org.mockito.Mockito.mock;
1817
import static org.mockito.Mockito.never;
@@ -90,6 +89,7 @@
9089
import org.whispersystems.textsecuregcm.storage.DeviceCapability;
9190
import org.whispersystems.textsecuregcm.storage.DeviceSpec;
9291
import org.whispersystems.textsecuregcm.storage.LinkDeviceTokenAlreadyUsedException;
92+
import org.whispersystems.textsecuregcm.storage.PersistentTimer;
9393
import org.whispersystems.textsecuregcm.tests.util.AccountsHelper;
9494
import org.whispersystems.textsecuregcm.tests.util.AuthHelper;
9595
import org.whispersystems.textsecuregcm.tests.util.KeysHelper;
@@ -104,6 +104,7 @@ class DeviceControllerTest {
104104

105105
private static final AccountsManager accountsManager = mock(AccountsManager.class);
106106
private static final ClientPublicKeysManager clientPublicKeysManager = mock(ClientPublicKeysManager.class);
107+
private static final PersistentTimer persistentTimer = mock(PersistentTimer.class);
107108
private static final RateLimiters rateLimiters = mock(RateLimiters.class);
108109
private static final RateLimiter rateLimiter = mock(RateLimiter.class);
109110
@SuppressWarnings("unchecked")
@@ -123,6 +124,7 @@ class DeviceControllerTest {
123124
accountsManager,
124125
clientPublicKeysManager,
125126
rateLimiters,
127+
persistentTimer,
126128
deviceConfiguration);
127129

128130
@RegisterExtension
@@ -161,6 +163,9 @@ void setup() {
161163
when(clientPublicKeysManager.setPublicKey(any(), anyByte(), any()))
162164
.thenReturn(CompletableFuture.completedFuture(null));
163165

166+
when(persistentTimer.start(anyString(), anyString()))
167+
.thenReturn(CompletableFuture.completedFuture(mock(PersistentTimer.Sample.class)));
168+
164169
AccountsHelper.setupMockUpdate(accountsManager);
165170
}
166171

0 commit comments

Comments
 (0)