Added manual snapshotting to create a new file tree with the current dataset

Author: dimitribouniol

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

@@ -1498,3 +1498,77 @@ extension DiskPersistence.Datastore.Index {
         return (manifest: manifest, removedPages: removedPages)
     }
 }
+
+// MARK: - Snapshotting
+
+extension DiskPersistence.Datastore.Index {
+    @discardableResult
+    func copy(
+        into newDatastore: DiskPersistence<ReadWrite>.Datastore,
+        rootObjectManifest: inout DatastoreRootManifest,
+        targetPageSize: Int
+    ) async throws -> DiskPersistence<ReadWrite>.Datastore.Index {
+        let actualPageSize = min(max(targetPageSize, Configuration.minimumPageSize), Configuration.maximumPageSize) - DiskPersistence.Datastore.Page.headerSize
+        
+        var newIndexManifest = DatastoreIndexManifest(
+            id: self.id.manifestID,
+            orderedPages: []
+        )
+        
+        let stream = AsyncThrowingBackpressureStream { continuation in
+            try await self.forwardScanEntries(after: self.firstInsertionCursor) { entry in
+                try await continuation.yield(entry)
+                return true
+            }
+        }
+        
+        var currentPageBlocks: [DatastorePageEntryBlock] = []
+        var remainingSpace = actualPageSize
+        
+        var count = 0
+        
+        for try await entry in stream {
+            count += 1
+            let blocks = entry.blocks(remainingPageSpace: remainingSpace, maxPageSpace: actualPageSize)
+            
+            for block in blocks {
+                let encodedSize = block.encodedSize
+                if encodedSize > remainingSpace {
+                    let newPage = DiskPersistence<ReadWrite>.Datastore.Page(
+                        datastore: newDatastore,
+                        id: .init(index: id, page: .init()),
+                        blocks: currentPageBlocks
+                    )
+                    newIndexManifest.orderedPages.append(.added(newPage.id.page))
+                    try await newPage.persistIfNeeded()
+                    
+                    currentPageBlocks.removeAll(keepingCapacity: true)
+                    remainingSpace = actualPageSize
+                }
+                
+                remainingSpace -= encodedSize
+                currentPageBlocks.append(block)
+            }
+        }
+        
+        if !currentPageBlocks.isEmpty {
+            let newPage = DiskPersistence<ReadWrite>.Datastore.Page(
+                datastore: newDatastore,
+                id: .init(index: id, page: .init()),
+                blocks: currentPageBlocks
+            )
+            newIndexManifest.orderedPages.append(.added(newPage.id.page))
+            try await newPage.persistIfNeeded()
+        }
+        
+        rootObjectManifest.addedIndexes.insert(.init(id))
+        rootObjectManifest.addedIndexManifests.insert(.init(id))
+        let newIndex = DiskPersistence<ReadWrite>.Datastore.Index(
+            datastore: newDatastore,
+            id: id,
+            manifest: newIndexManifest
+        )
+        try await newIndex.persistIfNeeded()
+        return newIndex
+    }
+}

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

