Updated Disk Persistence to use a cache instead of keeping all objects in memory indefinitely.

Author: dimitribouniol

README.md

@@ -53,7 +53,6 @@ targets: [
 As this project matures towards release, the project will focus on the functionality and work listed below:
 - Force migration methods
 - Composite indexes (via macros?)
-- Cleaning up old resources in memory
 - Cleaning up old resources on disk
 - Reversed ranged reads
 - Controls for the edit history
@@ -63,6 +62,7 @@ As this project matures towards release, the project will focus on the functiona
 - An example app
 - A memory persistence useful for testing apps with
 - A pre-configured data store tuned to storing pure Data, useful for types like Images
+- Cleaning up memory leaks
 
 The above list will be kept up to date during development and will likely see additions during that process.
 

Sources/CodableDatastore/Persistence/Disk Persistence/Datastore/DatastoreIndex.swift

@@ -13,7 +13,7 @@ typealias DatastoreIndexIdentifier = TypedIdentifier<DiskPersistence<ReadOnly>.D
 
 extension DiskPersistence.Datastore {
     actor Index: Identifiable {
-        unowned let datastore: DiskPersistence<AccessMode>.Datastore
+        let datastore: DiskPersistence<AccessMode>.Datastore
 
         let id: ID
 
@@ -34,6 +34,12 @@ extension DiskPersistence.Datastore {
             self._manifest = manifest
             self.isPersisted = manifest == nil
         }
+        
+        deinit {
+            Task { [id, datastore] in
+                await datastore.invalidate(id)
+            }
+        }
     }
 }
 

Sources/CodableDatastore/Persistence/Disk Persistence/Datastore/DatastorePage.swift

@@ -14,7 +14,7 @@ typealias DatastorePageIdentifier = DatedIdentifier<DiskPersistence<ReadOnly>.Da
 
 extension DiskPersistence.Datastore {
     actor Page: Identifiable {
-        unowned let datastore: DiskPersistence<AccessMode>.Datastore
+        let datastore: DiskPersistence<AccessMode>.Datastore
 
         let id: ID
 
@@ -36,6 +36,12 @@ extension DiskPersistence.Datastore {
             }
             self.isPersisted = blocks == nil
         }
+        
+        deinit {
+            Task { [id, datastore] in
+                await datastore.invalidate(id)
+            }
+        }
     }
 }
 

Sources/CodableDatastore/Persistence/Disk Persistence/Datastore/DatastoreRoot.swift

@@ -12,7 +12,7 @@ typealias DatastoreRootIdentifier = DatedIdentifier<DiskPersistence<ReadOnly>.Da
 
 extension DiskPersistence.Datastore {
     actor RootObject: Identifiable {
-        unowned let datastore: DiskPersistence<AccessMode>.Datastore
+        let datastore: DiskPersistence<AccessMode>.Datastore
 
         let id: DatastoreRootIdentifier
 
@@ -30,6 +30,12 @@ extension DiskPersistence.Datastore {
             self._rootObject = rootObject
             self.isPersisted = rootObject == nil
         }
+        
+        deinit {
+            Task { [id, datastore] in
+                await datastore.invalidate(id)
+            }
+        }
     }
 }
 

Sources/CodableDatastore/Persistence/Disk Persistence/Datastore/PersistenceDatastore.swift

@@ -10,6 +10,14 @@ import Foundation
 
 typealias DatastoreIdentifier = TypedIdentifier<DiskPersistence<ReadOnly>.Datastore>
 
