From 04123091102219054063f2e2a0a5d5f1cdaf24f7 Mon Sep 17 00:00:00 2001
From: patrick brisbin <pbrisbin@gmail.com>
Date: Mon, 16 Dec 2024 08:28:05 -0500
Subject: [PATCH] Retry docker-pulls

Exceptions occur here semi-frequently, for rate limits or transient
network issues it seems. This adds exponential backoff starting with a
base of 1s and retrying up to 5 times.

Example:

    Run Seconds Timestamp
    0   0.000   00:00:02
    1   1.000   00:00:03
    2   3.000   00:00:05
    3   7.000   00:00:09
    4   15.000  00:00:17
---
 package.yaml                 |  1 +
 restyler.cabal               |  1 +
 src/Restyler/Monad/Docker.hs | 30 +++++++++++++++++++++++++++++-
 3 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/package.yaml b/package.yaml
index cd9fcc0a..41752ccd 100644
--- a/package.yaml
+++ b/package.yaml
@@ -59,6 +59,7 @@ library:
     - opt-env-conf >= 0.6.0.2
     - path
     - relude
+    - retry
     - ronn
     - ronn-opt-env-conf
     - safe-coloured-text
diff --git a/restyler.cabal b/restyler.cabal
index f3ed22b0..576ac224 100644
--- a/restyler.cabal
+++ b/restyler.cabal
@@ -97,6 +97,7 @@ library
     , opt-env-conf >=0.6.0.2
     , path
     , relude
+    , retry
     , ronn
     , ronn-opt-env-conf
     , safe-coloured-text
diff --git a/src/Restyler/Monad/Docker.hs b/src/Restyler/Monad/Docker.hs
index bd84ba27..59632e86 100644
--- a/src/Restyler/Monad/Docker.hs
+++ b/src/Restyler/Monad/Docker.hs
@@ -21,6 +21,14 @@ import Restyler.Prelude
 import Data.Text qualified as T
 import Restyler.AnnotatedException
 import System.Process.Typed
+import UnliftIO.Retry
+  ( RetryPolicyM
+  , RetryStatus (..)
+  , exponentialBackoff
+  , limitRetries
+  , recovering
+  , skipAsyncExceptions
+  )
 
 class Monad m => MonadDocker m where
   dockerPull :: HasCallStack => String -> m ExitCode
@@ -46,11 +54,31 @@ instance
   (MonadUnliftIO m, MonadLogger m, MonadReader env m, HasLogger env)
   => MonadDocker (ActualDocker m)
   where
-  dockerPull image = runDocker ["pull", "--quiet", image]
+  dockerPull image =
+    recovering dockerPullRetryPolicy skipAsyncExceptions $ \status -> do
+      when (rsIterNumber status > 0) $ do
+        logWarn
+          $ "Retrying docker-pull"
+          :# [ "attempt" .= rsIterNumber status
+             , "limit" .= dockerPullRetryLimit
+             ]
+
+      runDocker ["pull", "--quiet", image]
   dockerRun args = runDocker $ ["run", "--rm"] <> args
   dockerRunStdout args = runDockerStdout $ ["run", "--rm"] <> args
   dockerImageRm image = runDocker_ ["image", "rm", "--force", image]
 
+dockerPullRetryPolicy :: Monad m => RetryPolicyM m
+dockerPullRetryPolicy =
+  exponentialBackoff dockerPullRetryBaseBackoff
+    <> limitRetries dockerPullRetryLimit
+
+dockerPullRetryBaseBackoff :: Int
+dockerPullRetryBaseBackoff = 1 * 1000000
+
+dockerPullRetryLimit :: Int
+dockerPullRetryLimit = 5
+
 runDocker
   :: (MonadUnliftIO m, MonadLogger m, MonadReader env m, HasLogger env, HasCallStack)
   => [String]