@@ -66,6 +66,14 @@ extension DatastoreRootManifest {
         case primary
         case direct(index: DatastoreIndexIdentifier)
         case secondary(index: DatastoreIndexIdentifier)
+        
+        init(_ id: PersistenceDatastoreIndexID) {
+            switch id {
+            case .primary:                  self = .primary
+            case .direct(let index, _):     self = .direct(index: index)
+            case .secondary(let index, _):  self = .secondary(index: index)
+            }
+        }
     }
 
     enum IndexManifestID: Codable, Hashable {

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

@@ -279,3 +279,80 @@ extension DiskPersistence.Datastore {
 
     var hasObservers: Bool { !observers.isEmpty }
 }
+
+// MARK: - Snapshotting
+
+extension DiskPersistence.Datastore {
+    @discardableResult
+    func copy(
+        rootIdentifier: DatastoreRootIdentifier?,
+        datastoreKey: DatastoreKey,
+        into newSnapshot: Snapshot<ReadWrite>,
+        iteration: inout SnapshotIteration,
+        targetPageSize: Int
+    ) async throws -> DiskPersistence<ReadWrite>.Datastore {
+        let newDatastore = DiskPersistence<ReadWrite>.Datastore(id: id, snapshot: newSnapshot)
+        
+        /// Copy the datastore over.
+        iteration.dataStores[datastoreKey] = SnapshotIteration.DatastoreInfo(key: datastoreKey, id: id, root: rootIdentifier)
+        
+        /// Record the addition of a new datastore since it lives in a new location on disk.
+        iteration.addedDatastores.insert(id)
+        
+        /// Stop here if there are no roots for the datastore — there is nothing else to migrate.
+        guard let datastoreRootIdentifier = rootIdentifier
+        else { return newDatastore }
+        
+        /// Record the addition of a new datastore since it lives in a new location on disk.
+        iteration.addedDatastoreRoots.insert(DatastoreRootReference(
+            datastoreID: id,
+            datastoreRootID: datastoreRootIdentifier
+        ))
+        
+        let rootObject = rootObject(for: datastoreRootIdentifier)
+        let rootObjectManifest = try await rootObject.manifest
+        var newRootObjectManifest = DatastoreRootManifest(
+            id: rootObjectManifest.id,
+            modificationDate: rootObjectManifest.modificationDate,
+            descriptor: rootObjectManifest.descriptor,
+            primaryIndexManifest: rootObjectManifest.primaryIndexManifest,
+            directIndexManifests: rootObjectManifest.directIndexManifests,
+            secondaryIndexManifests: rootObjectManifest.secondaryIndexManifests,
+            addedIndexes: [],
+            removedIndexes: [],
+            addedIndexManifests: [],
+            removedIndexManifests: []
+        )
+        
+        try await rootObject.primaryIndex.copy(
+            into: newDatastore,
+            rootObjectManifest: &newRootObjectManifest,
+            targetPageSize: targetPageSize
+        )
+        
+        for index in try await rootObject.directIndexes.values {
+            try await index.copy(
+                into: newDatastore,
+                rootObjectManifest: &newRootObjectManifest,
+                targetPageSize: targetPageSize
+            )
+        }
+        
+        for index in try await rootObject.secondaryIndexes.values {
+            try await index.copy(
+                into: newDatastore,
+                rootObjectManifest: &newRootObjectManifest,
+                targetPageSize: targetPageSize
+            )
+        }
+        
+        let newRootObject = DiskPersistence<ReadWrite>.Datastore.RootObject(
+            datastore: newDatastore,
+            id: datastoreRootIdentifier,
+            rootObject: newRootObjectManifest
+        )
+        try await newRootObject.persistIfNeeded()
+        
+        return newDatastore
+    }
+}

Sources/CodableDatastore/Persistence/Disk Persistence/DatedIdentifier.swift

@@ -72,6 +72,21 @@ struct DatedIdentifierComponents {
         "\(hour)-\(minute)"
     }
 
+    var date: Date? {
+        DateComponents(
+            calendar: Calendar(identifier: .gregorian),
+            timeZone: TimeZone(secondsFromGMT: 0),
+            year: Int(year),
+            month: Int(month),
+            day: Int(day),
+            hour: Int(hour),
+            minute: Int(minute),
+            second: Int(second),
+            nanosecond: Int(millisecond).map { $0*1_000_000 }
+        )
+        .date
+    }
+    
     static let size = 40
 }
 

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

