Don't batch message decryption.

// FREEBIE

SignalServiceKit/src/Messages/OWSBatchMessageProcessor.m

@@ -132,6 +132,7 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSBatchMessageProc
 
 - (void)addJobWithEnvelopeData:(NSData *)envelopeData plaintextData:(NSData *_Nullable)plaintextData
 {
+    // We need to persist the decrypted envelope data ASAP to prevent data loss.
     [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
         OWSMessageContentJob *job =
             [[OWSMessageContentJob alloc] initWithEnvelopeData:envelopeData plaintextData:plaintextData];
@@ -280,6 +281,7 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSBatchMessageProc
 {
     OWSAssert(envelopeData);
 
+    // We need to persist the decrypted envelope data ASAP to prevent data loss.
     [self.finder addJobWithEnvelopeData:envelopeData plaintextData:plaintextData];
 }
 
@@ -322,8 +324,16 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSBatchMessageProc
                        self.tag,
                        jobs.count,
                        [OWSMessageContentJob numberOfKeysInCollection]);
+
+                   // Wait a bit in hopes of increasing the batch size.
+                   // This delay won't affect the first message to arrive when this queue is idle,
+                   // so by definition we're receiving more than one message and can benefit from
+                   // batching.
+                   dispatch_after(
+                       dispatch_time(DISPATCH_TIME_NOW, (int64_t)0.1f * NSEC_PER_SEC), dispatch_get_main_queue(), ^{
                            [self drainQueueWorkStep];
                        });
+               });
            }];
 }
 
@@ -429,6 +439,7 @@ NSString *const OWSMessageContentJobFinderExtensionGroup = @"OWSBatchMessageProc
 {
     OWSAssert(envelopeData);
 
+    // We need to persist the decrypted envelope data ASAP to prevent data loss.
     [self.processingQueue enqueueEnvelopeData:envelopeData plaintextData:plaintextData];
     [self.processingQueue drainQueue];
 }

SignalServiceKit/src/Messages/OWSMessageReceiver.m

@@ -105,27 +105,16 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
     return self;
 }
 
-- (NSArray<OWSMessageDecryptJob *> *)nextJobsForBatchSize:(NSUInteger)maxBatchSize
+- (OWSMessageDecryptJob *_Nullable)nextJob
 {
-    NSMutableArray<OWSMessageDecryptJob *> *jobs = [NSMutableArray new];
+    __block OWSMessageDecryptJob *_Nullable job = nil;
     [self.dbConnection readWithBlock:^(YapDatabaseReadTransaction *_Nonnull transaction) {
         YapDatabaseViewTransaction *viewTransaction = [transaction ext:OWSMessageDecryptJobFinderExtensionName];
         OWSAssert(viewTransaction != nil);
-        [viewTransaction enumerateKeysAndObjectsInGroup:OWSMessageDecryptJobFinderExtensionGroup
+        job = [viewTransaction firstObjectInGroup:OWSMessageDecryptJobFinderExtensionGroup];
-                                             usingBlock:^(NSString *_Nonnull collection,
-                                                 NSString *_Nonnull key,
-                                                 id _Nonnull object,
-                                                 NSUInteger index,
-                                                 BOOL *_Nonnull stop) {
-                                                 OWSMessageDecryptJob *job = object;
-                                                 [jobs addObject:job];
-                                                 if (jobs.count >= maxBatchSize) {
-                                                     *stop = YES;
-                                                 }
-                                             }];
     }];
 
-    return [jobs copy];
+    return job;
 }
 
 - (void)addJobForEnvelope:(OWSSignalServiceProtosEnvelope *)envelope
@@ -135,10 +124,10 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
     }];
 }
 
-- (void)removeJobsWithIds:(NSArray<NSString *> *)uniqueIds
+- (void)removeJobWithId:(NSString *)uniqueId
 {
     [self.dbConnection readWriteWithBlock:^(YapDatabaseReadWriteTransaction *_Nonnull transaction) {
-        [transaction removeObjectsForKeys:uniqueIds inCollection:[OWSMessageDecryptJob collection]];
+        [transaction removeObjectForKey:uniqueId inCollection:[OWSMessageDecryptJob collection]];
     }];
 }
 
