Rework recipient device updates.

SignalServiceKit/src/Contacts/SignalRecipient.h

@@ -25,9 +25,9 @@ NS_ASSUME_NONNULL_BEGIN
 + (instancetype)getOrBuildUnsavedRecipientForRecipientId:(NSString *)recipientId
                                              transaction:(YapDatabaseReadTransaction *)transaction;
 
-- (void)addDevicesToRegisteredRecipient:(NSSet *)devices
-                            transaction:(YapDatabaseReadWriteTransaction *)transaction;
-- (void)removeDevicesFromRecipient:(NSSet *)devices transaction:(YapDatabaseReadWriteTransaction *)transaction;
+- (void)updateRegisteredRecipientWithDevicesToAdd:(nullable NSArray *)devicesToAdd
+                                  devicesToRemove:(nullable NSArray *)devicesToRemove
+                                      transaction:(YapDatabaseReadWriteTransaction *)transaction;
 
 - (NSString *)recipientId;
 

SignalServiceKit/src/Contacts/SignalRecipient.m

@@ -112,6 +112,23 @@ NS_ASSUME_NONNULL_BEGIN
     self.devices = [updatedDevices copy];
 }
 
+- (void)updateRegisteredRecipientWithDevicesToAdd:(nullable NSArray *)devicesToAdd
+                                  devicesToRemove:(nullable NSArray *)devicesToRemove
+                                      transaction:(YapDatabaseReadWriteTransaction *)transaction {
+    OWSAssertDebug(transaction);
+    OWSAssertDebug(devicesToAdd.count > 0 || devicesToRemove.count > 0);
+
+    // Add before we remove, since removeDevicesFromRecipient:...
+    // can removeUnregisteredRecipient:... if the recipient has
+    // no devices left.
+    if (devicesToAdd.count > 0) {
+        [self addDevicesToRegisteredRecipient:[NSSet setWithArray:devicesToAdd] transaction:transaction];
+    }
+    if (devicesToRemove.count > 0) {
+        [self removeDevicesFromRecipient:[NSSet setWithArray:devicesToRemove] transaction:transaction];
+    }
+}
+
 - (void)addDevicesToRegisteredRecipient:(NSSet *)devices transaction:(YapDatabaseReadWriteTransaction *)transaction
 {
     OWSAssertDebug(transaction);
@@ -119,8 +136,8 @@ NS_ASSUME_NONNULL_BEGIN
     
     [self addDevices:devices];
 
-    SignalRecipient *latest =
-        [SignalRecipient markRecipientAsRegisteredAndGet:self.recipientId transaction:transaction];
+    SignalRecipient *latest = [SignalRecipient markRecipientAsRegisteredAndGet:self.recipientId
+                                                                   transaction:transaction];
 
     if ([devices isSubsetOfSet:latest.devices.set]) {
         return;

SignalServiceKit/src/Messages/OWSMessageSender.m

@@ -1264,20 +1264,17 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
                 OWSProdFail([OWSAnalyticsEvents messageSenderErrorNoMissingOrExtraDevices]);
             }
 
-            if (missingDevices && missingDevices.count > 0) {
-                OWSLogInfo(@"Adding missing devices: %@", missingDevices);
-                [recipient addDevicesToRegisteredRecipient:[NSSet setWithArray:missingDevices] transaction:transaction];
-            }
+            [recipient updateRegisteredRecipientWithDevicesToAdd:missingDevices
+                                                 devicesToRemove:extraDevices
+                                                     transaction:transaction];
 
             if (extraDevices && extraDevices.count > 0) {
-                OWSLogInfo(@"removing extra devices: %@", extraDevices);
+                OWSLogInfo(@"Deleting sessions for extra devices: %@", extraDevices);
                 for (NSNumber *extraDeviceId in extraDevices) {
                     [self.primaryStorage deleteSessionForContact:recipient.uniqueId
                                                         deviceId:extraDeviceId.intValue
                                                  protocolContext:transaction];
                 }
-
-                [recipient removeDevicesFromRecipient:[NSSet setWithArray:extraDevices] transaction:transaction];
             }
 
             dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
@@ -1372,7 +1369,9 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
         } @catch (NSException *exception) {
             if ([exception.name isEqualToString:OWSMessageSenderInvalidDeviceException]) {
                 [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-                    [recipient removeDevicesFromRecipient:[NSSet setWithObject:deviceId] transaction:transaction];
+                    [recipient updateRegisteredRecipientWithDevicesToAdd:nil
+                                                         devicesToRemove:@[ deviceId ]
+                                                             transaction:transaction];
                 }];
                 continue;
             } else {
@@ -1409,7 +1408,9 @@ NSString *const OWSMessageSenderRateLimitedException = @"RateLimitedException";
         } @catch (NSException *exception) {
             if ([exception.name isEqualToString:OWSMessageSenderInvalidDeviceException]) {
                 [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *transaction) {
-                    [recipient removeDevicesFromRecipient:[NSSet setWithObject:deviceId] transaction:transaction];
+                    [recipient updateRegisteredRecipientWithDevicesToAdd:nil
+                                                         devicesToRemove:@[ deviceId ]
+                                                             transaction:transaction];
                 }];
             } else {
                 @throw exception;