tart push: re-try when encountering errors when pushing disk layers (#888)

* tart push: re-try when encountering errors when pushing disk layers

* Only re-try on URLError

Author: edigaryev

Sources/tart/OCI/Layerizer/DiskV2.swift

@@ -1,6 +1,7 @@
 import Foundation
 import Compression
 import System
+import Retry
 
 class DiskV2: Disk {
   private static let bufferSizeBytes = 4 * 1024 * 1024
@@ -28,8 +29,19 @@ class DiskV2: Disk {
           let compressedData = try (data as NSData).compressed(using: .lz4) as Data
           let compressedDataDigest = Digest.hash(compressedData)
 
-          if try await !registry.blobExists(compressedDataDigest) {
-            _ = try await registry.pushBlob(fromData: compressedData, chunkSizeMb: chunkSizeMb, digest: compressedDataDigest)
+          try await retry(maxAttempts: 5, backoff: .exponentialWithFullJitter(baseDelay: .seconds(5), maxDelay: .seconds(60))) {
+            if try await !registry.blobExists(compressedDataDigest) {
+              _ = try await registry.pushBlob(fromData: compressedData, chunkSizeMb: chunkSizeMb, digest: compressedDataDigest)
+            }
+          } recoverFromFailure: { error in
+            if error is URLError {
+              print("Error: \(error.localizedDescription)")
+              print("Attempting to re-try...")
+
+              return .retry
+            }
+
+            return .throw
           }
 
           // Update progress using a relative value

Sources/tart/VMStorageOCI.swift

@@ -208,14 +208,14 @@ class VMStorageOCI: PrunableStorage {
 
           try await tmpVMDir.pullFromRegistry(registry: registry, manifest: manifest, concurrency: concurrency, localLayerCache: localLayerCache)
         } recoverFromFailure: { error in
-          if error is RuntimeError {
-            return .throw
-          }
+          if error is Retryable {
+            print("Error: \(error.localizedDescription)")
+            print("Attempting to re-try...")
 
-          print("Error: \(error.localizedDescription)")
-          print("Attempting to re-try...")
+            return .retry
+          }
 
-          return .retry
+          return .throw
         }
         try move(digestName, from: tmpVMDir)
         transaction.finish()