From 594e12f23c95de34095da12f75ccace4948c6601 Mon Sep 17 00:00:00 2001
From: Matthew Chen <matthew@signal.org>
Date: Tue, 24 Apr 2018 16:02:35 -0400
Subject: [PATCH] Respond to CR.

---
 .../migrations/OWS109OutgoingMessageState.m   | 74 +++++++++++++------
 .../migrations/OWSDatabaseMigrationRunner.m   |  4 +-
 2 files changed, 54 insertions(+), 24 deletions(-)

diff --git a/SignalMessaging/environment/migrations/OWS109OutgoingMessageState.m b/SignalMessaging/environment/migrations/OWS109OutgoingMessageState.m
index aa1994272..6eef8fe8d 100644
--- a/SignalMessaging/environment/migrations/OWS109OutgoingMessageState.m
+++ b/SignalMessaging/environment/migrations/OWS109OutgoingMessageState.m
@@ -3,6 +3,7 @@
 //
 
 #import "OWS109OutgoingMessageState.h"
+#import <SignalServiceKit/OWSPrimaryStorage.h>
 #import <SignalServiceKit/TSOutgoingMessage.h>
 #import <YapDatabase/YapDatabaseTransaction.h>
 
@@ -18,34 +19,61 @@ static NSString *const OWS109OutgoingMessageStateMigrationId = @"109";
     return OWS109OutgoingMessageStateMigrationId;
 }
 
-- (void)runUpWithTransaction:(YapDatabaseReadWriteTransaction *)transaction
+// Override parent migration
+- (void)runUpWithCompletion:(OWSDatabaseMigrationCompletion)completion
 {
-    OWSAssert(transaction);
-
-    // Persist the migration of the outgoing message state.
-    // For performance, we want to upgrade all existing outgoing messages in
-    // a single transaction.
-    NSMutableArray<NSString *> *messageIds =
-        [[transaction allKeysInCollection:TSOutgoingMessage.collection] mutableCopy];
-    DDLogInfo(@"%@ Migrating %zd outgoing messages.", self.logTag, messageIds.count);
-    while (messageIds.count > 0) {
+    OWSAssert(completion);
+
+    OWSDatabaseConnection *dbConnection = (OWSDatabaseConnection *)self.primaryStorage.newDatabaseConnection;
+
+    [dbConnection asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
+        NSMutableArray<NSString *> *messageIds =
+            [[transaction allKeysInCollection:TSOutgoingMessage.collection] mutableCopy];
+        DDLogInfo(@"%@ Migrating %zd outgoing messages.", self.logTag, messageIds.count);
+
+        [self processBatch:messageIds
+              dbConnection:dbConnection
+                completion:^{
+                    DDLogInfo(@"Completed migration %@", self.uniqueId);
+
+                    [self save];
+
+                    completion();
+                }];
+    }];
+}
+
+- (void)processBatch:(NSMutableArray<NSString *> *)messageIds
+        dbConnection:(OWSDatabaseConnection *)dbConnection
+          completion:(OWSDatabaseMigrationCompletion)completion
+{
+    OWSAssert(dbConnection);
+    OWSAssert(completion);
+
+    DDLogVerbose(@"%@ %s: %zd", self.logTag, __PRETTY_FUNCTION__, messageIds.count);
+
+    if (messageIds.count < 1) {
+        completion();
+        return;
+    }
+
+    [dbConnection asyncReadWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
         const int kBatchSize = 1000;
-        @autoreleasepool {
-            for (int i = 0; i < kBatchSize; i++) {
-                if (messageIds.count == 0) {
-                    break;
-                }
-                NSString *messageId = [messageIds lastObject];
-                [messageIds removeLastObject];
-                id message = [transaction objectForKey:messageId inCollection:TSOutgoingMessage.collection];
-                if (![message isKindOfClass:[TSOutgoingMessage class]]) {
-                    return;
-                }
-                TSOutgoingMessage *outgoingMessage = (TSOutgoingMessage *)message;
-                [outgoingMessage saveWithTransaction:transaction];
+        for (int i = 0; i < kBatchSize && messageIds.count > 0; i++) {
+            NSString *messageId = [messageIds lastObject];
+            [messageIds removeLastObject];
+            id message = [transaction objectForKey:messageId inCollection:TSOutgoingMessage.collection];
+            if (![message isKindOfClass:[TSOutgoingMessage class]]) {
+                continue;
             }
+            TSOutgoingMessage *outgoingMessage = (TSOutgoingMessage *)message;
+            [outgoingMessage saveWithTransaction:transaction];
         }
     }
+        completionBlock:^{
+            // Process the next batch.
+            [self processBatch:messageIds dbConnection:dbConnection completion:completion];
+        }];
 }
 
 @end
diff --git a/SignalMessaging/environment/migrations/OWSDatabaseMigrationRunner.m b/SignalMessaging/environment/migrations/OWSDatabaseMigrationRunner.m
index d5537cb1c..7fff0cfef 100644
--- a/SignalMessaging/environment/migrations/OWSDatabaseMigrationRunner.m
+++ b/SignalMessaging/environment/migrations/OWSDatabaseMigrationRunner.m
@@ -10,6 +10,7 @@
 #import "OWS105AttachmentFilePaths.h"
 #import "OWS107LegacySounds.h"
 #import "OWS108CallLoggingPreference.h"
+#import "OWS109OutgoingMessageState.h"
 #import "OWSDatabaseMigration.h"
 #import <SignalMessaging/SignalMessaging-Swift.h>
 #import <SignalServiceKit/AppContext.h>
@@ -42,7 +43,8 @@ NS_ASSUME_NONNULL_BEGIN
         [[OWS105AttachmentFilePaths alloc] initWithPrimaryStorage:primaryStorage],
         [[OWS106EnsureProfileComplete alloc] initWithPrimaryStorage:primaryStorage],
         [[OWS107LegacySounds alloc] initWithPrimaryStorage:primaryStorage],
-        [[OWS108CallLoggingPreference alloc] initWithPrimaryStorage:primaryStorage]
+        [[OWS108CallLoggingPreference alloc] initWithPrimaryStorage:primaryStorage],
+        [[OWS109OutgoingMessageState alloc] initWithPrimaryStorage:primaryStorage]
     ];
 }