From 106eb5a2c8b5a6198178aaeb4722127006c6fa6d Mon Sep 17 00:00:00 2001
From: Nikolay Edigaryev <edigaryev@gmail.com>
Date: Sat, 10 Aug 2024 21:06:39 +0400
Subject: [PATCH] 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
---
 Sources/tart/OCI/Layerizer/DiskV2.swift | 16 ++++++++++++++--
 Sources/tart/VMStorageOCI.swift         | 12 ++++++------
 2 files changed, 20 insertions(+), 8 deletions(-)

diff --git a/Sources/tart/OCI/Layerizer/DiskV2.swift b/Sources/tart/OCI/Layerizer/DiskV2.swift
index 7c2217e..7a9e346 100644
--- a/Sources/tart/OCI/Layerizer/DiskV2.swift
+++ b/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
diff --git a/Sources/tart/VMStorageOCI.swift b/Sources/tart/VMStorageOCI.swift
index 6be1508..25283d6 100644
--- a/Sources/tart/VMStorageOCI.swift
+++ b/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...")
+
+            return .retry
           }
 
-          print("Error: \(error.localizedDescription)")
-          print("Attempting to re-try...")
-
-          return .retry
+          return .throw
         }
         try move(digestName, from: tmpVMDir)
         transaction.finish()