xtool/Sources/XUtils/TemporaryDirectory.swift

import Foundation

package struct TemporaryDirectory: ~Copyable {
    private static let debugTmp = ProcessInfo.processInfo.environment["XTL_DEBUG_TMP"] != nil

    private var shouldDelete: Bool
    private let base: URL
    package let url: URL

    /// Prepares a fresh tmpdir root.
    ///
    /// Optional, but try calling this at launch to clean up old resources.
    package static func prepare() {
        _ = TemporaryDirectoryRoot.shared
    }

    /// Creates a temporary directory where `lastPathComponent` is exactly `name`.
    ///
    /// The directory is deleted on deinit or (if the object never deinits) on next launch.
    /// To save the contents, move them elsewhere with ``persist(at:)``.
    package init(name: String) throws {
        do {
            let basename = name
                .replacingOccurrences(of: ".", with: "_")
                .replacingOccurrences(of: "/", with: "_")
                .prefix(32)
            self.base = try TemporaryDirectoryRoot.shared.url
                // ensures uniqueness
                .appendingPathComponent("dir-\(basename)-\(UUID().uuidString)")
            self.url = base.appendingPathComponent(name, isDirectory: true)
            self.shouldDelete = true
        } catch {
            // non-copyable types can't be partially initialized so we need a stub value
            self.base = URL(fileURLWithPath: "")
            self.url = base
            self.shouldDelete = false
            throw error
        }
        try FileManager.default.createDirectory(at: url, withIntermediateDirectories: true)

        if Self.debugTmp {
            stderrPrint("Created TemporaryDirectory: \(url.path)")
        }
    }

    private func _delete() {
        guard !Self.debugTmp else { return }
        try? FileManager.default.removeItem(at: base)
    }

    package consuming func persist(at location: URL) throws {
        try FileManager.default.moveItem(at: url, to: location)
        _delete()
        // we do this after moving, so that if the move fails we clean up
        shouldDelete = false
    }

    deinit {
        if shouldDelete { _delete() }
    }
}

private struct TemporaryDirectoryRoot {
    static let shared = TemporaryDirectoryRoot()

    private let _url: Result<URL, Errors>
    var url: URL {
        get throws(Errors) {
            try _url.get()
        }
    }

    private init() {
        let url = XToolCache.directory

        Self.pruneOrphans(in: url)

        self._url = Result {
            let childDir = try Self.claimDirectory(in: url)
            try FileManager.default.createDirectory(at: childDir, withIntermediateDirectories: true)
            return childDir
        }
        .mapError { $0 as? Errors ?? .tmpdirCreationFailed(url, $0) }
    }

    private static func claimDirectory(in url: URL) throws -> URL {
        // create a lockfile + tmpdir. while we hold the lock, other instaces of
        // xtool will not delete the directory during their prune step.

        try FileManager.default.createDirectory(at: url, withIntermediateDirectories: true)

        let pid = ProcessInfo.processInfo.processIdentifier

        for i in 0..<10 {
            let random = "pid-\(pid)-\(i == 0 ? "0" : UUID().uuidString)"
            let lockFile = url.appending(path: "\(random).lock")
            #if os(macOS)
            let fd = try FileDescriptor.open(
                FilePath(lockFile.path),
                .writeOnly,
                options: [.create, .exclusiveLock],
                permissions: [.ownerReadWrite, .groupRead, .otherRead],
            )
            #else
            let fd = try FileDescriptor.open(
                FilePath(lockFile.path),
                .writeOnly,
                options: [.create],
                permissions: [.ownerReadWrite, .groupRead, .otherRead],
            )
            guard try fd.tryLock(mode: .exclusive) else {
                // someone else raced us and claimed the right to prune between when we created the file and
                // when we tried to lock it
                continue
            }
            #endif
            // leak the file descriptor so that the lock is held until the process exits
            return url.appending(path: random)
        }
        throw Errors.tmpdirClaimFailed(url)
    }

    private static func pruneOrphans(in url: URL) {
        guard let children = try? FileManager.default.contentsOfDirectory(at: url, includingPropertiesForKeys: nil)
            else { return }
        for lock in children {
            do {
                let basename = lock.lastPathComponent
                if basename.hasPrefix("tmp-") {
                    // legacy tmpdir, remove it
                    try? FileManager.default.removeItem(at: lock)
                    continue
                }
                guard basename.hasPrefix("pid-") && basename.hasSuffix(".lock") else { continue }
                let lockFD = try FileDescriptor.open(FilePath(lock.path), .readWrite)
                defer { try? lockFD.close() }
                if try lockFD.tryLock(mode: .exclusive) {
                    // we must remove the directory first. if we instead removed the lock first,
                    // we could be killed after the lock was removed but before the dir was
                    // removed and therefore leave it hanging around.
                    do {
                        try FileManager.default.removeItem(at: lock.deletingPathExtension())
                    } catch CocoaError.fileNoSuchFile {
                        // pass
                    }
                    try FileManager.default.removeItem(at: lock)
                }
            } catch {
                // continue
            }
        }
    }

    enum Errors: Error, CustomStringConvertible {
        case tmpdirCreationFailed(URL, Error)
        case tmpdirClaimFailed(URL)

        var description: String {
            switch self {
            case let .tmpdirCreationFailed(url, error):
                "Could not create temporary directory in '\(url.path)': \(error)"
            case let .tmpdirClaimFailed(url):
                "Could not claim temporary directory in '\(url.path)'"
            }
        }
    }
}