Clean up code

src/Myriad/Docker.hs

@@ -5,43 +5,41 @@ module Myriad.Docker
     , startCleanup
     , setupContainer
     , killContainer
-    , killAllContainers
+    , killContainers
-    , killContainerMaybe
-    , killAllContainersMaybe
     , evalCode
     ) where
 
-import Control.Monad.Logger
 import Control.Monad.Reader
 
 import qualified Data.Map.Strict as M
+import Data.Snowflake
 
 import Control.Concurrent.Async.Lifted
 import Control.Concurrent.Lifted (fork, threadDelay)
 import Control.Concurrent.MVar.Lifted
 import Control.Concurrent.QSem.Lifted
 import Control.Exception.Lifted
-import Data.Snowflake
 import System.Process.Typed
 
 import Myriad.Core
 import Myriad.Util
 
-exec :: String -> MyriadIO ()
-exec = runProcess_ . shell
-
 buildImage :: LanguageConfig -> MyriadIO ()
 buildImage lang@LanguageConfig { name, concurrent } = do
-    logInfoN $ mconcat ["Building image ", cvs $ imageName lang]
+    logInfo ["Building image ", cvs $ imageName lang]
-    let cmd = mconcat ["docker build -t ", imageName lang, " ./languages/", cvs name]
+    exec_ ["docker build -t ", imageName lang, " ./languages/", cvs name]
-    runProcess_ . setStdout nullStream $ shell cmd
+    setupQSems
-    logInfoN $ mconcat ["Built image ", cvs $ imageName lang]
+    logInfo ["Built image ", cvs $ imageName lang]
-    Env { config = MyriadConfig { prepareContainers }, containerSems, evalSems } <- ask
+    MyriadConfig { prepareContainers } <- asks config
+    when_ prepareContainers $ setupContainer lang
+    where
+        setupQSems :: MyriadIO ()
+        setupQSems = do
+            Env { containerSems, evalSems } <- ask
             csem <- newQSem 1 -- We only want one container to be set up at a time
             esem <- newQSem $ fromIntegral concurrent
-    modifyMVar_ containerSems $ pure . M.insert name csem
+            mapMVar containerSems $ M.insert name csem
-    modifyMVar_ evalSems $ pure . M.insert name esem
+            mapMVar evalSems $ M.insert name esem
-    when_ prepareContainers $ setupContainer lang
 
 buildAllImages :: MyriadIO ()
 buildAllImages = do
@@ -60,19 +58,22 @@ startCleanup = do
         timer :: Int -> MyriadIO ()
         timer t = forever do
             threadDelay t
-            n <- killAllContainersMaybe
+            n <- killContainers
-            logInfoN $ mconcat ["Cleaned up ", cvs $ show n, " containers"]
+            logInfo ["Cleaned up ", cvs $ show n, " containers"]
             timer t
 
 setupContainer :: LanguageConfig -> MyriadIO ContainerName
 setupContainer lang@LanguageConfig { name, memory, cpus } = do
-    ref <- asks containers
+    cnts <- asks containers >>= readMVar
-    cnts <- readMVar ref
     case cnts M.!? name of
-        Just x  -> pure x
+        Nothing  -> setup
-        Nothing -> do
+        Just cnt -> pure cnt
+    where
+        setup :: MyriadIO ContainerName
+        setup = do
+            ref <- asks containers
             cnt <- newContainerName lang
-            let cmd = mconcat
+            exec_
                 [ "docker run --rm --name="
                 , cvs cnt
                 -- User 1000 will be for setting up the environment
@@ -86,100 +87,99 @@ setupContainer lang@LanguageConfig { name, memory, cpus } = do
                 , imageName lang
                 , " /bin/sh"
                 ]
-            runProcess_ . setStdout nullStream $ shell cmd
             -- The `eval` directory is where all the eval work is done
             -- 711 so that users can't traverse into other people's code
-            exec $ mconcat ["docker exec ", cnt, " mkdir eval"]
+            exec_ ["docker exec ", cnt, " mkdir eval"]
-            exec $ mconcat ["docker exec ", cnt, " chmod 711 eval"]
+            exec_ ["docker exec ", cnt, " chmod 711 eval"]
-            modifyMVar_ ref $ pure . M.insert name cnt
+            mapMVar ref $ M.insert name cnt
-            logInfoN $ mconcat ["Started container ", cvs cnt]
+            logInfo ["Started container ", cvs cnt]
             pure cnt
 
-killContainer :: Language -> MyriadIO ()
+killContainer :: Language -> MyriadIO Bool
 killContainer lang = do
-    ref <- asks containers
-    containers <- readMVar ref
-    case containers M.!? lang of
-        Nothing  -> pure ()
-        Just cnt -> do
-            modifyMVar_ ref $ pure . M.delete lang
-            let cmd = mconcat ["docker kill ", cnt]
-            runProcess_ . setStderr nullStream . setStdout nullStream $ shell cmd
-            logInfoN $ mconcat ["Killed container ", cvs cnt]
-
-killContainerMaybe :: Language -> MyriadIO Bool
-killContainerMaybe lang = do
     containers <- asks containers >>= readMVar
     case containers M.!? lang of
         Nothing  -> pure False
         Just cnt -> do
