From ca96739050b179e9bff04ddac3e540f2af839434 Mon Sep 17 00:00:00 2001
From: Merijn Verstraaten <merijn@inconsistent.nl>
Date: Mon, 15 Oct 2018 14:44:36 +0200
Subject: [PATCH] Clean up after SLURM.

---
 benchmark-analysis/benchmark-analysis.cabal  |  1 +
 benchmark-analysis/ingest-src/ProcessPool.hs | 37 +++++++++++++-------
 2 files changed, 26 insertions(+), 12 deletions(-)

diff --git a/benchmark-analysis/benchmark-analysis.cabal b/benchmark-analysis/benchmark-analysis.cabal
index ce2b778..c261dee 100644
--- a/benchmark-analysis/benchmark-analysis.cabal
+++ b/benchmark-analysis/benchmark-analysis.cabal
@@ -94,6 +94,7 @@ Executable Ingest
                ,        exceptions == 0.10.*
                ,        filepath == 1.4.*
                ,        haskeline == 0.7.4.2
+               ,        hostname == 1.0
                ,        microlens == 0.4.8.*
                ,        monad-logger == 0.3.28.*
                ,        memory == 0.14.*
diff --git a/benchmark-analysis/ingest-src/ProcessPool.hs b/benchmark-analysis/ingest-src/ProcessPool.hs
index b14c64a..3b1a844 100644
--- a/benchmark-analysis/ingest-src/ProcessPool.hs
+++ b/benchmark-analysis/ingest-src/ProcessPool.hs
@@ -10,8 +10,8 @@ module ProcessPool
     , withProcess
     ) where
 
-import Control.Exception (Exception)
-import Control.Monad (guard)
+import Control.Exception (Exception, SomeException, try)
+import Control.Monad (guard, void)
 import Control.Monad.Catch (MonadMask, bracket, throwM, uninterruptibleMask_)
 import Data.Acquire (mkAcquireType, withAcquire, ReleaseType(ReleaseException))
 import Data.List (intercalate)
@@ -20,7 +20,10 @@ import qualified Data.Pool as Pool
 import qualified Data.Text as T
 import qualified Data.Time.LocalTime as Time
 import Data.Time.Calendar (DayOfWeek(Saturday,Sunday), dayOfWeek)
+import Network.HostName (getHostName)
+import System.Directory (removeFile)
 import System.Exit (ExitCode(ExitFailure))
+import System.FilePath ((<.>))
 import System.IO (BufferMode(LineBuffering), Handle)
 import qualified System.IO as System
 import System.Posix.Signals (sigKILL, signalProcess)
@@ -59,15 +62,18 @@ getJobTimeOut = liftIO $ do
 withProcessPool
     :: forall a m . (MonadLogger m, MonadMask m, MonadUnliftIO m)
     => Int -> GPU -> (Pool Process -> m a) -> m a
-withProcessPool n (GPU name _) = bracket createProcessPool destroyProcessPool
+withProcessPool n (GPU name _) f = do
+    hostName <- liftIO getHostName
+    bracket (createProcessPool hostName) destroyProcessPool f
   where
-    createProcessPool :: MonadIO m => m (Pool Process)
-    createProcessPool = withUnliftIO $ \(UnliftIO runInIO) -> Pool.createPool
-        (runInIO allocateProcess)
-        (runInIO . destroyProcess)
-        1
-        3153600000
-        n
+    createProcessPool :: MonadIO m => String -> m (Pool Process)
+    createProcessPool hostName = withUnliftIO $ \(UnliftIO runInIO) ->
+        Pool.createPool
+            (runInIO allocateProcess)
+            (runInIO . destroyProcess hostName)
+            1
+            3153600000
+            n
 
     destroyProcessPool :: MonadIO m => Pool Process -> m ()
     destroyProcessPool = liftIO . Pool.destroyAllResources
@@ -95,15 +101,22 @@ withProcessPool n (GPU name _) = bracket createProcessPool destroyProcessPool
             , exePath, "-L", libPath, "-W", "-S"
             ]
 
-    destroyProcess :: (MonadLogger m, MonadUnliftIO m) => Process -> m ()
-    destroyProcess Process{..} = do
+    destroyProcess
+        :: (MonadLogger m, MonadUnliftIO m) => String -> Process -> m ()
+    destroyProcess hostName Process{..} = do
         liftIO . uninterruptibleMask_ $ do
             Proc.getPid procHandle >>= mapM_ (signalProcess sigKILL)
             System.hClose inHandle
             System.hClose outHandle
             () <$ Proc.waitForProcess procHandle
+            tryRemoveFile $ "main.0" <.> show procId <.> hostName
+            tryRemoveFile $ ".PRUN_ENVIRONMENT" <.> show procId <.> hostName
         logInfoN $ "Destroyed process: " <> showText procId
 
+    tryRemoveFile :: FilePath -> IO ()
+    tryRemoveFile path =
+        void (try $ removeFile path :: IO (Either SomeException ()))
+
 checkProcess :: MonadIO m => Process -> m ()
 checkProcess Process{..} = liftIO $ do
     result <- Proc.getProcessExitCode procHandle