From 3cc5c1437960a814b32cb7cc6c51a36a2bb0e79a Mon Sep 17 00:00:00 2001
From: Ben Gamari <ben@smart-cactus.org>
Date: Sat, 16 May 2020 21:24:31 -0400
Subject: [PATCH] Sort: Cleanup temporary files on failure

---
 src/GHC/RTS/Events/Sort.hs | 35 +++++++++++++++++++++++++++++------
 1 file changed, 29 insertions(+), 6 deletions(-)

diff --git a/src/GHC/RTS/Events/Sort.hs b/src/GHC/RTS/Events/Sort.hs
index 7168623..6832938 100644
--- a/src/GHC/RTS/Events/Sort.hs
+++ b/src/GHC/RTS/Events/Sort.hs
@@ -6,6 +6,8 @@ module GHC.RTS.Events.Sort
   ( GHC.RTS.Events.Sort.sortEvents
   ) where
 
+import Control.Exception
+import Data.IORef
 import Data.Traversable
 import Data.Coerce
 import Data.Function (on)
@@ -46,6 +48,25 @@ instance Ord OnTime where
 instance Eq OnTime where
   (==) = coerce ((==) `on` evTime)
 
+-- | C
+cleanupFiles :: ((FilePath -> IO ()) -> IO a)
+                -- ^ a continuation accepting an action to register a file to
+                -- be cleaned up.
+             -> IO a
+cleanupFiles cont =
+  bracket start finish (cont . register)
+  where
+    start :: IO (IORef [FilePath])
+    start = newIORef []
+    finish :: IORef [FilePath] -> IO ()
+    finish ref = do
+      files <- readIORef ref
+      mapM_ removeFile files
+    register :: IORef [FilePath] -> FilePath -> IO ()
+    register ref path = do
+      atomicModifyIORef ref (\paths -> (path : paths, ()))
+      return ()
+
 -- | @sortEvents tmpDir outPath eventlog@ sorts @eventlog@ via on-disk merge
 -- sort, using @tmpDir@ for intermediate data. The sorted eventlog is written
 -- to @eventlog@.
@@ -54,10 +75,11 @@ sortEvents :: FilePath  -- ^ temporary directory
            -> EventLog  -- ^ eventlog to sort
            -> IO ()
 sortEvents _tmpDir _outPath (EventLog _ (Data [])) = fail "sortEvents: no events"
-sortEvents tmpDir outPath (EventLog hdr (Data events0)) = do
+sortEvents tmpDir outPath (EventLog hdr (Data events0)) = cleanupFiles $ \cleanupFile -> do
     chunks <- toSortedChunks events0
+    mapM_ cleanupFile chunks
     hdl <- openBinaryFile outPath WriteMode
-    mergeChunks' hdl chunks
+    mergeChunks' cleanupFile hdl chunks
     hClose hdl
     return ()
   where
@@ -67,8 +89,8 @@ sortEvents tmpDir outPath (EventLog hdr (Data events0)) = do
       . mapM (writeTempChunk . sortEventsInMem)
       . chunksOf cHUNK_SIZE
 
-    mergeChunks' :: Handle -> S.Seq SortedChunk -> IO ()
-    mergeChunks' destFile chunks
+    mergeChunks' :: (FilePath -> IO ()) -> Handle -> S.Seq SortedChunk -> IO ()
+    mergeChunks' cleanupFile destFile chunks
       | S.null chunks =
         fail "sortEvents: this can't happen"
       | S.length chunks <= fAN_IN = do
@@ -78,10 +100,11 @@ sortEvents tmpDir outPath (EventLog hdr (Data events0)) = do
       | otherwise = do
         chunksss <- flip mapM (nChunks fAN_IN chunks) $ \fps -> do
           (fp, hdl) <- createTempChunk
-          mergeChunks' hdl fps
+          cleanupFile fp
+          mergeChunks' cleanupFile hdl fps
           mapM_ removeFile fps
           return fp
-        mergeChunks' destFile (S.fromList chunksss)
+        mergeChunks' cleanupFile destFile (S.fromList chunksss)
 
     readChunk :: SortedChunk -> IO [Event]
     readChunk fp = do