Handle StoryMessage disappearance in OWSDisappearingMessagesJob

Signal/src/ViewControllers/HomeView/Stories/IncomingStoryViewModel.swift

@@ -80,10 +80,10 @@ struct IncomingStoryViewModel: Dependencies {
         latestRecordTimestamp = latestRecord.timestamp
     }
 
-    func copy(updatedRecords: [StoryMessageRecord], deletedRecordIds: [Int64], transaction: SDSAnyReadTransaction) throws -> Self {
+    func copy(updatedRecords: [StoryMessageRecord], deletedRecordIds: [Int64], transaction: SDSAnyReadTransaction) throws -> Self? {
         var updatedRecords = updatedRecords
         let records = self.records.lazy
-            .filter { deletedRecordIds.contains($0.id ?? 0) }
+            .filter { !deletedRecordIds.contains($0.id ?? 0) }
             .map { oldRecord in
                 if let idx = updatedRecords.firstIndex(where: { $0.id == oldRecord.id }) {
                     return updatedRecords.remove(at: idx)
@@ -91,6 +91,7 @@ struct IncomingStoryViewModel: Dependencies {
                     return oldRecord
                 }
             } + updatedRecords
+        guard !records.isEmpty else { return nil }
         return try .init(records: records, identifier: identifier, transaction: transaction)
     }
 }

Signal/src/ViewControllers/HomeView/Stories/StoriesViewController.swift

@@ -120,7 +120,7 @@ class StoriesViewController: OWSViewController {
             var changedIdentifiers = [UUID]()
 
             let newModels = Self.databaseStorage.read { transaction in
-                self.models.map { model in
+                self.models.compactMap { model in
                     guard let latestRecord = model.records.first else { return model }
 
                     let modelDeletedRowIds = model.recordIds.filter { deletedRowIds.contains($0) }

SignalServiceKit/src/Messages/OWSDisappearingMessagesJob.h

@@ -1,5 +1,5 @@
 //
-//  Copyright (c) 2021 Open Whisper Systems. All rights reserved.
+//  Copyright (c) 2022 Open Whisper Systems. All rights reserved.
 //
 
 NS_ASSUME_NONNULL_BEGIN
@@ -19,6 +19,8 @@ NS_ASSUME_NONNULL_BEGIN
                  expirationStartedAt:(uint64_t)expirationStartedAt
                          transaction:(SDSAnyWriteTransaction *_Nonnull)transaction;
 
+- (void)scheduleRunByTimestamp:(uint64_t)timestamp;
+
 // Clean up any messages that expired since last launch immediately
 // and continue cleaning in the background.
 - (void)startIfNecessary;

SignalServiceKit/src/Messages/OWSDisappearingMessagesJob.m

@@ -130,26 +130,52 @@ void AssertIsOnDisappearingMessagesQueue()
     return expirationCount;
 }
 
+- (NSUInteger)deleteExpiredStories
+{
+    AssertIsOnDisappearingMessagesQueue();
+
+    OWSBackgroundTask *_Nullable backgroundTask = [OWSBackgroundTask backgroundTaskWithLabelStr:__PRETTY_FUNCTION__];
+
+    __block NSUInteger expirationCount = 0;
+    DatabaseStorageWrite(self.databaseStorage, ^(SDSAnyWriteTransaction *transaction) {
+        expirationCount = [StoryManager deleteExpiredStoriesWithTransaction:transaction];
+    });
+
+    OWSLogDebug(@"Removed %lu expired stories", (unsigned long)expirationCount);
+
+    OWSAssertDebug(backgroundTask);
+    backgroundTask = nil;
+    return expirationCount;
+}
+
 // deletes any expired messages and schedules the next run.
 - (NSUInteger)runLoop
 {
     OWSLogVerbose(@"in runLoop");
     AssertIsOnDisappearingMessagesQueue();
 
-    NSUInteger deletedCount = [self deleteExpiredMessages];
+    NSUInteger deletedCount = [self deleteExpiredMessages] + [self deleteExpiredStories];
 
-    __block NSNumber *nextExpirationTimestampNumber;
-    [self.databaseStorage readWithBlock:^(SDSAnyReadTransaction *transaction) {
-        nextExpirationTimestampNumber =
-            [self.disappearingMessagesFinder nextExpirationTimestampWithTransaction:transaction];
-    } file:__FILE__ function:__FUNCTION__ line:__LINE__];
+    __block NSNumber *nextMessageExpirationTimestampNumber;
+    __block NSNumber *nextStoryExpirationTimestampNumber;
+    [self.databaseStorage
+        readWithBlock:^(SDSAnyReadTransaction *transaction) {
+            nextMessageExpirationTimestampNumber =
+                [self.disappearingMessagesFinder nextExpirationTimestampWithTransaction:transaction];
+            nextStoryExpirationTimestampNumber = [StoryManager nextExpirationTimestampWithTransaction:transaction];
+        }
+                 file:__FILE__
+             function:__FUNCTION__
+                 line:__LINE__];
 
-    if (!nextExpirationTimestampNumber) {
+    if (!nextMessageExpirationTimestampNumber && !nextStoryExpirationTimestampNumber) {
         OWSLogDebug(@"No more expiring messages.");
         return deletedCount;
     }
 
-    uint64_t nextExpirationAt = nextExpirationTimestampNumber.unsignedLongLongValue;
+    uint64_t nextMessageExpirationAt = nextMessageExpirationTimestampNumber.unsignedLongLongValue;
+    uint64_t nextStoryExpirationAt = nextStoryExpirationTimestampNumber.unsignedLongLongValue;
+    uint64_t nextExpirationAt = MIN(nextMessageExpirationAt, nextStoryExpirationAt);
     NSDate *nextEpirationDate = [NSDate ows_dateWithMillisecondsSince1970:nextExpirationAt];
     [self scheduleRunByDate:nextEpirationDate];
 
@@ -181,6 +207,11 @@ void AssertIsOnDisappearingMessagesQueue()
     }];
 }
 
+- (void)scheduleRunByTimestamp:(uint64_t)timestamp
+{
+    [self scheduleRunByDate:[NSDate ows_dateWithMillisecondsSince1970:timestamp]];
+}
+
 #pragma mark -
 
 - (void)startIfNecessary

SignalServiceKit/src/Messages/Stories/StoryFinder.swift

@@ -91,4 +91,44 @@ public enum StoryFinder {
             return nil
         }
     }
+
+    // The stories should be enumerated in order from "next to expire" to "last to expire".
+    public static func enumerateExpiredStories(transaction: GRDBReadTransaction, block: @escaping (StoryMessageRecord, UnsafeMutablePointer<ObjCBool>) -> Void) {
+
+        let sql = """
+            SELECT *
+            FROM \(StoryMessageRecord.databaseTableName)
+            WHERE timestamp <= \(Date().ows_millisecondsSince1970 - StoryManager.storyLifetime)
+            ORDER BY timestamp ASC
+        """
+
+        do {
+            let cursor = try StoryMessageRecord.fetchCursor(transaction.database, sql: sql)
+            while let record = try cursor.next() {
+                var stop: ObjCBool = false
+                block(record, &stop)
+                if stop.boolValue {
+                    return
+                }
+            }
+        } catch {
+            owsFail("error: \(error)")
+        }
+    }
+
+    public static func oldestTimestamp(transaction: GRDBReadTransaction) -> UInt64? {
+        let sql = """
+            SELECT timestamp
+            FROM \(StoryMessageRecord.databaseTableName)
+            ORDER BY timestamp ASC
+            LIMIT 1
+        """
+
+        do {
+            return try UInt64.fetchOne(transaction.database, sql: sql)
+        } catch {
+            owsFailDebug("failed to lookup next story expiration \(error)")
+            return nil
+        }
+    }
 }

SignalServiceKit/src/Messages/Stories/StoryManager.swift

@@ -6,6 +6,8 @@ import Foundation
 
 @objc
 public class StoryManager: NSObject {
+    public static let storyLifetime = kDayInMs
+
     @objc
     public class func processIncomingStoryMessage(
         _ storyMessage: SSKProtoStoryMessage,
@@ -22,5 +24,28 @@ public class StoryManager: NSObject {
 
         // TODO: Optimistic downloading of story attachments.
         attachmentDownloads.enqueueDownloadOfAttachmentsForNewStoryMessage(record, transaction: transaction)
+
+        OWSDisappearingMessagesJob.shared.scheduleRun(byTimestamp: record.timestamp + storyLifetime)
+    }
+
+    @objc
+    public class func deleteExpiredStories(transaction: SDSAnyWriteTransaction) -> UInt {
+        var removedCount: UInt = 0
+        StoryFinder.enumerateExpiredStories(transaction: transaction.unwrapGrdbRead) { record, _ in
+            Logger.info("Removing StoryMessage \(record.timestamp) which expired at: \(record.timestamp + storyLifetime)")
+            do {
+                try record.delete(transaction.unwrapGrdbWrite.database)
+                removedCount += 1
+            } catch {
+                owsFailDebug("Failed to remove expired story with timestamp \(record.timestamp) \(error)")
+            }
+        }
+        return removedCount
+    }
+
+    @objc
+    public class func nextExpirationTimestamp(transaction: SDSAnyReadTransaction) -> NSNumber? {
+        guard let timestamp = StoryFinder.oldestTimestamp(transaction: transaction.unwrapGrdbRead) else { return nil }
+        return NSNumber(value: timestamp + storyLifetime)
     }
 }

SignalServiceKit/src/Messages/Stories/StoryMessageRecord.swift

@@ -116,6 +116,12 @@ public struct StoryMessageRecord: Dependencies, Codable, Identifiable, Fetchable
     mutating public func didInsert(with rowID: Int64, for column: String?) {
         self.id = rowID
     }
+
+    @discardableResult
+    public func delete(_ db: Database) throws -> Bool {
+        // TODO: Cleanup associated records
+        return try performDelete(db)
+    }
 }
 
 public enum StoryManifest: Codable {