Make stacked content pruning reference-aware

Author: yzhuang-oai

Sources/tart/Commands/Prune.swift

@@ -81,27 +81,34 @@ struct Prune: AsyncParsableCommand {
   }
 
   static func pruneSpaceBudget(prunableStorages: [PrunableStorage], spaceBudgetBytes: UInt64) throws {
-    let prunables: [Prunable] = try prunableStorages
-      .flatMap { try $0.prunables() }
-      .sorted { try $0.accessDate() > $1.accessDate() }
-
-    var spaceBudgetBytes = spaceBudgetBytes
-    var prunablesToDelete: [Prunable] = []
-
-    for prunable in prunables {
-      let prunableSizeBytes = UInt64(try prunable.allocatedSizeBytes())
-
-      if prunableSizeBytes <= spaceBudgetBytes {
-        // Don't mark for deletion as
-        // there's a budget available
-        spaceBudgetBytes -= prunableSizeBytes
-      } else {
-        // Mark for deletion
-        prunablesToDelete.append(prunable)
+    while true {
+      let prunables: [Prunable] = try prunableStorages
+        .flatMap { try $0.prunables() }
+        .sorted { try $0.accessDate() > $1.accessDate() }
+
+      var remainingBudgetBytes = spaceBudgetBytes
+      var prunableToDelete: Prunable?
+
+      for prunable in prunables {
+        let prunableSizeBytes = UInt64(try prunable.allocatedSizeBytes())
+
+        if prunableSizeBytes <= remainingBudgetBytes {
+          // Don't mark for deletion as there is budget available
+          remainingBudgetBytes -= prunableSizeBytes
+        } else {
+          prunableToDelete = prunable
+          break
+        }
+      }
+
+      guard let prunableToDelete else {
+        return
       }
-    }
 
-    try prunablesToDelete.forEach { try $0.delete() }
+      // Deleting one cached stacked image can change which remaining image
+      // owns shared immutable content. Rebuild before choosing another.
+      try prunableToDelete.delete()
+    }
   }
 
   static func reclaimIfNeeded(_ requiredBytes: UInt64, _ initiator: Prunable? = nil) throws {
@@ -145,46 +152,51 @@ struct Prune: AsyncParsableCommand {
     try Prune.reclaimIfPossible(requiredBytes - volumeAvailableCapacityCalculated, initiator)
   }
 
-  private static func reclaimIfPossible(_ reclaimBytes: UInt64, _ initiator: Prunable? = nil) throws {
+  static func reclaimIfPossible(_ reclaimBytes: UInt64, _ initiator: Prunable? = nil) throws {
     let span = OTel.shared.tracer.spanBuilder(spanName: "prune").startSpan()
     defer { span.end() }
 
     let prunableStorages: [PrunableStorage] = [try VMStorageOCI(), try IPSWCache()]
-    let prunables: [Prunable] = try prunableStorages
-      .flatMap { try $0.prunables() }
-      .sorted { try $0.accessDate() < $1.accessDate() }
+    let prunables = {
+      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.allocatedSizeBytes() }.reduce(0, +)
-    if cacheUsedBytes < reclaimBytes {
+    let initialPrunables = try prunables()
+    let initialCacheUsedBytes = try initialPrunables.map { try $0.allocatedSizeBytes() }.reduce(0, +)
+    guard let reclaimBytes = Int(exactly: reclaimBytes), initialCacheUsedBytes >= reclaimBytes else {
       return
     }
 
-    var cacheReclaimedBytes: Int = 0
-
-    var it = prunables.makeIterator()
+    let targetCacheUsedBytes = initialCacheUsedBytes - reclaimBytes
+    var currentCacheUsedBytes = initialCacheUsedBytes
+    let initiatorPath = initiator.map {
+      $0.url.resolvingSymlinksInPath().standardizedFileURL.path
+    }
 
-    while cacheReclaimedBytes <= reclaimBytes {
-      guard let prunable = it.next() else {
+    while currentCacheUsedBytes > targetCacheUsedBytes {
+      // Deleting one cached stacked image can transfer ownership of shared
+      // immutable content to another record without reclaiming those bytes.
+      // Rebuild the candidates after every deletion so automatic pruning
+      // measures the cache that remains rather than a stale ownership snapshot.
+      guard let prunable = try prunables().first(where: {
+        $0.url.resolvingSymlinksInPath().standardizedFileURL.path != initiatorPath
+      }) else {
         break
       }
 
-      if prunable.url == initiator?.url.resolvingSymlinksInPath() {
-        // do not prune the initiator
-        continue
-      }
-
       let allocatedSizeBytes = try prunable.allocatedSizeBytes()
 
       OpenTelemetry.instance.contextProvider.activeSpan?
         .addEvent(name: "Pruned \(allocatedSizeBytes) bytes for \(prunable.url.path)")
 
-      cacheReclaimedBytes += allocatedSizeBytes
-
       try prunable.delete()
+      currentCacheUsedBytes = try prunables().map { try $0.allocatedSizeBytes() }.reduce(0, +)
     }
 
     OpenTelemetry.instance.contextProvider.activeSpan?
-      .addEvent(name: "Reclaimed \(cacheReclaimedBytes) bytes")
+      .addEvent(name: "Reclaimed \(initialCacheUsedBytes - currentCacheUsedBytes) bytes")
   }
 }

Sources/tart/Commands/Run.swift

@@ -1039,20 +1039,24 @@ struct AdditionalDisk {
         // A cached stacked image has no writable top overlay. Create one in a
         // disposable directory for this additional-disk attachment.
         let temporaryVMDir = try VMDirectory.temporary()
-        try FileManager.default.copyItem(at: vmDir.configURL, to: temporaryVMDir.configURL)
-        try FileManager.default.copyItem(at: vmDir.nvramURL, to: temporaryVMDir.nvramURL)
-        try FileManager.default.copyItem(at: vmDir.manifestURL, to: temporaryVMDir.manifestURL)
-        let lock = try FileLock(lockURL: temporaryVMDir.baseURL)
-        try lock.lock()
+        let temporaryVMDirLock = try FileLock(lockURL: temporaryVMDir.baseURL)
+        try temporaryVMDirLock.lock()
+        try vmDir.cloneStacked(
+          to: temporaryVMDir,
+          copyWritableOverlay: false,
+          generateMAC: false
+        )
         let stack = try temporaryVMDir.diskImageStack()
-        try stack.createWritableOverlay()
         let attachment = try stack.makeAttachment(
           readOnly: diskReadOnly,
           cachingMode: try VZDiskImageCachingMode(cachingModeRaw) ?? .automatic,
           synchronizationMode: try VZDiskImageSynchronizationMode(syncModeRaw)
         )
 
-        return AdditionalDisk(configuration: VZVirtioBlockDeviceConfiguration(attachment: attachment), temporaryDiskLock: lock)
+        return AdditionalDisk(
+          configuration: VZVirtioBlockDeviceConfiguration(attachment: attachment),
+          temporaryDiskLock: temporaryVMDirLock
+        )
       }
 
       // Unfortunately, VZDiskImageStorageDeviceAttachment does not support

Sources/tart/Config.swift

@@ -33,7 +33,7 @@ struct Config {
         continue
       }
 
-      try FileManager.default.removeItem(at: entry)
+      try VMDirectory(baseURL: entry).removeFromDisk()
 
       try lock.unlock()
     }

Sources/tart/ContentStore.swift

@@ -16,6 +16,7 @@ struct ContentStore {
 
   let baseURL: URL
   private let digestDirectoryURL: URL
+  private let pruneLockURL: URL
 
   init() throws {
     try self.init(baseURL: Config().tartCacheDir.appendingPathComponent("content", isDirectory: true))
@@ -24,13 +25,41 @@ struct ContentStore {
   init(baseURL: URL) throws {
     self.baseURL = baseURL
     self.digestDirectoryURL = baseURL.appendingPathComponent(Self.digestAlgorithm, isDirectory: true)
+    self.pruneLockURL = baseURL.appendingPathComponent(".gc.lock")
     try FileManager.default.createDirectory(at: digestDirectoryURL, withIntermediateDirectories: true)
+    if !FileManager.default.fileExists(atPath: pruneLockURL.path) {
+      _ = FileManager.default.createFile(atPath: pruneLockURL.path, contents: Data())
+    }
+  }
+
+  /// Serializes reference publication with the final reference check and
+  /// deletion of immutable cache entries across Tart processes.
+  func withPruneLock<T>(_ body: () throws -> T) throws -> T {
+    let lock = try FileLock(lockURL: pruneLockURL)
+    try lock.lock()
+    defer { try? lock.unlock() }
+
+    return try body()
+  }
+
+  /// Waits for any prune already scanning references to finish. After this
+  /// returns, later prune runs can see a reference the caller already wrote.
+  func synchronizePublishedReferences() throws {
+    try withPruneLock {}
   }
 
   func contentURL(for contentDigest: String) throws -> URL {
+    try contentURL(for: contentDigest, under: baseURL)
+  }
+
+  /// Returns the canonical path for a digest under an arbitrary content-store
+  /// root without creating directories or lock files.
+  func contentURL(for contentDigest: String, under baseURL: URL) throws -> URL {
     let digestHex = try validatedDigestHex(contentDigest)
 
-    return digestDirectoryURL.appendingPathComponent(digestHex)
+    return baseURL
+      .appendingPathComponent(Self.digestAlgorithm, isDirectory: true)
+      .appendingPathComponent(digestHex)
   }
 
   func temporaryContentURL(for contentDigest: String) throws -> URL {
@@ -61,16 +90,18 @@ struct ContentStore {
     return lockURL
   }
 
-  /// Returns a digest-addressed entry without rereading it. Pull verifies
-  /// content hashes before accepting a cache hit; clone only needs a cheap
-  /// structural check, like Tart's existing disk.img path.
+  /// Returns an immutable digest-addressed entry without rereading it. Files
+  /// are verified when installed and when deciding whether a pull is a cache
+  /// hit; normal clone/run/push paths trust the store like Tart's disk.img.
   func contentURLIfPresent(for contentDigest: String) throws -> URL? {
     let url = try contentURL(for: contentDigest)
 
     guard FileManager.default.fileExists(atPath: url.path) else {
       return nil
     }
 
+    try url.updateAccessDate()
+
     return url
   }
 
@@ -88,6 +119,33 @@ struct ContentStore {
     return url
   }
 
+  /// Returns immutable content files that no retained cached image or local VM
+  /// references. Callers may prune these like other cache entries.
+  func prunables(excluding referencedContentDigests: Swift.Set<String>) throws -> [URL] {
+    guard let enumerator = FileManager.default.enumerator(
+      at: digestDirectoryURL,
+      includingPropertiesForKeys: [.isRegularFileKey],
+      options: [.skipsSubdirectoryDescendants]
+    ) else {
+      return []
+    }
+
+    return try enumerator.compactMap { element in
+      guard let url = element as? URL,
+            try url.resourceValues(forKeys: [.isRegularFileKey]).isRegularFile == true else {
+        return nil
+      }
+
+      let contentDigest = "\(Self.digestPrefix)\(url.lastPathComponent)"
+      guard (try? validatedDigestHex(contentDigest)) != nil,
+            !referencedContentDigests.contains(contentDigest) else {
+        return nil
+      }
+
+      return url
+    }
+  }
+
   /// Move a fully reconstructed temporary file into the cache after verifying
   /// its semantic identity. The caller should create the temporary file with
   /// temporaryContentURL(for:) or resumableContentURL(for:) so rename stays on

Sources/tart/OCI/Manifest.swift

@@ -37,6 +37,24 @@ struct TartDiskFileGroup: Equatable {
   /// Whole reconstructed-file digest. Existing flat manifests do not have
   /// this until a macOS 27 clone normalizes its local manifest copy.
   var contentDigest: String?
+
+  /// Expected size of the complete disk file reconstructed from these chunks.
+  func uncompressedSize() -> UInt64? {
+    var result: UInt64 = 0
+    for chunk in chunks {
+      guard let size = chunk.uncompressedSize() else {
+        return nil
+      }
+
+      let addition = result.addingReportingOverflow(size)
+      guard !addition.overflow else {
+        return nil
+      }
+      result = addition.partialValue
+    }
+
+    return result
+  }
 }
 
 enum TartDiskRepresentation: Equatable {
@@ -172,6 +190,16 @@ struct OCIManifest: Codable, Equatable {
     return .stacked(base: base, overlays: overlays)
   }
 
+  /// Returns content-store digests needed to reconstruct this disk stack.
+  func diskContentDigests() throws -> [String] {
+    switch try tartDiskRepresentation() {
+    case .flat(let base):
+      return base.contentDigest.map { [$0] } ?? []
+    case .stacked(let base, let overlays):
+      return ([base] + overlays).compactMap(\.contentDigest)
+    }
+  }
+
   private func validateChunkMetadata(_ chunks: [OCIManifestLayer]) throws {
     guard chunks.allSatisfy({ $0.uncompressedSize() != nil && $0.uncompressedContentDigest() != nil }) else {
       throw OCIManifestValidationError.invalidDiskMetadata("disk chunks need uncompressed size and content digest")

Sources/tart/VMDirectory+DiskImageStack.swift

@@ -1,6 +1,12 @@
 import Foundation
 
 extension VMDirectory {
+  /// Returns content-store digests needed to reconstruct this VM's disk stack.
+  func diskContentDigests() throws -> [String] {
+    let manifest = try OCIManifest(fromJSON: Data(contentsOf: manifestURL))
+    return try manifest.diskContentDigests()
+  }
+
   func diskImageStack(contentStore providedStore: ContentStore? = nil) throws -> DiskImageStack {
     let manifest = try OCIManifest(fromJSON: Data(contentsOf: manifestURL))
     let base: TartDiskFileGroup
@@ -42,13 +48,18 @@ extension VMDirectory {
     generateMAC: Bool,
     contentStore: ContentStore? = nil
   ) throws {
-    try FileManager.default.copyItem(at: configURL, to: destination.configURL)
-    try FileManager.default.copyItem(at: nvramURL, to: destination.nvramURL)
-    try FileManager.default.copyItem(at: manifestURL, to: destination.manifestURL)
+    let contentStore = try contentStore ?? ContentStore()
+    try contentStore.withPruneLock {
+      try FileManager.default.copyItem(at: configURL, to: destination.configURL)
+      try FileManager.default.copyItem(at: nvramURL, to: destination.nvramURL)
+      try FileManager.default.copyItem(at: manifestURL, to: destination.manifestURL)
+
+      if copyWritableOverlay {
+        try FileManager.default.copyItem(at: overlayURL, to: destination.overlayURL)
+      }
+    }
 
-    if copyWritableOverlay {
-      try FileManager.default.copyItem(at: overlayURL, to: destination.overlayURL)
-    } else {
+    if !copyWritableOverlay {
       try destination.diskImageStack(contentStore: contentStore).createWritableOverlay()
     }
 
@@ -67,17 +78,6 @@ extension VMDirectory {
     let contentDigest = try Digest.hash(diskURL)
     let contentStore = try providedStore ?? ContentStore()
 
-    if try contentStore.existingContentURL(for: contentDigest) == nil {
-      let temporaryURL = try contentStore.temporaryContentURL(for: contentDigest)
-      do {
-        try FileManager.default.copyItem(at: diskURL, to: temporaryURL)
-        _ = try contentStore.install(temporaryURL, contentDigest: contentDigest)
-      } catch {
-        try? FileManager.default.removeItem(at: temporaryURL)
-        throw error
-      }
-    }
-
     var manifest = try OCIManifest(fromJSON: Data(contentsOf: manifestURL))
     guard case .flat = try manifest.tartDiskRepresentation() else {
       throw RuntimeError.VMConfigurationError("--stacked cannot use an image that already has a stacked disk")
@@ -101,7 +101,24 @@ extension VMDirectory {
 
     try FileManager.default.copyItem(at: configURL, to: destination.configURL)
     try FileManager.default.copyItem(at: nvramURL, to: destination.nvramURL)
-    try manifest.toJSON().write(to: destination.manifestURL)
+    try contentStore.withPruneLock {
+      try manifest.toJSON().write(to: destination.manifestURL)
+    }
+
+    // Publish the temporary VM's manifest before installing the shared base.
+    // Reference-aware pruning includes in-progress manifests, so the content
+    // cannot be collected in the window before this VM is moved into place.
+    if try contentStore.contentURLIfPresent(for: contentDigest) == nil {
+      let temporaryURL = try contentStore.temporaryContentURL(for: contentDigest)
+      do {
+        try FileManager.default.copyItem(at: diskURL, to: temporaryURL)
+        _ = try contentStore.install(temporaryURL, contentDigest: contentDigest)
+      } catch {
+        try? FileManager.default.removeItem(at: temporaryURL)
+        throw error
+      }
+    }
+
     try destination.diskImageStack(contentStore: contentStore).createWritableOverlay()
 
     if generateMAC {