From 68a3f51a269af385e69057fb9f35994d1aa92e91 Mon Sep 17 00:00:00 2001
From: Justin Mazzocchi <2831158+jzzocc@users.noreply.github.com>
Date: Mon, 14 Nov 2022 19:09:44 -0800
Subject: [PATCH] Limit number of records cleaned

---
 DB/Sources/DB/Content/ContentDatabase.swift | 73 +++++++++++++--------
 1 file changed, 46 insertions(+), 27 deletions(-)

diff --git a/DB/Sources/DB/Content/ContentDatabase.swift b/DB/Sources/DB/Content/ContentDatabase.swift
index 8c38a98..8dfe741 100644
--- a/DB/Sources/DB/Content/ContentDatabase.swift
+++ b/DB/Sources/DB/Content/ContentDatabase.swift
@@ -651,12 +651,54 @@ public extension ContentDatabase {
 
 private extension ContentDatabase {
     static let cleanAfterLastReadIdCount = 40
+    static let cleanLimit = 100
     static let ephemeralTimelines = NSCountedSet()
 
     static func fileURL(id: Identity.Id, appGroup: String) throws -> URL {
         try FileManager.default.databaseDirectoryURL(name: id.uuidString, appGroup: appGroup)
     }
 
+    static func statusIdsToDeleteForPositionPreservingClean(db: Database) throws -> Set<Status.Id> {
+        var statusIds = try Status.Id.fetchAll(
+            db,
+            TimelineStatusJoin.select(TimelineStatusJoin.Columns.statusId)
+                .order(TimelineStatusJoin.Columns.statusId.desc))
+
+        if let lastReadId = try Status.Id.fetchOne(
+            db,
+            LastReadIdRecord.filter(LastReadIdRecord.Columns.timelineId == Timeline.home.id)
+                .select(LastReadIdRecord.Columns.id))
+            ?? statusIds.first,
+           let index = statusIds.firstIndex(of: lastReadId) {
+            statusIds = Array(statusIds.prefix(index + Self.cleanAfterLastReadIdCount))
+        }
+
+        let reblogStatusIds = try Status.Id.fetchSet(
+            db,
+            StatusRecord.filter(statusIds.contains(StatusRecord.Columns.id)
+                                    && StatusRecord.Columns.reblogId != nil)
+                .select(StatusRecord.Columns.reblogId))
+
+        let statusIdsToKeep = Set(statusIds).union(reblogStatusIds)
+        let allStatusIds = try Status.Id.fetchSet(db, StatusRecord.select(StatusRecord.Columns.id))
+        let staleStatusIds = allStatusIds.subtracting(statusIdsToKeep)
+
+        return Set(Array(staleStatusIds).prefix(Self.cleanLimit))
+    }
+
+    static func accountIdsToDeleteForPositionPreservingClean(db: Database) throws -> Set<Account.Id> {
+        var accountIdsToKeep = try Account.Id.fetchSet(db, StatusRecord.select(StatusRecord.Columns.accountId))
+        accountIdsToKeep.formUnion(try Account.Id.fetchSet(
+            db,
+            AccountRecord.filter(accountIdsToKeep.contains(AccountRecord.Columns.id)
+                                    && AccountRecord.Columns.movedId != nil)
+                .select(AccountRecord.Columns.movedId)))
+        let allAccountIds = try Account.Id.fetchSet(db, AccountRecord.select(AccountRecord.Columns.id))
+        let staleAccountIds = allAccountIds.subtracting(accountIdsToKeep)
+
+        return Set(Array(staleAccountIds).prefix(Self.cleanLimit))
+    }
+
     static func clean(_ databaseWriter: DatabaseWriter,
                       useHomeTimelineLastReadId: Bool) throws {
         try databaseWriter.write {
@@ -668,33 +710,10 @@ private extension ContentDatabase {
 
             if useHomeTimelineLastReadId {
                 try TimelineRecord.filter(TimelineRecord.Columns.id != Timeline.home.id).deleteAll($0)
-                var statusIds = try Status.Id.fetchAll(
-                    $0,
-                    TimelineStatusJoin.select(TimelineStatusJoin.Columns.statusId)
-                        .order(TimelineStatusJoin.Columns.statusId.desc))
-
-                if let lastReadId = try Status.Id.fetchOne(
-                    $0,
-                    LastReadIdRecord.filter(LastReadIdRecord.Columns.timelineId == Timeline.home.id)
-                        .select(LastReadIdRecord.Columns.id))
-                    ?? statusIds.first,
-                   let index = statusIds.firstIndex(of: lastReadId) {
-                    statusIds = Array(statusIds.prefix(index + Self.cleanAfterLastReadIdCount))
-                }
-
-                statusIds += try Status.Id.fetchAll(
-                    $0,
-                    StatusRecord.filter(statusIds.contains(StatusRecord.Columns.id)
-                                            && StatusRecord.Columns.reblogId != nil)
-                        .select(StatusRecord.Columns.reblogId))
-                try StatusRecord.filter(!statusIds.contains(StatusRecord.Columns.id)).deleteAll($0)
-                var accountIds = try Account.Id.fetchAll($0, StatusRecord.select(StatusRecord.Columns.accountId))
-                accountIds += try Account.Id.fetchAll(
-                    $0,
-                    AccountRecord.filter(accountIds.contains(AccountRecord.Columns.id)
-                                            && AccountRecord.Columns.movedId != nil)
-                        .select(AccountRecord.Columns.movedId))
-                try AccountRecord.filter(!accountIds.contains(AccountRecord.Columns.id)).deleteAll($0)
+                try StatusRecord.filter(statusIdsToDeleteForPositionPreservingClean(db: $0)
+                    .contains(StatusRecord.Columns.id)).deleteAll($0)
+                try AccountRecord.filter(accountIdsToDeleteForPositionPreservingClean(db: $0)
+                    .contains(AccountRecord.Columns.id)).deleteAll($0)
             } else {
                 try TimelineRecord.deleteAll($0)
                 try StatusRecord.deleteAll($0)