@@ -287,71 +276,45 @@ NSString *const OWSMessageDecryptJobFinderExtensionGroup = @"OWSMessageProcessin
 {
     AssertIsOnMainThread();
 
-    NSArray<OWSMessageDecryptJob *> *jobs = [self.finder nextJobsForBatchSize:kIncomingMessageBatchSize];
+    OWSMessageDecryptJob *_Nullable job = [self.finder nextJob];
-    OWSAssert(jobs);
+    if (!job) {
-    if (jobs.count < 1) {
         self.isDrainingQueue = NO;
         DDLogVerbose(@"%@ Queue is drained.", self.tag);
         return;
     }
 
-    [self processJobs:jobs
+    [self processJob:job
-           completion:^{
+          completion:^(BOOL success) {
-               dispatch_async(dispatch_get_main_queue(), ^{
+              [self.finder removeJobWithId:job.uniqueId];
-                   [self.finder removeJobsWithIds:jobs.uniqueIds];
+              DDLogVerbose(@"%@ %@ job. %lu jobs left.",
-                   DDLogVerbose(@"%@ completed %zd jobs. %zd jobs left.",
                   self.tag,
-                       jobs.count,
+                  success ? @"decrypted" : @"failed to decrypt",
-                       [OWSMessageDecryptJob numberOfKeysInCollection]);
+                  (unsigned long)[OWSMessageDecryptJob numberOfKeysInCollection]);
               [self drainQueueWorkStep];
-               });
           }];
 }
 
-- (void)processJobs:(NSArray<OWSMessageDecryptJob *> *)jobs completion:(void (^)())completion
+- (void)processJob:(OWSMessageDecryptJob *)job completion:(void (^)(BOOL))completion
 {
-    [self processJobs:jobs
+    OWSAssert(job);
-         unprocessedJobs:[jobs mutableCopy]
-        plaintextDataMap:[NSMutableDictionary new]
-              completion:completion];
-}
 
-- (void)processJobs:(NSArray<OWSMessageDecryptJob *> *)jobs
+    OWSSignalServiceProtosEnvelope *envelope = job.envelopeProto;
-     unprocessedJobs:(NSMutableArray<OWSMessageDecryptJob *> *)unprocessedJobs
+    dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
-    plaintextDataMap:(NSMutableDictionary<NSString *, NSData *> *)plaintextDataMap
+        [self.messageDecrypter decryptEnvelope:envelope
-          completion:(void (^)())completion
+            successBlock:^(NSData *_Nullable plaintextData) {
-{
-    OWSAssert(jobs.count > 0);
-    OWSAssert(unprocessedJobs.count <= jobs.count);
 
-    if (unprocessedJobs.count < 1) {
+                // We can't decrypt the same message twice, so we need to persist
-        for (OWSMessageDecryptJob *job in jobs) {
+                // the decrypted envelope data ASAP to prevent data loss.
-            NSData *_Nullable plaintextData = plaintextDataMap[job.uniqueId];
                 [self.batchMessageProcessor enqueueEnvelopeData:job.envelopeData plaintextData:plaintextData];
-        }
-        completion();
-        return;
-    }
 
-    dispatch_async(dispatch_get_global_queue(DISPATCH_QUEUE_PRIORITY_DEFAULT, 0), ^{
+                dispatch_async(dispatch_get_main_queue(), ^{
-        OWSAssert(unprocessedJobs.count > 0);
+                    completion(YES);
-        OWSMessageDecryptJob *job = unprocessedJobs.firstObject;
+                });
-        [unprocessedJobs removeObjectAtIndex:0];
-        [self.messageDecrypter decryptEnvelope:job.envelopeProto
-            successBlock:^(NSData *_Nullable plaintextData) {
-                if (plaintextData) {
-                    plaintextDataMap[job.uniqueId] = plaintextData;
-                }
-                [self processJobs:jobs
-                     unprocessedJobs:unprocessedJobs
-                    plaintextDataMap:plaintextDataMap
-                          completion:completion];
             }
             failureBlock:^{
-                [self processJobs:jobs
+                dispatch_async(dispatch_get_main_queue(), ^{
-                     unprocessedJobs:unprocessedJobs
+                    completion(NO);
-                    plaintextDataMap:plaintextDataMap
+                });
-                          completion:completion];
             }];
     });
 }