From 241fb0d8b7b13fb95858c5caf70a34b6180b5499 Mon Sep 17 00:00:00 2001
From: patrick brisbin <pbrisbin@gmail.com>
Date: Mon, 16 Oct 2023 09:40:47 -0400
Subject: [PATCH] Add IMAGE_CLEANUP setting

Our integration tests run all possible restylers in a space-restricted
environment, which has started to become a problem[^1]. This option will
tell Restyler to remove each image after using it, thus saving space.

[^1]: https://github.com/restyled-io/restylers/actions/runs/6522631491/job/17712406781
---
 restyle-path/main.hs         |  1 +
 src/Restyler/Options.hs      |  4 ++++
 src/Restyler/Restyler/Run.hs | 25 +++++++++++++++++++++++--
 test/SpecHelper.hs           |  1 +
 4 files changed, 29 insertions(+), 2 deletions(-)

diff --git a/restyle-path/main.hs b/restyle-path/main.hs
index 2b7523daf..08816f94b 100644
--- a/restyle-path/main.hs
+++ b/restyle-path/main.hs
@@ -64,6 +64,7 @@ main = do
                 , oRestrictions = eoRestrictions
                 , oStatsdHost = Nothing
                 , oStatsdPort = Nothing
+                , oImageCleanup = False
                 }
           }
 
diff --git a/src/Restyler/Options.hs b/src/Restyler/Options.hs
index 9e6f70f01..ca433a63a 100644
--- a/src/Restyler/Options.hs
+++ b/src/Restyler/Options.hs
@@ -22,6 +22,7 @@ data EnvOptions = EnvOptions
   , eoRestrictions :: Restrictions
   , eoStatsdHost :: Maybe String
   , eoStatsdPort :: Maybe Int
+  , eoImageCleanup :: Bool
   }
 
 data CLIOptions = CLIOptions
@@ -47,6 +48,7 @@ data Options = Options
   , oRestrictions :: Restrictions
   , oStatsdHost :: Maybe String
   , oStatsdPort :: Maybe Int
+  , oImageCleanup :: Bool
   }
 
 class HasOptions env where
@@ -84,6 +86,7 @@ parseOptions = do
       , oRestrictions = eoRestrictions
       , oStatsdHost = eoStatsdHost
       , oStatsdPort = eoStatsdPort
+      , oImageCleanup = eoImageCleanup
       }
 
 -- brittany-disable-next-binding
@@ -102,6 +105,7 @@ envParser =
     <*> envRestrictions
     <*> optional (Env.var Env.str "STATSD_HOST" mempty)
     <*> optional (Env.var Env.auto "STATSD_PORT" mempty)
+    <*> Env.switch "IMAGE_CLEANUP" mempty
 
 -- brittany-disable-next-binding
 
diff --git a/src/Restyler/Restyler/Run.hs b/src/Restyler/Restyler/Run.hs
index e30492536..f4abbb89e 100644
--- a/src/Restyler/Restyler/Run.hs
+++ b/src/Restyler/Restyler/Run.hs
@@ -36,6 +36,7 @@ import Restyler.Restyler
 import Restyler.RestylerResult
 import qualified Restyler.Wiki as Wiki
 import System.FilePath ((</>))
+import UnliftIO.Exception (tryAny)
 
 data RestylerExitFailure = RestylerExitFailure Restyler Int
   deriving stock (Show, Eq)
@@ -275,7 +276,7 @@ getDockerRunStyles Restyler {..} paths = case rRunStyle of
   RestylerRunStylePathOverwriteSep -> map (DockerRunPathOverwrite True) paths
 
 dockerRunRestyler
-  :: ( MonadIO m
+  :: ( MonadUnliftIO m
      , MonadLogger m
      , MonadSystem m
      , MonadProcess m
@@ -287,6 +288,7 @@ dockerRunRestyler
   -> m ()
 dockerRunRestyler r@Restyler {..} WithProgress {..} = do
   cwd <- getHostDirectory
+  imageCleanup <- oImageCleanup <$> view optionsL
   restrictions <- oRestrictions <$> view optionsL
 
   let
@@ -299,6 +301,11 @@ dockerRunRestyler r@Restyler {..} WithProgress {..} = do
     progress :: Text
     progress = pack (show pIndex) <> " of " <> pack (show pTotal)
 
+    -- Our integration tests run every restyler we support in a space-restricted
+    -- environment. This switch triggers removal of each image after running it,
+    -- to avoid out-of-space errors.
+    withImageCleanup f = if imageCleanup then f `finally` cleanupImage else f
+
   logInfo
     $ "Restyling"
     :# [ "restyler" .= rName
@@ -306,7 +313,7 @@ dockerRunRestyler r@Restyler {..} WithProgress {..} = do
        , "style" .= rRunStyle
        ]
 
-  ec <- case pItem of
+  ec <- withImageCleanup $ case pItem of
     DockerRunPathToStdout path -> do
       (ec, out) <- readProcessExitCode "docker" (args <> [prefix path])
       ec <$ writeFile path (fixNewline $ pack out)
@@ -325,6 +332,20 @@ dockerRunRestyler r@Restyler {..} WithProgress {..} = do
     | "./" `isPrefixOf` p = p
     | otherwise = "./" <> p
 
+  cleanupImage = do
+    eec <- tryAny $ callProcessExitCode "docker" ["image", "rm", "--force", rImage]
+    case eec of
+      Left ex ->
+        logWarn
+          $ "Exception removing Restyler image"
+          :# ["exception" .= displayException ex]
+      Right ExitSuccess ->
+        logInfo "Removed Restyler image"
+      Right (ExitFailure i) ->
+        logWarn
+          $ "Error removing Restyler image"
+          :# ["status" .= i]
+
 fixNewline :: Text -> Text
 fixNewline = (<> "\n") . T.dropWhileEnd (== '\n')
 
diff --git a/test/SpecHelper.hs b/test/SpecHelper.hs
index b64880779..c1e52011d 100644
--- a/test/SpecHelper.hs
+++ b/test/SpecHelper.hs
@@ -138,6 +138,7 @@ testOptions =
     , oRestrictions = fullRestrictions
     , oStatsdHost = Nothing
     , oStatsdPort = Nothing
+    , oImageCleanup = False
     }
 
 testAppExample :: TestAppT a -> TestAppT a