Skip to content

Commit 1442752

Browse files
ravi-signalravi-signal
committed
Add a tag->tag-set migration command
1 parent 236b049 commit 1442752

File tree

7 files changed

+294
-2
lines changed

7 files changed

+294
-2
lines changed

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

+59-2
Original file line numberDiff line numberDiff line change
@@ -21,24 +21,30 @@
2121
import java.util.List;
2222
import java.util.Map;
2323
import java.util.Objects;
24-
import java.util.UUID;
2524
import java.util.concurrent.CompletableFuture;
2625
import java.util.concurrent.CompletionException;
2726
import java.util.function.Consumer;
2827
import javax.annotation.Nonnull;
2928
import javax.crypto.Mac;
3029
import javax.crypto.spec.SecretKeySpec;
3130
import org.signal.libsignal.zkgroup.receipts.ReceiptCredentialRequest;
31+
import org.slf4j.Logger;
32+
import org.slf4j.LoggerFactory;
3233
import org.whispersystems.textsecuregcm.subscriptions.PaymentProvider;
33-
import org.whispersystems.textsecuregcm.util.AttributeValues;
34+
import org.whispersystems.textsecuregcm.util.ExceptionUtils;
35+
import org.whispersystems.textsecuregcm.util.Util;
36+
import reactor.core.publisher.Flux;
37+
import reactor.core.scheduler.Scheduler;
3438
import software.amazon.awssdk.core.SdkBytes;
3539
import software.amazon.awssdk.services.dynamodb.DynamoDbAsyncClient;
3640
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
3741
import software.amazon.awssdk.services.dynamodb.model.ConditionalCheckFailedException;
3842
import software.amazon.awssdk.services.dynamodb.model.ReturnValue;
43+
import software.amazon.awssdk.services.dynamodb.model.ScanRequest;
3944
import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
4045

4146
public class IssuedReceiptsManager {
47+
private static final Logger log = LoggerFactory.getLogger(IssuedReceiptsManager.class);
4248

4349
public static final String KEY_PROCESSOR_ITEM_ID = "A"; // S (HashKey)
4450
public static final String KEY_ISSUED_RECEIPT_TAG = "B"; // B
@@ -134,4 +140,55 @@ private byte[] generateHmac(String type, Consumer<Mac> byteConsumer) {
134140
throw new AssertionError(e);
135141
}
136142
}
143+
144+
public CompletableFuture<Void> migrateToTagSet(final IssuedReceipt issuedReceipt) {
145+
UpdateItemRequest updateItemRequest = UpdateItemRequest.builder()
146+
.tableName(table)
147+
.key(Map.of(KEY_PROCESSOR_ITEM_ID, s(issuedReceipt.itemId())))
148+
.conditionExpression("attribute_exists(#key) AND #tag = :tag")
149+
.returnValues(ReturnValue.NONE)
150+
.updateExpression("ADD #tags :singletonTag")
151+
.expressionAttributeNames(Map.of(
152+
"#key", KEY_PROCESSOR_ITEM_ID,
153+
"#tag", KEY_ISSUED_RECEIPT_TAG,
154+
"#tags", KEY_ISSUED_RECEIPT_TAG_SET))
155+
.expressionAttributeValues(Map.of(
156+
":tag", b(issuedReceipt.tag()),
157+
":singletonTag", AttributeValue.fromBs(Collections.singletonList(SdkBytes.fromByteArray(issuedReceipt.tag())))))
158+
.build();
159+
return dynamoDbAsyncClient.updateItem(updateItemRequest)
160+
.thenRun(Util.NOOP)
161+
.exceptionally(ExceptionUtils.exceptionallyHandler(ConditionalCheckFailedException.class, e -> {
162+
log.info("Not migrating item {}, because when we tried to migrate it was already deleted", issuedReceipt.itemId());
163+
return null;
164+
}));
165+
}
166+
167+
public record IssuedReceipt(String itemId, byte[] tag) {}
168+
public Flux<IssuedReceipt> receiptsWithoutTagSet(final int segments, final Scheduler scheduler) {
169+
if (segments < 1) {
170+
throw new IllegalArgumentException("Total number of segments must be positive");
171+
}
172+
173+
return Flux.range(0, segments)
174+
.parallel()
175+
.runOn(scheduler)
176+
.flatMap(segment -> dynamoDbAsyncClient.scanPaginator(ScanRequest.builder()
177+
.tableName(table)
178+
.consistentRead(true)
179+
.segment(segment)
180+
.totalSegments(segments)
181+
.filterExpression("attribute_not_exists(#tags)")
182+
.expressionAttributeNames(Map.of("#tags", KEY_ISSUED_RECEIPT_TAG_SET))
183+
.build())
184+
.items()
185+
.flatMapIterable(item -> {
186+
if (!item.containsKey(KEY_ISSUED_RECEIPT_TAG)) {
187+
log.error("Skipping item {} that was missing a receipt tag", item.get(KEY_PROCESSOR_ITEM_ID).s());
188+
return Collections.emptySet();
189+
}
190+
return List.of(new IssuedReceipt(item.get(KEY_PROCESSOR_ITEM_ID).s(), item.get(KEY_ISSUED_RECEIPT_TAG).b().asByteArray()));
191+
}))
192+
.sequential();
193+
}
137194
}