@@ -284,8 +284,16 @@ extension DiskPersistence {
 
 extension DiskPersistence {
     /// Load the default snapshot from disk, or create an empty one if such a file does not exist.
-    private func loadSnapshot(from storeInfo: StoreInfo) -> Snapshot<AccessMode> {
-        let snapshotID = storeInfo.currentSnapshot ?? SnapshotIdentifier()
+    ///
+    /// - Parameters:
+    ///   - storeInfo: The store infor to load from.
+    ///   - newSnapshotIdentifier: A new snapshot identifier to use if the store doesn't have one yet. If nil, a new one will be created automatically.
+    /// - Returns: A snapshot to start using.
+    private func loadSnapshot(
+        from storeInfo: StoreInfo,
+        newSnapshotIdentifier: SnapshotIdentifier? = nil
+    ) -> Snapshot<AccessMode> {
+        let snapshotID = storeInfo.currentSnapshot ?? newSnapshotIdentifier ?? SnapshotIdentifier()
 
         if let snapshot = snapshots[snapshotID] {
             return snapshot
@@ -378,6 +386,69 @@ extension DiskPersistence {
             try await updateCurrentSnapshot(accessor: escapingClosure).value
         }
     }
+    
+    /// Create a new snapshot from the current snapshot the persistence is pointing to.
+    ///
+    /// This method is temporarily publc — once all the components are public, you should use them directly before the next minor version.
+    public func _takeSnapshot() async throws where AccessMode == ReadWrite {
+        try await _takeSnapshot(newSnapshotIdentifier: nil)
+    }
+    
+    func _takeSnapshot(
+        newSnapshotIdentifier: SnapshotIdentifier?
+    ) async throws where AccessMode == ReadWrite { // TODO: return new snapshot iteration
+        let readSnapshot = try await currentSnapshot
+        let newSnapshot = try await createSnapshot(from: readSnapshot, newSnapshotIdentifier: newSnapshotIdentifier)
+        try await setCurrentSnapshot(snapshot: newSnapshot)
+    }
+    
+    /// Load the current snapshot the persistence is reading and writing to.
+    var currentSnapshot: Snapshot<AccessMode> {
+        // TODO: This should return a readonly snapshot, but we need to be able to make a read-only copy from the persistence first.
+        get async throws {
+            try await withStoreInfo { await loadSnapshot(from: $0) }
+        }
+    }
+    
+    func setCurrentSnapshot(
+        snapshot: Snapshot<ReadWrite>,
+        dateUpdate: ModificationUpdate = .updateOnWrite
+    ) async throws where AccessMode == ReadWrite {
+        try await withStoreInfo { storeInfo in
+            guard snapshot.persistence === self
+            else { throw DiskPersistenceError.wrongPersistence }
+            
+            /// Update the store info with snapshot and modification date to use
+            storeInfo.currentSnapshot = snapshot.id
+            storeInfo.modificationDate = dateUpdate.modificationDate(for: storeInfo.modificationDate)
+        }
+    }
+    
+    func loadSnapshot(id: SnapshotIdentifier) async throws -> Snapshot<ReadOnly> {
+        preconditionFailure("Unimplemented")
+    }
+    
+//    var allSnapshots: AsyncStream<Snapshot<ReadOnly>> {
+//        /// Crawl the `Snapshots` directory, and collect all the .snapshot folders. Return the manifest structure.
+//        preconditionFailure("Unimplemented")
+//    }
+    
+    func createSnapshot(
+        from snapshot: Snapshot<ReadWrite>, // TODO: Shouldn't need to be readwrite
+        actionName: String? = nil,
+        newSnapshotIdentifier: SnapshotIdentifier? = nil
+    ) async throws -> Snapshot<ReadWrite> where AccessMode == ReadWrite {
+        let newSnapshot = try await snapshot.copy(
+            into: self,
+            actionName: actionName,
+            newSnapshotIdentifier: newSnapshotIdentifier,
+            targetPageSize: Configuration.defaultPageSize
+        )
+        
+        /// Save a reference to this new snapshot, and return it.
+        snapshots[newSnapshot.id] = newSnapshot
+        return newSnapshot
+    }
 }
 
 // MARK: - Persistence Creation

Sources/CodableDatastore/Persistence/Disk Persistence/DiskPersistenceError.swift

@@ -33,6 +33,9 @@ public enum DiskPersistenceError: LocalizedError, Equatable {
     /// The persistence is read only and cannot be written to.
     case cannotWrite
 
+    /// The persistence of another object is not in common with the persistence being used.
+    case wrongPersistence
+    
     public var errorDescription: String? {
         switch self {
         case .notFileURL:
@@ -49,6 +52,8 @@ public enum DiskPersistenceError: LocalizedError, Equatable {
             "The entry was in a format that could not be understood."
         case .cannotWrite:
             "The persistence is read only and cannot be written to."
+        case .wrongPersistence:
+            "The persistence of another object is not in common with the persistence being used."
         }
     }
 }

Sources/CodableDatastore/Persistence/Disk Persistence/Snapshot/Snapshot.swift

@@ -340,3 +340,59 @@ extension Snapshot {
         return (datastore, datastoreInfo.root)
     }
 }
+
+// MARK: - Snapshotting
+
+extension Snapshot {
+    @discardableResult
+    func copy(
+        into persistence: DiskPersistence<ReadWrite>,
+        actionName: String? = nil,
+        newSnapshotIdentifier: SnapshotIdentifier? = nil,
+        targetPageSize: Int
+    ) async throws -> Snapshot<ReadWrite> {
+        try await readingManifest { manifest, iteration in
+            /// Create a new snapshot and iteration to load data into
+            let newSnapshot = Snapshot<ReadWrite>(id: newSnapshotIdentifier ?? SnapshotIdentifier(), persistence: persistence)
+            
+            let creationDate = (try? newSnapshot.id.components)?.date ?? Date()
+            var newIteration = SnapshotIteration(
+                id: SnapshotIterationIdentifier(rawValue: newSnapshot.id.rawValue),
+                creationDate: creationDate,
+                precedingIteration: iteration.id,
+                precedingSnapshot: id,
+                successiveIterations: [],
+                actionName: actionName,
+                dataStores: [:],
+                addedDatastores: [],
+                removedDatastores: [],
+                addedDatastoreRoots: [],
+                removedDatastoreRoots: []
+            )
+            
+            /// Iterate through each datastore and copy the data over
+            for (_, datastoreInfo) in iteration.dataStores {
+                let (datastore, _) = await loadDatastore(for: datastoreInfo.key, from: iteration)
+                try await datastore.copy(
+                    rootIdentifier: datastoreInfo.root,
+                    datastoreKey: datastoreInfo.key,
+                    into: newSnapshot,
+                    iteration: &newIteration,
+                    targetPageSize: targetPageSize
+                )
+            }
+            
+            /// Create a new manifest with our written data.
+            let newManifest = SnapshotManifest(
+                id: newSnapshot.id,
+                modificationDate: creationDate,
+                currentIteration: newIteration.id
+            )
+            
+            /// Write the iteration and manifest records so the persistence is complete.
+            try await newSnapshot.write(iteration: newIteration)
+            try await newSnapshot.write(manifest: newManifest)
+            return newSnapshot
+        }
+    }
+}