From 79084555f66f5cf1e3bbf52504267c64a2491b0d Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Fri, 12 Apr 2024 21:32:03 +0400
Subject: [PATCH] tart pull: retry VM pull with exponential backoff (#788)

---
 Package.resolved                | 20 +++++++++++++++++++-
 Package.swift                   |  4 +++-
 Sources/tart/VMStorageOCI.swift | 10 +++++++++-
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/Package.resolved b/Package.resolved
index 8ac986e..14a84f2 100644
--- a/Package.resolved
+++ b/Package.resolved
@@ -1,5 +1,5 @@
 {
-  "originHash" : "b6263fb821fd8d1b501b388d41e163542fe1aa79613c2842b2495eb23f3424f3",
+  "originHash" : "6d48639bc0ea02002de0b4f38fe3fce0ddc9d174f2e56180c2ffcbedb7391ef8",
   "pins" : [
     {
       "identity" : "antlr4",
@@ -82,6 +82,15 @@
         "version" : "1.0.3"
       }
     },
+    {
+      "identity" : "swift-log",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/apple/swift-log.git",
+      "state" : {
+        "revision" : "e97a6fcb1ab07462881ac165fdbb37f067e205d5",
+        "version" : "1.5.4"
+      }
+    },
     {
       "identity" : "swift-numerics",
       "kind" : "remoteSourceControl",
@@ -91,6 +100,15 @@
         "version" : "1.0.2"
       }
     },
+    {
+      "identity" : "swift-retry",
+      "kind" : "remoteSourceControl",
+      "location" : "https://github.com/fumoboy007/swift-retry",
+      "state" : {
+        "revision" : "9f133487ffc2ab4539688c29efe57bb1ba31d7b0",
+        "version" : "0.2.3"
+      }
+    },
     {
       "identity" : "swift-sysctl",
       "kind" : "remoteSourceControl",
diff --git a/Package.swift b/Package.swift
index 9ce8271..5de36ff 100644
--- a/Package.swift
+++ b/Package.swift
@@ -22,7 +22,8 @@ let package = Package(
     .package(url: "https://github.com/cfilipov/TextTable", branch: "master"),
     .package(url: "https://github.com/sersoft-gmbh/swift-sysctl.git", from: "1.8.0"),
     .package(url: "https://github.com/orchetect/SwiftRadix", from: "1.3.1"),
-    .package(url: "https://github.com/groue/Semaphore", from: "0.0.8")
+    .package(url: "https://github.com/groue/Semaphore", from: "0.0.8"),
+    .package(url: "https://github.com/fumoboy007/swift-retry", from: "0.2.3"),
   ],
   targets: [
     .executableTarget(name: "tart", dependencies: [
@@ -38,6 +39,7 @@ let package = Package(
       .product(name: "Sysctl", package: "swift-sysctl"),
       .product(name: "SwiftRadix", package: "SwiftRadix"),
       .product(name: "Semaphore", package: "Semaphore"),
+      .product(name: "DMRetry", package: "swift-retry"),
     ], exclude: [
       "OCI/Reference/Makefile",
       "OCI/Reference/Reference.g4",
diff --git a/Sources/tart/VMStorageOCI.swift b/Sources/tart/VMStorageOCI.swift
index 307e7f0..b098802 100644
--- a/Sources/tart/VMStorageOCI.swift
+++ b/Sources/tart/VMStorageOCI.swift
@@ -1,5 +1,6 @@
 import Foundation
 import Sentry
+import Retry
 
 class VMStorageOCI: PrunableStorage {
   let baseURL = try! Config().tartCacheDir.appendingPathComponent("OCIs", isDirectory: true)
@@ -188,7 +189,14 @@ class VMStorageOCI: PrunableStorage {
       }
 
       try await withTaskCancellationHandler(operation: {
-        try await tmpVMDir.pullFromRegistry(registry: registry, manifest: manifest, concurrency: concurrency)
+        try await retry(maxAttempts: 5, backoff: .exponentialWithFullJitter(baseDelay: .seconds(5), maxDelay: .seconds(60))) {
+          try await tmpVMDir.pullFromRegistry(registry: registry, manifest: manifest, concurrency: concurrency)
+        } recoverFromFailure: { error in
+          print("Error: \(error.localizedDescription)")
+          print("Attempting to re-try...")
+
+          return .retry
+        }
         try move(digestName, from: tmpVMDir)
         transaction.finish()
       }, onCancel: {