‎service/src/main/java/org/whispersystems/textsecuregcm/workers/CommandDependencies.java

+10
Original file line numberDiff line numberDiff line change
@@ -52,6 +52,7 @@
5252
import org.whispersystems.textsecuregcm.storage.ClientPublicKeys;
5353
import org.whispersystems.textsecuregcm.storage.ClientPublicKeysManager;
5454
import org.whispersystems.textsecuregcm.storage.DynamicConfigurationManager;
55+
import org.whispersystems.textsecuregcm.storage.IssuedReceiptsManager;
5556
import org.whispersystems.textsecuregcm.storage.KeysManager;
5657
import org.whispersystems.textsecuregcm.storage.MessagesCache;
5758
import org.whispersystems.textsecuregcm.storage.MessagesDynamoDb;
@@ -89,6 +90,7 @@ record CommandDependencies(
8990
FaultTolerantRedisClusterClient pushSchedulerCluster,
9091
ClientResources.Builder redisClusterClientResourcesBuilder,
9192
BackupManager backupManager,
93+
IssuedReceiptsManager issuedReceiptsManager,
9294
DynamicConfigurationManager<DynamicConfiguration> dynamicConfigurationManager,
9395
DynamoDbAsyncClient dynamoDbAsyncClient,
9496
PhoneNumberIdentifiers phoneNumberIdentifiers) {
@@ -261,6 +263,13 @@ static CommandDependencies build(
261263
remoteStorageRetryExecutor,
262264
configuration.getCdn3StorageManagerConfiguration()),
263265
clock);
266+
267+
final IssuedReceiptsManager issuedReceiptsManager = new IssuedReceiptsManager(
268+
configuration.getDynamoDbTables().getIssuedReceipts().getTableName(),
269+
configuration.getDynamoDbTables().getIssuedReceipts().getExpiration(),
270+
dynamoDbAsyncClient,
271+
configuration.getDynamoDbTables().getIssuedReceipts().getGenerator());
272+
264273
APNSender apnSender = new APNSender(apnSenderExecutor, configuration.getApnConfiguration());
265274
FcmSender fcmSender = new FcmSender(fcmSenderExecutor, configuration.getFcmConfiguration().credentials().value());
266275
PushNotificationScheduler pushNotificationScheduler = new PushNotificationScheduler(pushSchedulerCluster,
@@ -296,6 +305,7 @@ static CommandDependencies build(
296305
pushSchedulerCluster,
297306
redisClientResourcesBuilder,
298307
backupManager,
308+
issuedReceiptsManager,
299309
dynamicConfigurationManager,
300310
dynamoDbAsyncClient,
301311
phoneNumberIdentifiers

‎service/src/main/java/org/whispersystems/textsecuregcm/workers/IssuedReceiptMigrationCommand.java

+139
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,139 @@
1+
/*
2+
* Copyright 2024 Signal Messenger, LLC
3+
* SPDX-License-Identifier: AGPL-3.0-only
4+
*/
5+
6+
package org.whispersystems.textsecuregcm.workers;
7+
8+
import io.dropwizard.core.Application;
9+
import io.dropwizard.core.setup.Environment;
10+
import io.micrometer.core.instrument.Counter;
11+
import io.micrometer.core.instrument.Metrics;
12+
import net.sourceforge.argparse4j.inf.Namespace;
13+
import net.sourceforge.argparse4j.inf.Subparser;
14+
import org.slf4j.Logger;
15+
import org.slf4j.LoggerFactory;
16+
import org.whispersystems.textsecuregcm.WhisperServerConfiguration;
17+
import org.whispersystems.textsecuregcm.metrics.MetricsUtil;
18+
import org.whispersystems.textsecuregcm.storage.IssuedReceiptsManager;
19+
import reactor.core.publisher.Flux;
20+
import reactor.core.publisher.Mono;
21+
import reactor.core.scheduler.Schedulers;
22+
23+
import java.time.Clock;
24+
import java.util.ArrayList;
25+
import java.util.Collections;
26+
import java.util.Objects;
27+
import java.util.concurrent.CompletableFuture;
28+
import java.util.function.Function;
29+
30+
public class IssuedReceiptMigrationCommand extends AbstractCommandWithDependencies {
31+
32+
private final Logger logger = LoggerFactory.getLogger(getClass());
33+
34+
private static final String SEGMENT_COUNT_ARGUMENT = "segments";
35+
private static final String DRY_RUN_ARGUMENT = "dry-run";
36+
private static final String MAX_CONCURRENCY_ARGUMENT = "max-concurrency";
37+
private static final String BUFFER_ARGUMENT = "buffer";
38+
39+
private static final String INSPECTED_ISSUED_RECEIPTS = MetricsUtil.name(IssuedReceiptMigrationCommand.class,
40+
"inspectedIssuedReceipts");
41+
private static final String MIGRATED_ISSUED_RECEIPTS = MetricsUtil.name(IssuedReceiptMigrationCommand.class,
42+
"migratedIssuedReceipts");
43+
44+
private final Clock clock;
45+
46+
public IssuedReceiptMigrationCommand(final Clock clock) {
47+
super(new Application<>() {
48+
@Override
49+
public void run(final WhisperServerConfiguration configuration, final Environment environment) {
50+
}
51+
}, "migrate-issued-receipts", "Migrates columns in the issued receipts table");
52+
this.clock = clock;
53+
}
54+
55+
@Override
56+
public void configure(final Subparser subparser) {
57+
super.configure(subparser);
58+
59+
subparser.addArgument("--segments")
60+
.type(Integer.class)
61+
.dest(SEGMENT_COUNT_ARGUMENT)
62+
.required(false)
63+
.setDefault(1)
64+
.help("The total number of segments for a DynamoDB scan");
65+
66+
subparser.addArgument("--max-concurrency")
67+
.type(Integer.class)
68+
.dest(MAX_CONCURRENCY_ARGUMENT)
69+
.required(false)
70+
.setDefault(16)
71+
.help("Max concurrency for migration operations");
72+
73+
subparser.addArgument("--dry-run")
74+
.type(Boolean.class)
75+
.dest(DRY_RUN_ARGUMENT)
76+
.required(false)
77+
.setDefault(true)
78+
.help("If true, don’t actually perform migration");
79+
80+
subparser.addArgument("--buffer")
81+
.type(Integer.class)
82+
.dest(BUFFER_ARGUMENT)
83+
.setDefault(16_384)
84+
.help("Records to buffer");
85+
}
86+
87+
@Override
88+
protected void run(final Environment environment, final Namespace namespace,
89+
final WhisperServerConfiguration configuration, final CommandDependencies commandDependencies) throws Exception {
90+
final int bufferSize = namespace.getInt(BUFFER_ARGUMENT);
91+
final int segments = Objects.requireNonNull(namespace.getInt(SEGMENT_COUNT_ARGUMENT));
92+
final int concurrency = Objects.requireNonNull(namespace.getInt(MAX_CONCURRENCY_ARGUMENT));
93+
final boolean dryRun = namespace.getBoolean(DRY_RUN_ARGUMENT);
94+
95+
logger.info("Crawling issuedReceipts with {} segments and {} processors",
96+
segments,
97+
Runtime.getRuntime().availableProcessors());
98+
99+
final Counter inspected = Metrics.counter(INSPECTED_ISSUED_RECEIPTS,
100+
"dryRun", Boolean.toString(dryRun));
101+
final Counter migrated = Metrics.counter(MIGRATED_ISSUED_RECEIPTS,
102+
"dryRun", Boolean.toString(dryRun));
103+
104+
final IssuedReceiptsManager issuedReceiptsManager = commandDependencies.issuedReceiptsManager();
105+
final Flux<IssuedReceiptsManager.IssuedReceipt> receipts =
106+
issuedReceiptsManager.receiptsWithoutTagSet(segments, Schedulers.parallel());
107+
final long count = bufferShuffle(receipts, bufferSize)
108+
.doOnNext(issuedReceipt -> inspected.increment())
109+
.flatMap(issuedReceipt -> Mono
110+
.fromCompletionStage(() -> dryRun
111+
? CompletableFuture.completedFuture(null)
112+
: issuedReceiptsManager.migrateToTagSet(issuedReceipt))
113+
.thenReturn(true)
114+
.retry(3)
115+
.onErrorResume(throwable -> {
116+
logger.error("Failed to migrate {} after 3 attempts, giving up", issuedReceipt.itemId(), throwable);
117+
return Mono.just(false);
118+
}),
119+
concurrency)
120+
.doOnNext(success ->
121+
Metrics.counter(MIGRATED_ISSUED_RECEIPTS,
122+
"dryRun", Boolean.toString(dryRun),
123+
"success", Boolean.toString(success)))
124+
.count()
125+
.block();
126+
logger.info("Attempted to migrate {} issued receipts", count);
127+
}
128+
129+
private static <T> Flux<T> bufferShuffle(Flux<T> f, int bufferSize) {
130+
return f.buffer(bufferSize)
131+
.map(source -> {
132+
final ArrayList<T> shuffled = new ArrayList<>(source);
133+
Collections.shuffle(shuffled);
134+
return shuffled;
135+
})
136+
.limitRate(2)
137+
.flatMapIterable(Function.identity());
138+
}
139+
}

‎service/src/test/java/org/whispersystems/textsecuregcm/storage/IssuedReceiptsManagerTest.java

+83
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,8 @@
1717
import java.util.Set;
1818
import java.util.concurrent.CompletableFuture;
1919
import java.util.stream.Collectors;
20+
import java.util.stream.IntStream;
21+
import java.util.stream.Stream;
2022
import org.assertj.core.api.Condition;
2123
import org.junit.jupiter.api.BeforeEach;
2224
import org.junit.jupiter.api.Test;
@@ -26,11 +28,13 @@
2628
import org.whispersystems.textsecuregcm.subscriptions.PaymentProvider;
2729
import org.whispersystems.textsecuregcm.util.AttributeValues;
2830
import org.whispersystems.textsecuregcm.util.TestRandomUtil;
31+
import reactor.core.scheduler.Schedulers;
2932
import software.amazon.awssdk.core.SdkBytes;
3033
import software.amazon.awssdk.services.dynamodb.DynamoDbClient;
3134
import software.amazon.awssdk.services.dynamodb.model.AttributeValue;
3235
import software.amazon.awssdk.services.dynamodb.model.GetItemRequest;
3336
import software.amazon.awssdk.services.dynamodb.model.GetItemResponse;
37+
import software.amazon.awssdk.services.dynamodb.model.UpdateItemRequest;
3438

3539
class IssuedReceiptsManagerTest {
3640

@@ -89,6 +93,74 @@ void testRecordIssuance() {
8993
assertThat(future).succeedsWithin(Duration.ofSeconds(3));
9094
}
9195

96+
@Test
97+
void testMigrateToTagSet() {
98+
Instant now = Instant.ofEpochSecond(NOW_EPOCH_SECONDS);
99+
100+
issuedReceiptsManager
101+
.recordIssuance("itemId", PaymentProvider.STRIPE, randomReceiptCredentialRequest(), now)
102+
.join();
103+
removeTagSet("itemId");
104+
105+
assertThat(getItem("itemId").item()).doesNotContainKey(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET);
106+
107+
final IssuedReceiptsManager.IssuedReceipt issuedReceipt = issuedReceiptsManager
108+
.receiptsWithoutTagSet(1, Schedulers.immediate())
109+
.blockFirst();
110+
111+
issuedReceiptsManager.migrateToTagSet(issuedReceipt).join();
112+
113+
final Map<String, AttributeValue> item = getItem("itemId").item();
114+
assertThat(item)
115+
.containsKey(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET)
116+
.containsKey(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG);
117+
118+
final List<byte[]> tags = item
119+
.get(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET).bs()
120+
.stream()
121+
.map(SdkBytes::asByteArray)
122+
.toList();
123+
assertThat(tags).hasSize(1);
124+
125+
final byte[] tag = item.get(IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG).b().asByteArray();
126+
assertThat(tags).first().isEqualTo(tag);
127+
}
128+
129+
130+
@Test
131+
void testReceiptsWithoutTagSet() {
132+
Instant now = Instant.ofEpochSecond(NOW_EPOCH_SECONDS);
133+
134+
final int numItems = 100;
135+
final List<String> expectedNoTagSet = IntStream.range(0, numItems)
136+
.boxed()
137+
.flatMap(i -> {
138+
final String itemId = "item-%s".formatted(i);
139+
issuedReceiptsManager.recordIssuance(itemId, PaymentProvider.STRIPE, randomReceiptCredentialRequest(), now).join();
140+
141+
if (i % 2 == 0) {
142+
removeTagSet(itemId);
143+
return Stream.of(itemId);
144+
} else {
145+
return Stream.empty();
146+
}
147+
}).toList();
148+
final List<String> items = issuedReceiptsManager
149+
.receiptsWithoutTagSet(1, Schedulers.immediate())
150+
.map(IssuedReceiptsManager.IssuedReceipt::itemId)
151+
.collectList().block();
152+
assertThat(items).hasSize(numItems / 2);
153+
assertThat(items).containsExactlyInAnyOrderElementsOf(expectedNoTagSet);
154+
}
155+
156+
@Test
157+
void testMigrateAfterRecordExpires() {
158+
final IssuedReceiptsManager.IssuedReceipt issued = new IssuedReceiptsManager.IssuedReceipt("itemId",
159+
TestRandomUtil.nextBytes(32));
160+
// We should succeed but do nothing if the item is deleted by the time we try to migrate it
161+
issuedReceiptsManager.migrateToTagSet(issued).join();
162+
assertThat(getItem("itemId").hasItem()).isFalse();
163+
}
92164

93165
private GetItemResponse getItem(final String itemId) {
94166
final DynamoDbClient client = DYNAMO_DB_EXTENSION.getDynamoDbClient();
@@ -104,4 +176,15 @@ private static ReceiptCredentialRequest randomReceiptCredentialRequest() {
104176
when(request.serialize()).thenReturn(bytes);
105177
return request;
106178
}
179+
180+
private void removeTagSet(final String itemId) {
181+
final DynamoDbClient client = DYNAMO_DB_EXTENSION.getDynamoDbClient();
182+
// Simulate an entry that was written before we wrote the tag set field
183+
client.updateItem(UpdateItemRequest.builder()
184+
.tableName(Tables.ISSUED_RECEIPTS.tableName())
185+
.key(Map.of(IssuedReceiptsManager.KEY_PROCESSOR_ITEM_ID, AttributeValues.s(itemId)))
186+
.updateExpression("REMOVE #tags")
187+
.expressionAttributeNames(Map.of("#tags", IssuedReceiptsManager.KEY_ISSUED_RECEIPT_TAG_SET))
188+
.build());
189+
}
107190
}

‎service/src/test/java/org/whispersystems/textsecuregcm/workers/FinishPushNotificationExperimentCommandTest.java

+1
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,7 @@ void setUp() {
8181
null,
8282
null,
8383
null,
84+
null,
8485
null);
8586

8687
//noinspection unchecked

‎service/src/test/java/org/whispersystems/textsecuregcm/workers/NotifyIdleDevicesCommandTest.java

+1
Original file line numberDiff line numberDiff line change
@@ -68,6 +68,7 @@ private TestNotifyIdleDevicesCommand(final MessagesManager messagesManager,
6868
null,
6969
null,
7070
null,
71+
null,
7172
null);
7273

7374
this.idleDeviceNotificationScheduler = idleDeviceNotificationScheduler;

0 commit comments

Comments
 (0)