+struct WeakValue<T: AnyObject> {
+    weak var value: T?
+    
+    init(_ value: T) {
+        self.value = value
+    }
+}
+
 extension DiskPersistence {
     actor Datastore {
         let id: DatastoreIdentifier
@@ -21,9 +29,9 @@ extension DiskPersistence {
         var lastUpdateDescriptorTask: Task<Any, Error>?
 
         /// The root objects that are being tracked in memory.
-        var trackedRootObjects: [RootObject.ID : RootObject] = [:]
-        var trackedIndexes: [Index.ID : Index] = [:]
-        var trackedPages: [Page.ID : Page] = [:]
+        var trackedRootObjects: [RootObject.ID : WeakValue<RootObject>] = [:]
+        var trackedIndexes: [Index.ID : WeakValue<Index>] = [:]
+        var trackedPages: [Page.ID : WeakValue<Page>] = [:]
 
         /// The root objects on the file system that are actively loaded in memory.
         var loadedRootObjects: Set<RootObject.ID> = []
@@ -101,16 +109,23 @@ extension DiskPersistence.Datastore {
 
 extension DiskPersistence.Datastore {
     func rootObject(for identifier: RootObject.ID) -> RootObject {
-        if let rootObject = trackedRootObjects[identifier] {
+        if let rootObject = trackedRootObjects[identifier]?.value {
             return rootObject
         }
+//        print("🤷 Cache Miss: Root \(identifier)")
         let rootObject = RootObject(datastore: self, id: identifier)
-        trackedRootObjects[identifier] = rootObject
+        trackedRootObjects[identifier] = WeakValue(rootObject)
+        Task { await snapshot.persistence.cache(rootObject) }
         return rootObject
     }
 
     func adopt(rootObject: RootObject) {
-        trackedRootObjects[rootObject.id] = rootObject
+        trackedRootObjects[rootObject.id] = WeakValue(rootObject)
+        Task { await snapshot.persistence.cache(rootObject) }
+    }
+    
+    func invalidate(_ identifier: RootObject.ID) {
+        trackedRootObjects.removeValue(forKey: identifier)
     }
 
     func mark(identifier: RootObject.ID, asLoaded: Bool) {
@@ -122,16 +137,23 @@ extension DiskPersistence.Datastore {
     }
 
     func index(for identifier: Index.ID) -> Index {
-        if let index = trackedIndexes[identifier] {
+        if let index = trackedIndexes[identifier]?.value {
             return index
         }
+//        print("🤷 Cache Miss: Index \(identifier)")
         let index = Index(datastore: self, id: identifier)
-        trackedIndexes[identifier] = index
+        trackedIndexes[identifier] = WeakValue(index)
+        Task { await snapshot.persistence.cache(index) }
         return index
     }
 
     func adopt(index: Index) {
-        trackedIndexes[index.id] = index
+        trackedIndexes[index.id] = WeakValue(index)
+        Task { await snapshot.persistence.cache(index) }
+    }
+    
+    func invalidate(_ identifier: Index.ID) {
+        trackedIndexes.removeValue(forKey: identifier)
     }
 
     func mark(identifier: Index.ID, asLoaded: Bool) {
@@ -143,16 +165,23 @@ extension DiskPersistence.Datastore {
     }
 
     func page(for identifier: Page.ID) -> Page {
-        if let page = trackedPages[identifier.withoutManifest] {
+        if let page = trackedPages[identifier.withoutManifest]?.value {
             return page
         }
+//        print("🤷 Cache Miss: Page \(identifier.page)")
         let page = Page(datastore: self, id: identifier)
-        trackedPages[identifier.withoutManifest] = page
+        trackedPages[identifier.withoutManifest] = WeakValue(page)
+        Task { await snapshot.persistence.cache(page) }
         return page
     }
 
     func adopt(page: Page) {
-        trackedPages[page.id.withoutManifest] = page
+        trackedPages[page.id.withoutManifest] = WeakValue(page)
+        Task { await snapshot.persistence.cache(page) }
+    }
+    
+    func invalidate(_ identifier: Page.ID) {
+        trackedPages.removeValue(forKey: identifier.withoutManifest)
     }
 
     func mark(identifier: Page.ID, asLoaded: Bool) {

Sources/CodableDatastore/Persistence/Disk Persistence/DiskPersistence.swift

@@ -25,6 +25,16 @@ public actor DiskPersistence<AccessMode: _AccessMode>: Persistence {
 
     var lastTransaction: Transaction?
 
+    /// Shared caches across all snapshots and datastores.
+    var rollingRootObjectCacheIndex = 0
+    var rollingRootObjectCache: [Datastore.RootObject] = []
+    
+    var rollingIndexCacheIndex = 0
+    var rollingIndexCache: [Datastore.Index] = []
+    
+    var rollingPageCacheIndex = 0
+    var rollingPageCache: [Datastore.Page] = []
+    
     /// Initialize a ``DiskPersistence`` with a read-write URL.
     ///
     /// Use this initializer when creating a persistence from the main process that will access it, such as your app. To access the same persistence from another process, use ``init(readOnlyURL:)`` instead.
@@ -501,6 +511,43 @@ extension DiskPersistence {
     }
 }
 
+// MARK: - Persistence-wide Caches
+
+extension DiskPersistence {
+    func cache(_ rootObject: Datastore.RootObject) {
+        if rollingRootObjectCache.count <= rollingRootObjectCacheIndex {
+            rollingRootObjectCache.append(rootObject)
+        } else {
+            rollingRootObjectCache[rollingRootObjectCacheIndex] = rootObject
+        }
+        /// Limit cache to 16 recent root objects. We only really need one per active datastore.
+        /// Note more-recently accessed entries may be represented multiple times in the cache, and are more likely to survive.
+        rollingRootObjectCacheIndex = (rollingRootObjectCacheIndex + 1) % 16
+    }
+    
+    func cache(_ index: Datastore.Index) {
+        if rollingIndexCache.count <= rollingIndexCacheIndex {
+            rollingIndexCache.append(index)
+        } else {
+            rollingIndexCache[rollingIndexCacheIndex] = index
+        }
+        /// Limit cache to 128 recent indexes, which is 8 per datastore.
+        /// Note more-recently accessed entries may be represented multiple times in the cache, and are more likely to survive.
+        rollingIndexCacheIndex = (rollingIndexCacheIndex + 1) % 128
+    }
+    
+    func cache(_ page: Datastore.Page) {
+        if rollingPageCache.count <= rollingPageCacheIndex {
+            rollingPageCache.append(page)
+        } else {
+            rollingPageCache[rollingPageCacheIndex] = page
+        }
+        /// Limit cache to 4096 recent pages, which is up to 16MB.
+        /// Note more-recently accessed entries may be represented multiple times in the cache, and are more likely to survive.
+        rollingPageCacheIndex = (rollingPageCacheIndex + 1) % 4096
+    }
+}
+
 // MARK: - Helper Types
 
 class WeakDatastore {

Sources/CodableDatastore/Persistence/Disk Persistence/Transaction/Transaction.swift

@@ -124,6 +124,18 @@ extension DiskPersistence {
         }
 
         private func persist() async throws {
+            defer {
+                rootObjects.removeAll()
+                entryMutations.removeAll()
+                createdRootObjects.removeAll()
+                createdIndexes.removeAll()
+                createdPages.removeAll()
+                deletedRootObjects.removeAll()
+                deletedIndexes.removeAll()
+                deletedPages.removeAll()
+                childTransactions.removeAll()
+            }
+            
             if let parent {
                 try await parent.apply(
                     rootObjects: rootObjects,