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
-    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
-    modifyMVar_ evalSems $ pure . M.insert name esem
     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
+            mapMVar containerSems $ M.insert name csem
+            mapMVar evalSems $ M.insert name esem
 
 buildAllImages :: MyriadIO ()
 buildAllImages = do
@@ -60,126 +58,128 @@ 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="
+                [ "docker run --rm --name="
-                    , cvs cnt
+                , cvs cnt
-                    -- User 1000 will be for setting up the environment
+                -- User 1000 will be for setting up the environment
-                    , " -u1000:1000 -w/tmp/ -dt --net=none --cpus="
+                , " -u1000:1000 -w/tmp/ -dt --net=none --cpus="
-                    , cvs cpus
+                , cvs cpus
-                    , " -m="
+                , " -m="
-                    , cvs memory
+                , cvs memory
-                    , " --memory-swap="
+                , " --memory-swap="
-                    , cvs memory
+                , cvs memory
-                    , " "
+                , " "
-                    , imageName lang
+                , imageName lang
-                    , " /bin/sh"
+                , " /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
+    doneRef <- newMVar False -- For keeping track of if the evaluation is done, i.e. succeeded or timed out.
-    csem <- (M.! name) <$> readMVar containerSems
+    void . fork $ timer doneRef -- `race` could not have been used here since some evals can't be cancelled.
-    esem <- (M.! name) <$> readMVar evalSems
+    snowflakeGen <- asks snowflakeGen
-    bracket_ (waitQSem esem) (signalQSem esem) $ do
+    snowflake <- liftIO $ nextSnowflake snowflakeGen
-        cnt <- bracket_ (waitQSem csem) (signalQSem csem) $ setupContainer lang
+    res <- try $ eval cnt snowflake
-        doneRef <- newMVar False -- For keeping track of if the evaluation is done, i.e. succeeded or timed out.
+    case res of
-        void . fork $ timer doneRef -- `race` could not have been used here since some evals can't be cancelled.
+        Left (SomeException err) -> do
-        res <- try $ eval cnt
+            void $ killContainer name
-        case res of
+            done <- readMVar doneRef
-            Left (SomeException err) -> do
+            if done
-                void $ killContainerMaybe name
+                -- If we find the eval is done from an exception, then it was timed out.
-                done <- readMVar doneRef
+                then do
-                if done
+                    logError ["Code timed out in container ", cvs cnt, ", evaluation ", cvs $ show snowflake]
-                    -- If we find the eval is done from an exception, then it was timed out.
+                    pure EvalTimedOut
-                    then do
+                -- Otherwise, the container was killed from another eval, so we should retry.
-                        logInfoN $ mconcat ["Code timed out in container ", cvs cnt]
+                else do
-                        pure EvalTimedOut
+                    writeMVar doneRef True
-                    -- Otherwise, the container was killed from another eval, so we should retry.
+                    if numRetries < fromIntegral retries
-                    else do
+                        then do
-                        modifyMVar_ doneRef $ pure . const True
+                            logError ["An exception occured in ", cvs cnt, ", evaluation ", cvs $ show snowflake, ", retrying:\n", cvs $ show err]
-                        if numRetries < fromIntegral retries
+                            evalCode lang (numRetries + 1) code
-                            then evalCode lang (numRetries + 1) code
+                        else do
-                            else do
+                            logError ["An exception occured in ", cvs cnt, ", evaluation ", cvs $ show snowflake, ":\n", cvs $ show err]
-                                logErrorN $ mconcat ["An exception occured when evaluating in ", cvs cnt, ":\n", cvs $ show err]
+                            pure EvalErrored
-                                pure EvalErrored
+        Right x -> do
-            Right x -> do
+            writeMVar doneRef True
-                modifyMVar_ doneRef $ pure . const True
+            pure x
-                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