From f241e216146f75874a036c135cc990ff54c23056 Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Mon, 12 Sep 2022 20:56:52 +0400
Subject: [PATCH] Self-hosted temporary directory (#238)

---
 Sources/tart/Commands/Clone.swift  |  5 +++++
 Sources/tart/Commands/Create.swift |  5 +++++
 Sources/tart/Config.swift          | 21 +++++++++++++++++++--
 Sources/tart/Root.swift            |  3 +++
 Sources/tart/VMDirectory.swift     |  2 +-
 Sources/tart/VMStorageOCI.swift    |  4 ++++
 6 files changed, 37 insertions(+), 3 deletions(-)

diff --git a/Sources/tart/Commands/Clone.swift b/Sources/tart/Commands/Clone.swift
index e56daeb..46bb6ba 100644
--- a/Sources/tart/Commands/Clone.swift
+++ b/Sources/tart/Commands/Clone.swift
@@ -34,6 +34,11 @@ struct Clone: AsyncParsableCommand {
       let sourceVM = try VMStorageHelper.open(sourceName)
 
       let tmpVMDir = try VMDirectory.temporary()
+
+      // Lock the temporary VM directory to prevent it's garbage collection
+      let tmpVMDirLock = try FileLock(lockURL: tmpVMDir.baseURL)
+      try tmpVMDirLock.lock()
+
       try await withTaskCancellationHandler(operation: {
         let lock = try FileLock(lockURL: Config().tartHomeDir)
         try lock.lock()
diff --git a/Sources/tart/Commands/Create.swift b/Sources/tart/Commands/Create.swift
index 0a2bc88..2034771 100644
--- a/Sources/tart/Commands/Create.swift
+++ b/Sources/tart/Commands/Create.swift
@@ -27,6 +27,11 @@ struct Create: AsyncParsableCommand {
   func run() async throws {
     do {
       let tmpVMDir = try VMDirectory.temporary()
+
+      // Lock the temporary VM directory to prevent it's garbage collection
+      let tmpVMDirLock = try FileLock(lockURL: tmpVMDir.baseURL)
+      try tmpVMDirLock.lock()
+
       try await withTaskCancellationHandler(operation: {
         if let fromIPSW = fromIPSW {
           if fromIPSW == "latest" {
diff --git a/Sources/tart/Config.swift b/Sources/tart/Config.swift
index c6d6f8d..34799c2 100644
--- a/Sources/tart/Config.swift
+++ b/Sources/tart/Config.swift
@@ -3,6 +3,7 @@ import Foundation
 struct Config {
   let tartHomeDir: URL
   let tartCacheDir: URL
+  let tartTmpDir: URL
 
   init() throws {
     var tartHomeDir: URL
@@ -14,11 +15,27 @@ struct Config {
               .homeDirectoryForCurrentUser
               .appendingPathComponent(".tart", isDirectory: true)
     }
-
     self.tartHomeDir = tartHomeDir
-    tartCacheDir = tartHomeDir.appendingPathComponent("cache", isDirectory: true)
 
+    tartCacheDir = tartHomeDir.appendingPathComponent("cache", isDirectory: true)
     try FileManager.default.createDirectory(at: tartCacheDir, withIntermediateDirectories: true)
+
+    tartTmpDir = tartHomeDir.appendingPathComponent("tmp", isDirectory: true)
+    try FileManager.default.createDirectory(at: tartTmpDir, withIntermediateDirectories: true)
+  }
+
+  func gc() throws {
+    for entry in try FileManager.default.contentsOfDirectory(at: tartTmpDir,
+        includingPropertiesForKeys: [], options: []) {
+      let lock = try FileLock(lockURL: entry)
+      if try !lock.trylock() {
+        continue
+      }
+
+      try FileManager.default.removeItem(at: entry)
+
+      try lock.unlock()
+    }
   }
 
   static func jsonEncoder() -> JSONEncoder {
diff --git a/Sources/tart/Root.swift b/Sources/tart/Root.swift
index 20a0b69..530afbd 100644
--- a/Sources/tart/Root.swift
+++ b/Sources/tart/Root.swift
@@ -36,6 +36,9 @@ struct Root: AsyncParsableCommand {
     do {
       var command = try parseAsRoot()
 
+      // Run garbage-collection before each command (shouldn't take too long)
+      try Config().gc()
+
       if var asyncCommand = command as? AsyncParsableCommand {
         try await asyncCommand.run()
       } else {
diff --git a/Sources/tart/VMDirectory.swift b/Sources/tart/VMDirectory.swift
index 7caeb9b..1058e06 100644
--- a/Sources/tart/VMDirectory.swift
+++ b/Sources/tart/VMDirectory.swift
@@ -29,7 +29,7 @@ struct VMDirectory: Prunable {
   }
 
   static func temporary() throws -> VMDirectory {
-    let tmpDir = FileManager.default.temporaryDirectory.appendingPathComponent(UUID().uuidString)
+    let tmpDir = try Config().tartTmpDir.appendingPathComponent(UUID().uuidString)
     try FileManager.default.createDirectory(at: tmpDir, withIntermediateDirectories: false)
 
     return VMDirectory(baseURL: tmpDir)
diff --git a/Sources/tart/VMStorageOCI.swift b/Sources/tart/VMStorageOCI.swift
index 3ab4e3e..d0b5bf5 100644
--- a/Sources/tart/VMStorageOCI.swift
+++ b/Sources/tart/VMStorageOCI.swift
@@ -150,6 +150,10 @@ class VMStorageOCI: PrunableStorage {
     if !exists(digestName) {
       let tmpVMDir = try VMDirectory.temporary()
 
+      // Lock the temporary VM directory to prevent it's garbage collection
+      let tmpVMDirLock = try FileLock(lockURL: tmpVMDir.baseURL)
+      try tmpVMDirLock.lock()
+
       // 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)