-            res <- try @_ @SomeException $ killContainer lang
+            res <- kill cnt
             case res of
-                Left err -> do
+                Nothing  -> pure True
-                    logErrorN $ mconcat ["An exception occured when killing ", cvs cnt, ":\n", cvs $ show err]
+                Just err -> do
+                    logError ["An exception occured when killing ", cvs cnt, ":\n", cvs $ show err]
                     pure False
-                Right _  -> pure True
+    where
+        kill :: ContainerName -> MyriadIO (Maybe SomeException)
+        kill cnt = do
+            ref <- asks containers
+            mapMVar ref $ M.delete lang
+            res <- try $ exec_ ["docker kill ", cnt]
+            case res of
+                Left err -> pure $ Just err
+                Right _  -> do
+                    logInfo ["Killed container ", cvs cnt]
+                    pure Nothing
 
-killAllContainers :: MyriadIO ()
+killContainers :: MyriadIO [ContainerName]
-killAllContainers = do
+killContainers = do
     containers <- asks containers >>= readMVar
-    forConcurrently_ (M.keys containers) $ killContainer
+    xs <- forConcurrently (M.toList containers) \(k, v) -> (v,) <$> killContainer k
-
-killAllContainersMaybe :: MyriadIO [ContainerName]
-killAllContainersMaybe = do
-    containers <- asks containers >>= readMVar
-    xs <- forConcurrently (M.toList containers) \(k, v) -> (v,) <$> killContainerMaybe k
     pure . map fst $ filter snd xs
 
 evalCode :: LanguageConfig -> Int -> String -> MyriadIO EvalResult
-evalCode lang@LanguageConfig { name, timeout, retries } numRetries code = do
+evalCode lang@LanguageConfig { name, timeout, retries } numRetries code = withContainer \cnt -> do
-    Env { containerSems, evalSems } <- ask
-    csem <- (M.! name) <$> readMVar containerSems
-    esem <- (M.! name) <$> readMVar evalSems
-    bracket_ (waitQSem esem) (signalQSem esem) $ do
-        cnt <- bracket_ (waitQSem csem) (signalQSem csem) $ setupContainer lang
     doneRef <- newMVar False -- For keeping track of if the evaluation is done, i.e. succeeded or timed out.
     void . fork $ timer doneRef -- `race` could not have been used here since some evals can't be cancelled.
-        res <- try $ eval cnt
+    snowflakeGen <- asks snowflakeGen
+    snowflake <- liftIO $ nextSnowflake snowflakeGen
+    res <- try $ eval cnt snowflake
     case res of
         Left (SomeException err) -> do
-                void $ killContainerMaybe name
+            void $ killContainer name
             done <- readMVar doneRef
             if done
                 -- If we find the eval is done from an exception, then it was timed out.
                 then do
-                        logInfoN $ mconcat ["Code timed out in container ", cvs cnt]
+                    logError ["Code timed out in container ", cvs cnt, ", evaluation ", cvs $ show snowflake]
                     pure EvalTimedOut
                 -- Otherwise, the container was killed from another eval, so we should retry.
                 else do
-                        modifyMVar_ doneRef $ pure . const True
+                    writeMVar doneRef True
                     if numRetries < fromIntegral retries
-                            then evalCode lang (numRetries + 1) code
+                        then do
+                            logError ["An exception occured in ", cvs cnt, ", evaluation ", cvs $ show snowflake, ", retrying:\n", cvs $ show err]
+                            evalCode lang (numRetries + 1) code
                         else do
-                                logErrorN $ mconcat ["An exception occured when evaluating in ", cvs cnt, ":\n", cvs $ show err]
+                            logError ["An exception occured in ", cvs cnt, ", evaluation ", cvs $ show snowflake, ":\n", cvs $ show err]
                             pure EvalErrored
         Right x -> do
-                modifyMVar_ doneRef $ pure . const True
+            writeMVar doneRef True
             pure x
     where
+        withContainer :: (ContainerName -> MyriadIO a) -> MyriadIO a
+        withContainer f = do
+            Env { containerSems, evalSems } <- ask
+            csem <- (M.! name) <$> readMVar containerSems
+            esem <- (M.! name) <$> readMVar evalSems
+            bracket_ (waitQSem esem) (signalQSem esem) do
+                cnt <- bracket_ (waitQSem csem) (signalQSem csem) $ setupContainer lang
+                f cnt
+
         timer :: MVar Bool -> MyriadIO ()
         timer doneRef = do
             threadDelay $ fromIntegral timeout * 1000000
             done <- readMVar doneRef
             unless_ done do
