Download artifacts as parallel byte ranges that resume after a dropped connection

Sources/OmarchyAppleInstaller/VerifiedArtifactStager.swift

@@ -231,7 +231,8 @@ public enum ArtifactStageError: Error, Equatable, Sendable {
 /// Multi-part payloads are downloaded part by part, then streamed together into
 /// a single pending file whose SHA-256 is computed in the same pass. Assembly
 /// therefore needs about twice the payload size on disk until the parts are
-/// deleted, and interrupted transfers resume at part granularity.
+/// deleted. Within one run a dropped connection resumes at byte granularity
+/// (see `RangedArtifactDownloader`); across runs, at part granularity.
 public struct VerifiedArtifactStager: Sendable {
   /// Disk room required while a payload is assembled from parts (payload plus
   /// the pending concatenation) before the parts are deleted.
@@ -257,7 +258,7 @@ public struct VerifiedArtifactStager: Sendable {
   private let bundledEngineDirectory: URL?
 
   public init() {
-    downloader = ProgressReportingArtifactDownloader()
+    downloader = RangedArtifactDownloader()
     promoter = AtomicArtifactFilePromoter()
     bundledEngineDirectory = Bundle.main.resourceURL?.appendingPathComponent("Engine/artifacts")
   }
@@ -1009,3 +1010,340 @@ final class ProgressReportingArtifactDownloader: NSObject, ArtifactDownloading,
     transfer.continuation.resume(with: result)
   }
 }
+
+/// Downloads an artifact as parallel HTTP byte ranges written in place, each
+/// over its own session, so one slow or dropped connection neither stalls nor
+/// restarts the whole file. A dropped range resumes from its last byte. A
+/// server that ignores Range falls back to the single-stream downloader. The
+/// stager still checks the size and SHA-256 of the result.
+final class RangedArtifactDownloader: NSObject, ArtifactDownloading, URLSessionDataDelegate,
+  @unchecked Sendable
+{
+  struct Tuning: Sendable {
+    var rangeBytes: UInt64 = 64 * 1_048_576
+    var connections = 4
+    var attemptsWithoutProgress = 4
+    var retryDelay: Duration = .seconds(1)
+  }
+
+  private enum RangeFailure: Error {
+    case ignored
+  }
+
+  private struct Attempt {
+    let written: UInt64
+    let failure: (any Error)?
+  }
+
+  private struct Transfer {
+    let continuation: CheckedContinuation<Attempt, Never>
+    let job: Job
+    let range: ClosedRange<UInt64>
+    var written: UInt64 = 0
+    var failure: (any Error)?
+  }
+
+  private let configuration: URLSessionConfiguration
+  private let tuning: Tuning
+  private let fallback: any ArtifactDownloading
+  private let lock = NSLock()
+  private var transfers = [ObjectIdentifier: Transfer]()
+
+  init(
+    configuration: URLSessionConfiguration = .ephemeral,
+    tuning: Tuning = Tuning(),
+    fallback: (any ArtifactDownloading)? = nil
+  ) {
+    self.configuration = configuration
+    self.tuning = tuning
+    self.fallback =
+      fallback ?? ProgressReportingArtifactDownloader(configuration: configuration)
+    super.init()
+  }
+
+  func download(from sourceURL: URL) async throws -> URL {
+    try await fallback.download(from: sourceURL)
+  }
+
+  func download(
+    from sourceURL: URL,
+    expectedSizeBytes: UInt64,
+    onBytes: ArtifactByteProgressHandler?
+  ) async throws -> URL {
+    guard expectedSizeBytes > 0, expectedSizeBytes != UInt64.max else {
+      return try await fallback.download(
+        from: sourceURL, expectedSizeBytes: expectedSizeBytes, onBytes: onBytes)
+    }
+    try Task.checkCancellation()
+
+    let output = FileManager.default.temporaryDirectory.appendingPathComponent(
+      "omarchy-artifact-download-\(UUID().uuidString.lowercased())",
+      isDirectory: false
+    )
+    let descriptor = Darwin.open(output.path, O_WRONLY | O_CREAT | O_EXCL, 0o600)
+    guard descriptor >= 0 else {
+      throw POSIXError(POSIXErrorCode(rawValue: Darwin.errno) ?? .EIO)
+    }
+    var keepOutput = false
+    defer {
+      Darwin.close(descriptor)
+      if !keepOutput {
+        try? FileManager.default.removeItem(at: output)
+      }
+    }
+    guard Darwin.ftruncate(descriptor, off_t(expectedSizeBytes)) == 0 else {
+      throw POSIXError(POSIXErrorCode(rawValue: Darwin.errno) ?? .EIO)
+    }
+
+    let job = Job(
+      descriptor: descriptor,
+      sizeBytes: expectedSizeBytes,
+      rangeBytes: max(1, tuning.rangeBytes),
+      onBytes: onBytes
+    )
+    do {
+      try await withThrowingTaskGroup(of: Void.self) { group in
+        for _ in 0..<min(max(1, tuning.connections), job.rangeCount) {
+          group.addTask { try await self.work(job, url: sourceURL) }
+        }
+        try await group.waitForAll()
+      }
+    } catch RangeFailure.ignored {
+      return try await fallback.download(
+        from: sourceURL, expectedSizeBytes: expectedSizeBytes, onBytes: onBytes)
+    }
+    keepOutput = true
+    return output
+  }
+
+  private func work(_ job: Job, url: URL) async throws {
+    let session = URLSession(configuration: configuration, delegate: self, delegateQueue: nil)
+    defer { session.finishTasksAndInvalidate() }
+
+    while let range = job.claim() {
+      var start = range.lowerBound
+      var failuresWithoutProgress = 0
+      while start <= range.upperBound {
+        try Task.checkCancellation()
+        let attempt = await fetch(start...range.upperBound, of: url, for: job, in: session)
+        start += attempt.written
+        guard let failure = attempt.failure else {
+          continue
+        }
+        guard PayloadPrefetchFailure.isTransient(failure) else {
+          throw failure
+        }
+        failuresWithoutProgress = attempt.written > 0 ? 1 : failuresWithoutProgress + 1
+        guard failuresWithoutProgress < tuning.attemptsWithoutProgress else {
+          throw failure
+        }
+        try await Task.sleep(for: tuning.retryDelay)
+      }
+    }
+  }
+
+  private func fetch(
+    _ range: ClosedRange<UInt64>,
+    of url: URL,
+    for job: Job,
+    in session: URLSession
+  ) async -> Attempt {
+    var request = URLRequest(url: url)
+    request.setValue(
+      "bytes=\(range.lowerBound)-\(range.upperBound)", forHTTPHeaderField: "Range")
+    let task = session.dataTask(with: request)
+    let key = ObjectIdentifier(task)
+
+    return await withTaskCancellationHandler {
+      await withCheckedContinuation { continuation in
+        lock.withLock {
+          transfers[key] = Transfer(continuation: continuation, job: job, range: range)
+        }
+        task.resume()
+      }
+    } onCancel: {
+      task.cancel()
+    }
+  }
+
+  func urlSession(
+    _ session: URLSession,
+    dataTask: URLSessionDataTask,
+    didReceive response: URLResponse,
+    completionHandler: @escaping @Sendable (URLSession.ResponseDisposition) -> Void
+  ) {
+    let key = ObjectIdentifier(dataTask)
+    let failure: (any Error)? = lock.withLock {
+      guard var transfer = transfers[key] else {
+        return nil
+      }
+      transfer.failure = Self.rangeFailure(
+        response, range: transfer.range, sizeBytes: transfer.job.sizeBytes)
+      transfers[key] = transfer
+      return transfer.failure
+    }
+    completionHandler(failure == nil ? .allow : .cancel)
+  }
+
+  func urlSession(
+    _ session: URLSession,
+    dataTask: URLSessionDataTask,
+    didReceive data: Data
+  ) {
+    let key = ObjectIdentifier(dataTask)
+    guard let transfer = lock.withLock({ transfers[key] }), transfer.failure == nil else {
+      return
+    }
+    let length = UInt64(transfer.range.count)
+    let incoming = UInt64(data.count)
+    guard transfer.written + incoming <= length else {
+      record(
+        ArtifactStageError.sizeMismatch(expected: length, actual: transfer.written + incoming),
+        for: key)
+      dataTask.cancel()
+      return
+    }
+    let offset = transfer.range.lowerBound + transfer.written
+    if let failure = transfer.job.write(data, at: offset) {
+      record(failure, for: key)
+      dataTask.cancel()
+      return
+    }
+    lock.withLock { transfers[key]?.written += incoming }
+    transfer.job.received(incoming)
+  }
+
+  func urlSession(
+    _ session: URLSession,
+    task: URLSessionTask,
+    didCompleteWithError error: (any Error)?
+  ) {
+    let key = ObjectIdentifier(task)
+    guard let transfer = lock.withLock({ transfers.removeValue(forKey: key) }) else {
+      return
+    }
+    let failure: (any Error)?
+    if let recorded = transfer.failure {
+      failure = recorded
+    } else if let error {
+      failure = (error as? URLError)?.code == .cancelled ? CancellationError() : error
+    } else if transfer.written != UInt64(transfer.range.count) {
+      failure = URLError(.networkConnectionLost)
+    } else {
+      failure = nil
+    }
+    transfer.continuation.resume(returning: Attempt(written: transfer.written, failure: failure))
+  }
+
+  private func record(_ failure: any Error, for key: ObjectIdentifier) {
+    lock.withLock { transfers[key]?.failure = failure }
+  }
+
+  /// Accepts only a 206 for exactly the requested bytes of an object of the
+  /// pinned size. A 200 means the server ignored Range.
+  private static func rangeFailure(
+    _ response: URLResponse,
+    range: ClosedRange<UInt64>,
+    sizeBytes: UInt64
+  ) -> (any Error)? {
+    guard let http = response as? HTTPURLResponse else {
+      return ArtifactStageError.unexpectedHTTPStatus(0)
+    }
+    if http.statusCode == 200 {
+      return RangeFailure.ignored
+    }
+    guard http.statusCode == 206 else {
+      return ArtifactStageError.unexpectedHTTPStatus(http.statusCode)
+    }
+    // Content-Range: bytes <first>-<last>/<total>
+    let fields = (http.value(forHTTPHeaderField: "Content-Range") ?? "")
+      .split(whereSeparator: { " -/".contains($0) })
+    guard fields.count == 4, fields[0] == "bytes",
+      let first = UInt64(fields[1]), let last = UInt64(fields[2]),
+      let total = UInt64(fields[3])
+    else {
+      return RangeFailure.ignored
+    }
+    guard total == sizeBytes else {
+      return ArtifactStageError.sizeMismatch(expected: sizeBytes, actual: total)
+    }
+    return first == range.lowerBound && last == range.upperBound ? nil : RangeFailure.ignored
+  }
+
+  /// One download: hands out ranges, writes bytes in place, and meters progress.
+  private final class Job: @unchecked Sendable {
+    private static let reportIntervalSeconds: TimeInterval = 0.25
+
+    let sizeBytes: UInt64
+    let rangeCount: Int
+    private let descriptor: Int32
+    private let rangeBytes: UInt64
+    private let onBytes: ArtifactByteProgressHandler?
+    private let lock = NSLock()
+    private var nextRange = 0
+    private var receivedBytes: UInt64 = 0
+    private var lastReportedAt: Date?
+
+    init(
+      descriptor: Int32,
+      sizeBytes: UInt64,
+      rangeBytes: UInt64,
+      onBytes: ArtifactByteProgressHandler?
+    ) {
+      self.descriptor = descriptor
+      self.sizeBytes = sizeBytes
+      self.rangeBytes = rangeBytes
+      self.onBytes = onBytes
+      rangeCount = Int((sizeBytes + rangeBytes - 1) / rangeBytes)
+    }
+
+    func claim() -> ClosedRange<UInt64>? {
+      lock.withLock {
+        guard nextRange < rangeCount else {
+          return nil
+        }
+        let first = UInt64(nextRange) * rangeBytes
+        nextRange += 1
+        return first...(min(first + rangeBytes, sizeBytes) - 1)
+      }
+    }
+
+    func write(_ data: Data, at offset: UInt64) -> (any Error)? {
+      data.withUnsafeBytes { buffer -> (any Error)? in
+        var done = 0
+        while done < buffer.count {
+          let count = Darwin.pwrite(
+            descriptor, buffer.baseAddress! + done, buffer.count - done,
+            off_t(offset) + off_t(done))
+          guard count > 0 else {
+            return POSIXError(POSIXErrorCode(rawValue: Darwin.errno) ?? .EIO)
+          }
+          done += count
+        }
+        return nil
+      }
+    }
+
+    func received(_ count: UInt64) {
+      guard let onBytes else {
+        return
+      }
+      let report: UInt64? = lock.withLock {
+        receivedBytes += count
+        let now = Date()
+        let due =
+          receivedBytes == sizeBytes
+          || lastReportedAt.map { now.timeIntervalSince($0) >= Self.reportIntervalSeconds }
+            ?? true
+        guard due else {
+          return nil
+        }
+        lastReportedAt = now
+        return receivedBytes
+      }
+      if let report {
+        onBytes(report)
+      }
+    }
+  }
+}