Make key deletion operations asynchronous

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

@@ -111,11 +111,14 @@ public class DeviceController {
       throw new WebApplicationException(Response.Status.UNAUTHORIZED);
     }
 
+    final CompletableFuture<Void> deleteKeysFuture = keys.delete(account.getUuid(), deviceId);
+
     messages.clear(account.getUuid(), deviceId);
     account = accounts.update(account, a -> a.removeDevice(deviceId));
-    keys.delete(account.getUuid(), deviceId);
     // ensure any messages that came in after the first clear() are also removed
     messages.clear(account.getUuid(), deviceId);
+
+    deleteKeysFuture.join();
   }
 
   @Timed
@@ -336,10 +339,13 @@ public class DeviceController {
     final Account updatedAccount = accounts.update(account, a -> {
       device.setId(a.getNextDeviceId());
 
+      final CompletableFuture<Void> deleteKeysFuture = CompletableFuture.allOf(
+          keys.delete(a.getUuid(), device.getId()),
+          keys.delete(a.getPhoneNumberIdentifier(), device.getId()));
+
       messages.clear(a.getUuid(), device.getId());
 
-      keys.delete(a.getUuid(), device.getId());
+      deleteKeysFuture.join();
-      keys.delete(a.getPhoneNumberIdentifier(), device.getId());
 
       maybeDeviceActivationRequest.ifPresent(deviceActivationRequest -> CompletableFuture.allOf(
               keys.storeEcSignedPreKeys(a.getUuid(),

service/src/main/java/org/whispersystems/textsecuregcm/storage/AccountsManager.java

@@ -228,10 +228,15 @@ public class AccountsManager {
         // confident that everything has already been deleted. In the second case, though, we're taking over an existing
         // account and need to clear out messages and keys that may have been stored for the old account.
         if (!originalUuid.equals(actualUuid)) {
+          final CompletableFuture<Void> deleteKeysFuture = CompletableFuture.allOf(
+              keysManager.delete(actualUuid),
+              keysManager.delete(account.getPhoneNumberIdentifier()));
+
           messagesManager.clear(actualUuid);
-          keysManager.delete(actualUuid);
-          keysManager.delete(account.getPhoneNumberIdentifier());
           profilesManager.deleteAll(actualUuid);
+
+          deleteKeysFuture.join();
+
           clientPresenceManager.disconnectAllPresencesForUuid(actualUuid);
         }
 
@@ -324,8 +329,10 @@ public class AccountsManager {
 
       updatedAccount.set(numberChangedAccount);
 
-      keysManager.delete(phoneNumberIdentifier);
+      CompletableFuture.allOf(
-      keysManager.delete(originalPhoneNumberIdentifier);
+              keysManager.delete(phoneNumberIdentifier),
+              keysManager.delete(originalPhoneNumberIdentifier))
+          .join();
 
       keysManager.storeEcSignedPreKeys(phoneNumberIdentifier, pniSignedPreKeys);
 
@@ -369,9 +376,9 @@ public class AccountsManager {
 
     final List<Long> pqEnabledDeviceIDs = keysManager.getPqEnabledDevices(pni).join();
     keysManager.delete(pni);
-    keysManager.storeEcSignedPreKeys(pni, pniSignedPreKeys);
+    keysManager.storeEcSignedPreKeys(pni, pniSignedPreKeys).join();
     if (pniPqLastResortPreKeys != null) {
-      keysManager.storePqLastResort(pni, pqEnabledDeviceIDs.stream().collect(Collectors.toMap(Function.identity(), pniPqLastResortPreKeys::get)));
+      keysManager.storePqLastResort(pni, pqEnabledDeviceIDs.stream().collect(Collectors.toMap(Function.identity(), pniPqLastResortPreKeys::get))).join();
     }
 
     return updatedAccount;
@@ -755,13 +762,16 @@ public class AccountsManager {
     final CompletableFuture<Void> deleteSecureValueRecoveryServiceDataFuture = secureValueRecovery2Client.deleteBackups(
         account.getUuid());
 
+    final CompletableFuture<Void> deleteKeysFuture = CompletableFuture.allOf(
+        keysManager.delete(account.getUuid()),
+        keysManager.delete(account.getPhoneNumberIdentifier()));
+
     profilesManager.deleteAll(account.getUuid());
-    keysManager.delete(account.getUuid());
-    keysManager.delete(account.getPhoneNumberIdentifier());
     messagesManager.clear(account.getUuid());
     messagesManager.clear(account.getPhoneNumberIdentifier());
     registrationRecoveryPasswordsManager.removeForNumber(account.getNumber());
 
+    deleteKeysFuture.join();
     deleteStorageServiceDataFuture.join();
     deleteBackupServiceDataFuture.join();
     deleteSecureValueRecoveryServiceDataFuture.join();

service/src/main/java/org/whispersystems/textsecuregcm/storage/KeysManager.java

@@ -122,25 +122,23 @@ public class KeysManager {
     return pqPreKeys.getCount(identifier, deviceId);
   }
   
-  public void delete(final UUID accountUuid) {
+  public CompletableFuture<Void> delete(final UUID accountUuid) {
-    CompletableFuture.allOf(
+    return CompletableFuture.allOf(
             ecPreKeys.delete(accountUuid),
             pqPreKeys.delete(accountUuid),
             dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().deleteEcSignedPreKeys()
                 ? ecSignedPreKeys.delete(accountUuid)
                 : CompletableFuture.completedFuture(null),
-            pqLastResortKeys.delete(accountUuid))
+            pqLastResortKeys.delete(accountUuid));
-        .join();
   }
 
-  public void delete(final UUID accountUuid, final long deviceId) {
+  public CompletableFuture<Void> delete(final UUID accountUuid, final long deviceId) {
-    CompletableFuture.allOf(
+    return CompletableFuture.allOf(
             ecPreKeys.delete(accountUuid, deviceId),
             pqPreKeys.delete(accountUuid, deviceId),
             dynamicConfigurationManager.getConfiguration().getEcPreKeyMigrationConfiguration().deleteEcSignedPreKeys()
                 ? ecSignedPreKeys.delete(accountUuid, deviceId)
                 : CompletableFuture.completedFuture(null),
-            pqLastResortKeys.delete(accountUuid, deviceId))
+            pqLastResortKeys.delete(accountUuid, deviceId));
-        .join();
   }
 }

service/src/main/java/org/whispersystems/textsecuregcm/workers/UnlinkDeviceCommand.java

@@ -65,7 +65,7 @@ public class UnlinkDeviceCommand extends EnvironmentCommand<WhisperServerConfigu
     account = deps.accountsManager().update(account, a -> a.removeDevice(deviceId));
 
     System.out.format("Removing keys for device %s::%d\n", aci, deviceId);
-    deps.keysManager().delete(account.getUuid(), deviceId);
+    deps.keysManager().delete(account.getUuid(), deviceId).join();
 
     System.out.format("Clearing additional messages for %s::%d\n", aci, deviceId);
     deps.messagesManager().clear(account.getUuid(), deviceId);

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerChangeNumberIntegrationTest.java

@@ -100,13 +100,16 @@ class AccountsManagerChangeNumberIntegrationTest {
       final PhoneNumberIdentifiers phoneNumberIdentifiers =
           new PhoneNumberIdentifiers(DYNAMO_DB_EXTENSION.getDynamoDbClient(), Tables.PNI.tableName());
 
+      final KeysManager keysManager = mock(KeysManager.class);
+      when(keysManager.delete(any())).thenReturn(CompletableFuture.completedFuture(null));
+
       accountsManager = new AccountsManager(
           accounts,
           phoneNumberIdentifiers,
           CACHE_CLUSTER_EXTENSION.getRedisCluster(),
           accountLockManager,
           deletedAccounts,
-          mock(KeysManager.class),
+          keysManager,
           mock(MessagesManager.class),
           mock(ProfilesManager.class),
           mock(StoredVerificationCodeManager.class),

service/src/test/java/org/whispersystems/textsecuregcm/storage/AccountsManagerTest.java

@@ -157,6 +157,8 @@ class AccountsManagerTest {
       return null;
     }).when(accountLockManager).withLock(any(), any());
 
+    when(keysManager.delete(any())).thenReturn(CompletableFuture.completedFuture(null));
+
     accountsManager = new AccountsManager(
         accounts,
         phoneNumberIdentifiers,
@@ -785,6 +787,7 @@ class AccountsManagerTest {
     final IdentityKey pniIdentityKey = new IdentityKey(Curve.generateKeyPair().getPublicKey());
 
     when(keysManager.getPqEnabledDevices(any())).thenReturn(CompletableFuture.completedFuture(Collections.emptyList()));
+    when(keysManager.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
 
     final Account updatedAccount = accountsManager.updatePniKeys(account, pniIdentityKey, newSignedKeys, null, newRegistrationIds);
 
@@ -829,6 +832,8 @@ class AccountsManagerTest {
     UUID oldPni = account.getPhoneNumberIdentifier();
 
     when(keysManager.getPqEnabledDevices(oldPni)).thenReturn(CompletableFuture.completedFuture(List.of(1L)));
+    when(keysManager.storeEcSignedPreKeys(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
+    when(keysManager.storePqLastResort(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
 
     Map<Long, ECSignedPreKey> oldSignedPreKeys = account.getDevices().stream().collect(Collectors.toMap(Device::getId, Device::getSignedPreKey));
 

service/src/test/java/org/whispersystems/textsecuregcm/storage/KeysManagerTest.java

@@ -185,7 +185,7 @@ class KeysManagerTest {
     assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
     assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
 
-    keysManager.delete(ACCOUNT_UUID);
+    keysManager.delete(ACCOUNT_UUID).join();
 
     assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID).join());
     assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID).join());
@@ -222,7 +222,7 @@ class KeysManagerTest {
     assertTrue(keysManager.getEcSignedPreKey(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
     assertTrue(keysManager.getLastResort(ACCOUNT_UUID, DEVICE_ID + 1).join().isPresent());
 
-    keysManager.delete(ACCOUNT_UUID, DEVICE_ID);
+    keysManager.delete(ACCOUNT_UUID, DEVICE_ID).join();
 
     assertEquals(0, keysManager.getEcCount(ACCOUNT_UUID, DEVICE_ID).join());
     assertEquals(0, keysManager.getPqCount(ACCOUNT_UUID, DEVICE_ID).join());

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

@@ -8,6 +8,7 @@ import static org.assertj.core.api.Assertions.assertThat;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertNull;
 import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.anyLong;
 import static org.mockito.Mockito.clearInvocations;
 import static org.mockito.Mockito.eq;
 import static org.mockito.Mockito.mock;
@@ -155,6 +156,9 @@ class DeviceControllerTest {
     when(accountsManager.getByE164(AuthHelper.VALID_NUMBER_TWO)).thenReturn(Optional.of(maxedAccount));
 
     AccountsHelper.setupMockUpdate(accountsManager);
+
+    when(keysManager.storePqLastResort(any(), any())).thenReturn(CompletableFuture.completedFuture(null));
+    when(keysManager.delete(any(), anyLong())).thenReturn(CompletableFuture.completedFuture(null));
   }
 
   @AfterEach