-                modifyMVar_ doneRef $ pure . const True
+                writeMVar doneRef True
-                killContainerMaybe name
+                killContainer name
 
-        eval :: ContainerName -> MyriadIO EvalResult
+        eval :: ContainerName -> Snowflake -> MyriadIO EvalResult
-        eval cnt = do
+        eval cnt snowflake = do
-            logInfoN $ mconcat ["Running code in container ", cvs cnt, ":\n", cvs code]
+            logInfo ["Running code in container ", cvs cnt, ", evaluation ", cvs $ show snowflake, ":\n", cvs code]
-            snowflakeGen <- asks snowflakeGen
+            exec_ ["docker exec ", cvs cnt, " mkdir eval/", show snowflake]
-            snowflake <- liftIO $ nextSnowflake snowflakeGen
+            exec_ ["docker exec ", cvs cnt, " chmod 777 eval/", show snowflake]
-            exec $ mconcat ["docker exec ", cvs cnt, " mkdir eval/", show snowflake]
-            exec $ mconcat ["docker exec ", cvs cnt, " chmod 777 eval/", show snowflake]
             -- User 1001 will be used for the actual execution so that they can't access `eval` itself
             let args = ["exec", "-u1001:1001", "-w/tmp/eval/" <> show snowflake, cnt, "/bin/sh", "/var/run/run.sh", code]
             output <- readProcessInterleaved_ $ proc "docker" args
-            exec $ mconcat ["docker exec ", cnt, " rm -rf eval/", show snowflake]
+            exec_ ["docker exec ", cnt, " rm -rf eval/", show snowflake]
-            logInfoN $ mconcat ["Ran code in container ", cvs cnt]
+            logInfo ["Ran code in container ", cvs cnt, ", evaluation ", cvs $ show snowflake]
             pure $ EvalOk output

src/Myriad/Server.hs

@@ -5,7 +5,6 @@ module Myriad.Server
     ) where
 
 import Control.Monad.Except
-import Control.Monad.Logger
 import Control.Monad.Reader
 
 import Data.Aeson
@@ -42,13 +41,13 @@ serverT = handleLanguages :<|> handleEval :<|> handleContainers :<|> handleClean
     where
         handleLanguages :: MyriadT m [T.Text]
         handleLanguages = do
-            logInfoN $ mconcat ["GET /languages"]
+            logInfo ["GET /languages"]
             MyriadConfig { languages } <- asks config
             pure . map name $ languages
 
         handleEval :: EvalRequest -> MyriadT m EvalResponse
         handleEval EvalRequest { language, code } = do
-            logInfoN $ mconcat ["POST /eval"]
+            logInfo ["POST /eval"]
             MyriadConfig { languages } <- asks config
             case find (\x -> name x == language) languages of
                 Nothing  -> throwError $ err404 { errBody = "Language " <> cvs language <> " was not found" }
@@ -61,11 +60,11 @@ serverT = handleLanguages :<|> handleEval :<|> handleContainers :<|> handleClean
 
         handleContainers :: MyriadT m [T.Text]
         handleContainers = do
-            logInfoN $ mconcat ["GET /containers"]
+            logInfo ["GET /containers"]
             containers <- asks containers >>= readMVar
             pure . map cvs $ M.elems containers
 
         handleCleanup :: MyriadT m [T.Text]
         handleCleanup = do
-            logInfoN $ mconcat ["POST /cleanup"]
+            logInfo ["POST /cleanup"]
-            map cvs <$> killAllContainersMaybe
+            map cvs <$> killContainers

src/Myriad/Util.hs

@@ -4,12 +4,24 @@ module Myriad.Util
     , cvs
     , when_
     , unless_
+    , exec
+    , exec_
+    , logInfo
+    , logError
+    , mapMVar
+    , writeMVar
     ) where
 
+import qualified Control.Monad.Logger as L
 import Control.Monad.Reader
 
+import qualified Data.ByteString.Lazy as BL
 import Data.Snowflake
 import Data.String.Conversions
+import qualified Data.Text as T
+
+import Control.Concurrent.MVar.Lifted
+import System.Process.Typed
 
 import Myriad.Core
 
@@ -31,3 +43,21 @@ when_ p = when p . void
 
 unless_ :: Applicative f => Bool -> f a -> f ()
 unless_ p = unless p . void 
+
+exec :: [String] -> MyriadIO BL.ByteString
+exec = readProcessInterleaved_ . shell . mconcat
+
+exec_ :: [String] -> MyriadIO ()
+exec_ = void . exec
+
+logInfo :: [T.Text] -> MyriadIO ()
+logInfo = L.logInfoN . mconcat
+
+logError :: [T.Text] -> MyriadIO ()
+logError = L.logErrorN . mconcat
+
+mapMVar :: MVar a -> (a -> a) -> MyriadIO ()
+mapMVar var f = modifyMVar_ var (pure . f)
+
+writeMVar :: MVar a -> a -> MyriadIO ()
+writeMVar var x = mapMVar var $ const x