From 166d203e8eb0d62f8631aa72969161d850719dd5 Mon Sep 17 00:00:00 2001 From: Chris Eager Date: Mon, 19 Apr 2021 13:10:26 -0500 Subject: [PATCH] =?UTF-8?q?Don=E2=80=99t=20PUT=20unmigrated=20accounts=20i?= =?UTF-8?q?n=20update()?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .../storage/AccountsDynamoDb.java | 21 +++++----- .../storage/AccountsManager.java | 9 ++++- .../storage/AccountsDynamoDbTest.java | 7 +++- .../tests/storage/AccountsManagerTest.java | 38 +++++++++++++++++++ 4 files changed, 62 insertions(+), 13 deletions(-) diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java index 441c3d1ab..6a55de8cf 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDb.java @@ -3,13 +3,11 @@ package org.whispersystems.textsecuregcm.storage; import static com.codahale.metrics.MetricRegistry.name; import com.amazonaws.services.dynamodbv2.AmazonDynamoDB; -import com.amazonaws.services.dynamodbv2.document.AttributeUpdate; import com.amazonaws.services.dynamodbv2.document.DynamoDB; import com.amazonaws.services.dynamodbv2.document.Item; import com.amazonaws.services.dynamodbv2.document.PrimaryKey; import com.amazonaws.services.dynamodbv2.document.Table; import com.amazonaws.services.dynamodbv2.document.spec.GetItemSpec; -import com.amazonaws.services.dynamodbv2.document.spec.UpdateItemSpec; import com.amazonaws.services.dynamodbv2.model.AttributeValue; import com.amazonaws.services.dynamodbv2.model.CancellationReason; import com.amazonaws.services.dynamodbv2.model.Delete; @@ -19,6 +17,7 @@ import com.amazonaws.services.dynamodbv2.model.ReturnValuesOnConditionCheckFailu import com.amazonaws.services.dynamodbv2.model.TransactWriteItem; import com.amazonaws.services.dynamodbv2.model.TransactWriteItemsRequest; import com.amazonaws.services.dynamodbv2.model.TransactionCanceledException; +import com.amazonaws.services.dynamodbv2.model.UpdateItemRequest; import com.fasterxml.jackson.core.JsonProcessingException; import com.google.common.annotations.VisibleForTesting; import io.micrometer.core.instrument.Metrics; @@ -145,20 +144,22 @@ public class AccountsDynamoDb extends AbstractDynamoDbStore implements AccountSt @Override public void update(Account account) { UPDATE_TIMER.record(() -> { - UpdateItemSpec updateItemSpec; + UpdateItemRequest updateItemRequest; try { - updateItemSpec = new UpdateItemSpec() - .withPrimaryKey( - new PrimaryKey(KEY_ACCOUNT_UUID, UUIDUtil.toByteBuffer(account.getUuid()))) - .withAttributeUpdate( - new AttributeUpdate(ATTR_ACCOUNT_DATA).put(SystemMapper.getMapper().writeValueAsBytes(account)), - new AttributeUpdate(ATTR_MIGRATION_VERSION).put(String.valueOf(account.getDynamoDbMigrationVersion()))); + updateItemRequest = new UpdateItemRequest() + .withTableName(accountsTable.getTableName()) + .withKey(Map.of(KEY_ACCOUNT_UUID, new AttributeValue().withB(UUIDUtil.toByteBuffer(account.getUuid())))) + .withUpdateExpression("SET #data=:data") + .withConditionExpression("attribute_exists(#number)") + .withExpressionAttributeNames(Map.of("#number", ATTR_ACCOUNT_E164, + "#data", ATTR_ACCOUNT_DATA)) + .withExpressionAttributeValues(Map.of(":data", new AttributeValue().withB(ByteBuffer.wrap(SystemMapper.getMapper().writeValueAsBytes(account))))); } catch (JsonProcessingException e) { throw new IllegalArgumentException(e); } - accountsTable.updateItem(updateItemSpec); + client.updateItem(updateItemRequest); }); } diff --git a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java index f5a57d27f..e21b100a4 100644 --- a/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java +++ b/service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java @@ -7,6 +7,7 @@ package org.whispersystems.textsecuregcm.storage; import static com.codahale.metrics.MetricRegistry.name; +import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException; import com.codahale.metrics.MetricRegistry; import com.codahale.metrics.SharedMetricRegistries; import com.codahale.metrics.Timer; @@ -130,7 +131,11 @@ public class AccountsManager { if (dynamoWriteEnabled()) { runSafelyAndRecordMetrics(() -> { - dynamoUpdate(account); + try { + dynamoUpdate(account); + } catch (final ConditionalCheckFailedException e) { + dynamoCreate(account); + } return true; }, Optional.of(account.getUuid()), true, Boolean::compareTo, "update"); } @@ -404,7 +409,7 @@ public class AccountsManager { compare(databaseResult, dynamoResult, comparator); } catch (final Exception e) { - logger.error("Error running " + action + " ih Dynamo", e); + logger.error("Error running " + action + " in Dynamo", e); Metrics.counter(DYNAMO_MIGRATION_ERROR_COUNTER, "action", action).increment(); } diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDbTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDbTest.java index 8f99f5c3b..b352ba15d 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDbTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsDynamoDbTest.java @@ -17,6 +17,7 @@ import com.amazonaws.services.dynamodbv2.document.Item; import com.amazonaws.services.dynamodbv2.document.Table; import com.amazonaws.services.dynamodbv2.document.spec.GetItemSpec; import com.amazonaws.services.dynamodbv2.model.AttributeDefinition; +import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException; import com.amazonaws.services.dynamodbv2.model.CreateTableRequest; import com.amazonaws.services.dynamodbv2.model.KeySchemaElement; import com.amazonaws.services.dynamodbv2.model.KeyType; @@ -172,6 +173,11 @@ class AccountsDynamoDbTest { assertThat(retrieved.isPresent()).isTrue(); verifyStoredState("+14151112222", account.getUuid(), retrieved.get(), account); + + device = generateDevice(1); + Account unknownAccount = generateAccount("+14151113333", UUID.randomUUID(), Collections.singleton(device)); + + assertThatThrownBy(() -> accountsDynamoDb.update(unknownAccount)).isInstanceOfAny(ConditionalCheckFailedException.class); } @Test @@ -301,7 +307,6 @@ class AccountsDynamoDbTest { assertThat(migrated).isFalse(); verifyStoredState("+14151112222", firstUuid, account); - account.setDynamoDbMigrationVersion(account.getDynamoDbMigrationVersion() + 1); migrated = accountsDynamoDb.migrate(account); diff --git a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java index ae9cb5fad..d1cd5471e 100644 --- a/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java +++ b/service/src/test/java/org/whispersystems/textsecuregcm/tests/storage/AccountsManagerTest.java @@ -11,6 +11,7 @@ import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; @@ -19,6 +20,7 @@ import static org.mockito.Mockito.verifyNoMoreInteractions; import static org.mockito.Mockito.verifyZeroInteractions; import static org.mockito.Mockito.when; +import com.amazonaws.services.dynamodbv2.model.ConditionalCheckFailedException; import io.lettuce.core.RedisException; import io.lettuce.core.cluster.api.sync.RedisAdvancedClusterCommands; import java.util.HashSet; @@ -343,6 +345,42 @@ class AccountsManagerTest { verifyNoMoreInteractions(accountsDynamoDb); } + @Test + void testUpdate_dynamoConditionFailed() { + RedisAdvancedClusterCommands commands = mock(RedisAdvancedClusterCommands.class); + FaultTolerantRedisCluster cacheCluster = RedisClusterHelper.buildMockRedisCluster(commands); + Accounts accounts = mock(Accounts.class); + AccountsDynamoDb accountsDynamoDb = mock(AccountsDynamoDb.class); + DirectoryQueue directoryQueue = mock(DirectoryQueue.class); + KeysDynamoDb keysDynamoDb = mock(KeysDynamoDb.class); + MessagesManager messagesManager = mock(MessagesManager.class); + UsernamesManager usernamesManager = mock(UsernamesManager.class); + ProfilesManager profilesManager = mock(ProfilesManager.class); + SecureBackupClient secureBackupClient = mock(SecureBackupClient.class); + SecureStorageClient secureStorageClient = mock(SecureStorageClient.class); + UUID uuid = UUID.randomUUID(); + Account account = new Account("+14152222222", uuid, new HashSet<>(), new byte[16]); + + enableDynamo(true); + + when(commands.get(eq("Account3::" + uuid))).thenReturn(null); + doThrow(ConditionalCheckFailedException.class).when(accountsDynamoDb).update(any(Account.class)); + + AccountsManager accountsManager = new AccountsManager(accounts, accountsDynamoDb, cacheCluster, directoryQueue, keysDynamoDb, messagesManager, usernamesManager, profilesManager, secureStorageClient, secureBackupClient, experimentEnrollmentManager, dynamicConfigurationManager); + + assertEquals(0, account.getDynamoDbMigrationVersion()); + + accountsManager.update(account); + + assertEquals(1, account.getDynamoDbMigrationVersion()); + + verify(accounts, times(1)).update(account); + verifyNoMoreInteractions(accounts); + + verify(accountsDynamoDb, times(1)).update(account); + verify(accountsDynamoDb, times(1)).create(account); + verifyNoMoreInteractions(accountsDynamoDb); + } @Test void testCompareAccounts() {