From 440681320a76e671419e4b5efd8dae3d48ce0764 Mon Sep 17 00:00:00 2001 From: Nikolay Edigaryev Date: Fri, 12 Aug 2022 16:43:04 +0300 Subject: [PATCH] Introduce "tart prune" command (#164) --- Package.resolved | 9 ++ Package.swift | 2 + Sources/tart/Commands/List.swift | 2 +- Sources/tart/Commands/Prune.swift | 101 +++++++++++++++++++++++ Sources/tart/IPSWCache.swift | 20 +++++ Sources/tart/Prunable.swift | 11 +++ Sources/tart/Root.swift | 1 + Sources/tart/URL+AccessDate.swift | 25 ++++++ Sources/tart/URL+Prunable.swift | 11 +++ Sources/tart/VM.swift | 7 +- Sources/tart/VMDirectory.swift | 14 +++- Sources/tart/VMStorageOCI.swift | 30 +++++-- Tests/TartTests/URLAccessDateTests.swift | 22 +++++ 13 files changed, 243 insertions(+), 12 deletions(-) create mode 100644 Sources/tart/Commands/Prune.swift create mode 100644 Sources/tart/IPSWCache.swift create mode 100644 Sources/tart/Prunable.swift create mode 100644 Sources/tart/URL+AccessDate.swift create mode 100644 Sources/tart/URL+Prunable.swift create mode 100644 Tests/TartTests/URLAccessDateTests.swift diff --git a/Package.resolved b/Package.resolved index 3c84d2c4..caf85107 100644 --- a/Package.resolved +++ b/Package.resolved @@ -126,6 +126,15 @@ "version" : "0.9.2" } }, + { + "identity" : "swiftdate", + "kind" : "remoteSourceControl", + "location" : "https://github.com/malcommac/SwiftDate", + "state" : { + "revision" : "6190d0cefff3013e77ed567e6b074f324e5c5bf5", + "version" : "6.3.1" + } + }, { "identity" : "xctest-dynamic-overlay", "kind" : "remoteSourceControl", diff --git a/Package.swift b/Package.swift index 1b4abe91..e719352f 100644 --- a/Package.swift +++ b/Package.swift @@ -15,6 +15,7 @@ let package = Package( .package(url: "https://github.com/pointfreeco/swift-parsing", from: "0.9.2"), .package(url: "https://github.com/swift-server/async-http-client", from: "1.11.4"), .package(url: "https://github.com/apple/swift-algorithms", from: "1.0.0"), + .package(url: "https://github.com/malcommac/SwiftDate", from: "6.3.1") ], targets: [ .executableTarget(name: "tart", dependencies: [ @@ -23,6 +24,7 @@ let package = Package( .product(name: "AsyncHTTPClient", package: "async-http-client"), .product(name: "Dynamic", package: "Dynamic"), .product(name: "Parsing", package: "swift-parsing"), + .product(name: "SwiftDate", package: "SwiftDate"), ]), .testTarget(name: "TartTests", dependencies: ["tart"]) ] diff --git a/Sources/tart/Commands/List.swift b/Sources/tart/Commands/List.swift index 4e73b239..69a71e1b 100644 --- a/Sources/tart/Commands/List.swift +++ b/Sources/tart/Commands/List.swift @@ -10,7 +10,7 @@ struct List: AsyncParsableCommand { print("Source\tName") displayTable("local", try VMStorageLocal().list()) - displayTable("oci", try VMStorageOCI().list()) + displayTable("oci", try VMStorageOCI().list().map { (name, vmDir, _) in (name, vmDir) }) Foundation.exit(0) } catch { diff --git a/Sources/tart/Commands/Prune.swift b/Sources/tart/Commands/Prune.swift new file mode 100644 index 00000000..39037d8c --- /dev/null +++ b/Sources/tart/Commands/Prune.swift @@ -0,0 +1,101 @@ +import ArgumentParser +import Dispatch +import SwiftUI +import SwiftDate + +struct Prune: AsyncParsableCommand { + static var configuration = CommandConfiguration(abstract: "Prune OCI and IPSW caches") + + @Option(help: ArgumentHelp("Remove cache entries last accessed more than n days ago", + discussion: "For example, --older-than=7 will remove entries that weren't accessed by Tart in the last 7 days.", + valueName: "n")) + var olderThan: UInt? + + @Option(help: ArgumentHelp("Remove least recently used cache entries that do not fit the specified cache size budget n, expressed in gigabytes", + discussion: "For example, --cache-budget=50 will effectively shrink all caches to a total size of 50 gigabytes.", + valueName: "n")) + var cacheBudget: UInt? + + func validate() throws { + if olderThan == nil && cacheBudget == nil { + throw ValidationError("at least one criteria must be specified") + } + } + + func run() async throws { + do { + // Clean up cache entries based on last accessed date + if let olderThan = olderThan { + let olderThanInterval = Int(exactly: olderThan)!.days.timeInterval + let olderThanDate = Date().addingTimeInterval(olderThanInterval) + + try Prune.pruneOlderThan(olderThanDate: olderThanDate) + } + + // Clean up cache entries based on imposed cache size limit and entry's last accessed date + if let cacheBudget = cacheBudget { + try Prune.pruneCacheBudget(cacheBudgetBytes: UInt64(cacheBudget) * 1024 * 1024 * 1024) + } + + Foundation.exit(0) + } catch { + print(error) + + Foundation.exit(1) + } + } + + static func pruneOlderThan(olderThanDate: Date) throws { + let prunableStorages: [PrunableStorage] = [VMStorageOCI(), try IPSWCache()] + let prunables: [Prunable] = try prunableStorages.flatMap { try $0.prunables() } + + try prunables.filter { try $0.accessDate() <= olderThanDate }.forEach { try $0.delete() } + } + + static func pruneCacheBudget(cacheBudgetBytes: UInt64) throws { + let prunableStorages: [PrunableStorage] = [VMStorageOCI(), try IPSWCache()] + let prunables: [Prunable] = try prunableStorages + .flatMap { try $0.prunables() } + .sorted { try $0.accessDate() < $1.accessDate() } + + let cacheUsedBytes = try prunables.map { try $0.sizeBytes() }.reduce(0, +) + var cacheReclaimedBytes: Int = 0 + + var it = prunables.makeIterator() + + while (cacheUsedBytes - cacheReclaimedBytes) > cacheBudgetBytes { + guard let prunable = it.next() else { + break + } + + cacheReclaimedBytes -= try prunable.sizeBytes() + try prunable.delete() + } + } + + static func pruneReclaim(reclaimBytes: UInt64) throws { + let prunableStorages: [PrunableStorage] = [VMStorageOCI(), try IPSWCache()] + let prunables: [Prunable] = try prunableStorages + .flatMap { try $0.prunables() } + .sorted { try $0.accessDate() < $1.accessDate() } + + // Does it even make sense to start? + let cacheUsedBytes = try prunables.map { try $0.sizeBytes() }.reduce(0, +) + if cacheUsedBytes < reclaimBytes { + return + } + + var cacheReclaimedBytes: Int = 0 + + var it = prunables.makeIterator() + + while cacheReclaimedBytes <= reclaimBytes { + guard let prunable = it.next() else { + break + } + + cacheReclaimedBytes -= try prunable.sizeBytes() + try prunable.delete() + } + } +} diff --git a/Sources/tart/IPSWCache.swift b/Sources/tart/IPSWCache.swift new file mode 100644 index 00000000..17ae7522 --- /dev/null +++ b/Sources/tart/IPSWCache.swift @@ -0,0 +1,20 @@ +import Foundation +import Virtualization + +class IPSWCache: PrunableStorage { + let baseURL: URL + + init() throws { + baseURL = Config().tartCacheDir.appendingPathComponent("IPSWs", isDirectory: true) + try FileManager.default.createDirectory(at: baseURL, withIntermediateDirectories: true) + } + + func locationFor(image: VZMacOSRestoreImage) -> URL { + baseURL.appendingPathComponent("\(image.buildVersion).ipsw", isDirectory: false) + } + + func prunables() throws -> [Prunable] { + try FileManager.default.contentsOfDirectory(at: baseURL, includingPropertiesForKeys: nil) + .filter { $0.lastPathComponent.hasSuffix(".ipsw")} + } +} diff --git a/Sources/tart/Prunable.swift b/Sources/tart/Prunable.swift new file mode 100644 index 00000000..7ba24914 --- /dev/null +++ b/Sources/tart/Prunable.swift @@ -0,0 +1,11 @@ +import Foundation + +protocol PrunableStorage { + func prunables() throws -> [Prunable] +} + +protocol Prunable { + func delete() throws + func accessDate() throws -> Date + func sizeBytes() throws -> Int +} diff --git a/Sources/tart/Root.swift b/Sources/tart/Root.swift index 159ee4d0..20a0b694 100644 --- a/Sources/tart/Root.swift +++ b/Sources/tart/Root.swift @@ -16,6 +16,7 @@ struct Root: AsyncParsableCommand { IP.self, Pull.self, Push.self, + Prune.self, Delete.self, ]) diff --git a/Sources/tart/URL+AccessDate.swift b/Sources/tart/URL+AccessDate.swift new file mode 100644 index 00000000..18517daf --- /dev/null +++ b/Sources/tart/URL+AccessDate.swift @@ -0,0 +1,25 @@ +import Foundation + +extension URL { + func accessDate() throws -> Date { + let attrs = try resourceValues(forKeys: [.contentAccessDateKey]) + return attrs.contentAccessDate! + } + + func updateAccessDate(_ accessDate: Date = Date()) throws { + let attrs = try resourceValues(forKeys: [.contentAccessDateKey]) + let modificationDate = attrs.contentAccessDate! + + let times = [accessDate.asTimeval(), modificationDate.asTimeval()] + let ret = utimes(path, times) + if ret != 0 { + throw RuntimeError("utimes(2) failed: \(ret.explanation())") + } + } +} + +extension Date { + func asTimeval() -> timeval { + timeval(tv_sec: timeIntervalSince1970.toUnit(.second)!, tv_usec: 0) + } +} diff --git a/Sources/tart/URL+Prunable.swift b/Sources/tart/URL+Prunable.swift new file mode 100644 index 00000000..49daa086 --- /dev/null +++ b/Sources/tart/URL+Prunable.swift @@ -0,0 +1,11 @@ +import Foundation + +extension URL: Prunable { + func delete() throws { + try FileManager.default.removeItem(at: self) + } + + func sizeBytes() throws -> Int { + try resourceValues(forKeys: [.totalFileAllocatedSizeKey]).totalFileAllocatedSize! + } +} diff --git a/Sources/tart/VM.swift b/Sources/tart/VM.swift index edd4ffbd..9864275e 100644 --- a/Sources/tart/VM.swift +++ b/Sources/tart/VM.swift @@ -55,14 +55,11 @@ class VM: NSObject, VZVirtualMachineDelegate, ObservableObject { } } - - let ipswCacheFolder = Config().tartCacheDir.appendingPathComponent("IPSWs", isDirectory: true) - try FileManager.default.createDirectory(at: ipswCacheFolder, withIntermediateDirectories: true) - - let expectedIPSWLocation = ipswCacheFolder.appendingPathComponent("\(image.buildVersion).ipsw", isDirectory: false) + let expectedIPSWLocation = try IPSWCache().locationFor(image: image) if FileManager.default.fileExists(atPath: expectedIPSWLocation.path) { defaultLogger.appendNewLine("Using cached *.ipsw file...") + try expectedIPSWLocation.updateAccessDate() return expectedIPSWLocation } diff --git a/Sources/tart/VMDirectory.swift b/Sources/tart/VMDirectory.swift index 13eb7432..cdd2ebf3 100644 --- a/Sources/tart/VMDirectory.swift +++ b/Sources/tart/VMDirectory.swift @@ -7,7 +7,7 @@ struct UninitializedVMDirectoryError: Error { struct AlreadyInitializedVMDirectoryError: Error { } -struct VMDirectory { +struct VMDirectory: Prunable { var baseURL: URL var configURL: URL { @@ -77,4 +77,16 @@ struct VMDirectory { try diskFileHandle.truncate(atOffset: UInt64(sizeGB) * 1000 * 1000 * 1000) try diskFileHandle.close() } + + func delete() throws { + try FileManager.default.removeItem(at: baseURL) + } + + func accessDate() throws -> Date { + try baseURL.accessDate() + } + + func sizeBytes() throws -> Int { + try configURL.sizeBytes() + diskURL.sizeBytes() + nvramURL.sizeBytes() + } } diff --git a/Sources/tart/VMStorageOCI.swift b/Sources/tart/VMStorageOCI.swift index 48fc79d2..5c0890f9 100644 --- a/Sources/tart/VMStorageOCI.swift +++ b/Sources/tart/VMStorageOCI.swift @@ -1,6 +1,6 @@ import Foundation -class VMStorageOCI { +class VMStorageOCI: PrunableStorage { let baseURL = Config().tartCacheDir.appendingPathComponent("OCIs", isDirectory: true) private func vmURL(_ name: RemoteName) -> URL { @@ -16,6 +16,8 @@ class VMStorageOCI { try vmDir.validate() + try vmDir.baseURL.updateAccessDate() + return vmDir } @@ -42,8 +44,8 @@ class VMStorageOCI { try FileManager.default.removeItem(at: vmURL(name)) } - func list() throws -> [(String, VMDirectory)] { - var result: [(String, VMDirectory)] = Array() + func list() throws -> [(String, VMDirectory, Bool)] { + var result: [(String, VMDirectory, Bool)] = Array() guard let enumerator = FileManager.default.enumerator(at: baseURL, includingPropertiesForKeys: [.isSymbolicLinkKey], options: [.producesRelativePathURLs]) else { @@ -60,18 +62,23 @@ class VMStorageOCI { let parts = [foundURL.deletingLastPathComponent().relativePath, foundURL.lastPathComponent] var name: String - if try foundURL.resourceValues(forKeys: [.isSymbolicLinkKey]).isSymbolicLink! { + let isSymlink = try foundURL.resourceValues(forKeys: [.isSymbolicLinkKey]).isSymbolicLink! + if isSymlink { name = parts.joined(separator: ":") } else { name = parts.joined(separator: "@") } - result.append((name, vmDir)) + result.append((name, vmDir, isSymlink)) } return result } + func prunables() throws -> [Prunable] { + try list().filter { (_, _, isSymlink) in !isSymlink }.map { (_, vmDir, _) in vmDir } + } + func pull(_ name: RemoteName, registry: Registry) async throws { defaultLogger.appendNewLine("pulling manifest...") @@ -82,6 +89,19 @@ class VMStorageOCI { if !exists(digestName) { let tmpVMDir = try VMDirectory.temporary() + + // Try to reclaim some cache space if we know the VM size in advance + if let uncompressedDiskSize = manifest.uncompressedDiskSize() { + let requiredCapacityBytes = UInt64(uncompressedDiskSize + 128 * 1024 * 1024) + + let attrs = try tmpVMDir.baseURL.resourceValues(forKeys: [.volumeAvailableCapacityForImportantUsageKey]) + let availableCapacityBytes = UInt64(attrs.volumeAvailableCapacityForImportantUsage!) + + if availableCapacityBytes < requiredCapacityBytes { + try Prune.pruneReclaim(reclaimBytes: requiredCapacityBytes - availableCapacityBytes) + } + } + try await withTaskCancellationHandler(operation: { try await tmpVMDir.pullFromRegistry(registry: registry, manifest: manifest) try move(digestName, from: tmpVMDir) diff --git a/Tests/TartTests/URLAccessDateTests.swift b/Tests/TartTests/URLAccessDateTests.swift new file mode 100644 index 00000000..638f9d58 --- /dev/null +++ b/Tests/TartTests/URLAccessDateTests.swift @@ -0,0 +1,22 @@ +import XCTest +@testable import tart + +final class URLAccessDateTests: XCTestCase { + func testGetAndSetAccessTime() throws { + // Create a temporary file + let tmpDir = URL(fileURLWithPath: NSTemporaryDirectory(), isDirectory: true) + var tmpFile = tmpDir.appendingPathComponent(UUID().uuidString) + FileManager.default.createFile(atPath: tmpFile.path, contents: nil) + + // Ensure it's access date is different than our desired access date + let arbitraryDate = Date.init(year: 2008, month: 09, day: 28, hour: 23, minute: 15) + XCTAssertNotEqual(arbitraryDate, try tmpFile.accessDate()) + + // Set our desired access date for a file + try tmpFile.updateAccessDate(arbitraryDate) + + // Ensure the access date has changed to our value + tmpFile.removeCachedResourceValue(forKey: .contentAccessDateKey) + XCTAssertEqual(arbitraryDate, try tmpFile.accessDate()) + } +}