myriad/src/Myriad/Docker.hs

259 lines
10 KiB
Haskell

module Myriad.Docker
( EvalResult(..)
, buildImage
, buildAllImages
, startCleanup
, setupContainer
, killContainer
, killContainers
, evalCode
) where
import Control.Monad.Reader
import qualified Data.ByteString.Lazy as BL
import qualified Data.Map.Strict as M
import Data.Snowflake
import Data.String.Conversions
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 System.FilePath ((</>))
import System.Process.Typed
import Optics
import Myriad.Config
import Myriad.Core
type Myriad = MyriadT IO
data EvalResult
= EvalOk BL.ByteString
| EvalTimedOut
| EvalErrored
deriving (Show)
buildImage :: Language -> Myriad ()
buildImage lang = do
env <- ask
logInfo ["Checking for image ", cs $ imageName lang]
res <- try $ exec ["docker images -q ", imageName lang]
case res of
Left (SomeException err) -> logError ["An exception occured when checking for image ", cs $ imageName lang, ":\n", cs $ show err]
Right s -> do
when (BL.null s) . void $ do -- If string is empty that means the image does not yet exist
logInfo ["Building image ", cs $ imageName lang]
exec_ ["docker build -t ", imageName lang, " ", cs (env ^. #languagesDir) </> cs (lang ^. #name)]
logInfo ["Built image ", cs $ imageName lang]
setupQSems
when (env ^. #config % #prepareContainers) . void $ setupContainer lang
where
setupQSems :: Myriad ()
setupQSems = do
env <- ask
csem <- newQSem 1 -- We only want one container to be set up at a time
esem <- newQSem $ fromIntegral (lang ^. #concurrent)
mapMVar (env ^. #containerSems) $ M.insert (lang ^. #name) csem
mapMVar (env ^. #evalSems) $ M.insert (lang ^. #name) esem
buildAllImages :: Myriad ()
buildAllImages = do
config <- gview #config
if config ^. #buildConcurrently
then do
logInfo ["Building all images concurrently"]
forConcurrently_ (config ^. #languages) buildImage
else do
logInfo ["Building all images sequentially"]
forM_ (config ^. #languages) buildImage
startCleanup :: Myriad ()
startCleanup = do
config <- gview #config
when (config ^. #cleanupInterval > 0) . void $ do
let t = fromIntegral (config ^. #cleanupInterval)
fork $ timer t
where
-- Given time in minutes
timer :: Int -> Myriad ()
timer t = forever $ do
-- Takes time in microseconds
threadDelay $ t * 60000000
logInfo ["Starting cleanup of containers"]
n <- killContainers
logInfo ["Cleaned up ", cs $ show n, " containers, next in ", cs $ show t, " minutes"]
timer t
setupContainer :: Language -> Myriad ContainerName
setupContainer lang = do
cnts <- gview #containers >>= readMVar
case cnts M.!? (lang ^. #name) of
Nothing -> setup
Just cnt -> pure cnt
where
setup :: Myriad ContainerName
setup = do
ref <- gview #containers
cnt <- newContainerName lang
logInfo ["Setting up new container ", cs cnt]
exec_
[ "docker run --runtime="
, cs $ lang ^. #runtime
, " --rm --name="
, cs cnt
-- User 1000 will be for setting up the environment
, " -u1000:1000 -w/tmp/ -dt --net=none --cpus="
, show $ lang ^. #cpus
, " -m="
, cs $ lang ^. #memory
, " --memory-swap="
, cs $ lang ^. #memory
, " "
, imageName lang
, " /bin/sh"
]
-- The `eval` directory is where all the eval work is done
-- 711 so that users can't traverse into other people's code
exec_ ["docker exec ", cnt, " mkdir eval"]
exec_ ["docker exec ", cnt, " chmod 711 eval"]
mapMVar ref $ M.insert (lang ^. #name) cnt
logInfo ["Started new container ", cs cnt]
pure cnt
killContainer :: LanguageName -> Myriad Bool
killContainer lang = do
containers <- gview #containers >>= readMVar
case containers M.!? lang of
Nothing -> pure False
Just cnt -> do
logInfo ["Killing container ", cs cnt]
res <- kill cnt
case res of
Nothing -> do
logInfo ["Killed container ", cs cnt]
pure True
Just err -> do
logError ["An exception occured when killing ", cs cnt, ":\n", cs $ show err]
pure False
where
kill :: ContainerName -> Myriad (Maybe SomeException)
kill cnt = do
ref <- gview #containers
mapMVar ref $ M.delete lang
res <- try $ exec_ ["docker kill ", cnt]
case res of
Left err -> pure $ Just err
Right _ -> pure Nothing
killContainers :: Myriad [ContainerName]
killContainers = do
containers <- gview #containers >>= readMVar
logInfo ["Starting killing of containers"]
xs <- forConcurrently (M.toList containers) $ \(k, v) -> (v,) <$> killContainer k
logInfo ["Finished killing of containers"]
pure . map fst $ filter snd xs
evalCode :: Language -> Int -> String -> Myriad EvalResult
evalCode lang numRetries code = withContainer $ \cnt -> do
doneRef <- newMVar False -- For keeping track of if the evaluation is done, i.e. succeeded or timed out.
void . fork $ timer cnt doneRef -- `race` could not have been used here since some evals can't be cancelled.
snowflakeGen <- gview #snowflakeGen
snowflake <- liftIO $ nextSnowflake snowflakeGen
res <- try $ eval cnt snowflake
case res of
Left (SomeException err) -> do
void . killContainer $ lang ^. #name
done <- readMVar doneRef
if done
-- If we find the eval is done from an exception, then it was timed out.
then do
logInfo ["Code timed out in container ", cs cnt, ", evaluation ", cs $ show snowflake]
pure EvalTimedOut
-- Otherwise, the container was killed from another eval, so we should retry.
else do
writeMVar doneRef True
if numRetries < fromIntegral (lang ^. #retries)
then do
logInfo
[ "An exception occured in "
, cs cnt
, ", evaluation "
, cs $ show snowflake
, "retrying:\n"
, cs $ show err
]
evalCode lang (numRetries + 1) code
else do
logInfo
[ "An exception occured in "
, cs cnt
, ", evaluation "
, cs $ show snowflake
, ":\n"
, cs $ show err
]
pure EvalErrored
Right x -> do
writeMVar doneRef True
pure x
where
withContainer :: (ContainerName -> Myriad a) -> Myriad a
withContainer f = do
env <- ask
csem <- (M.! (lang ^. #name)) <$> readMVar (env ^. #containerSems)
esem <- (M.! (lang ^. #name)) <$> readMVar (env ^. #evalSems)
bracket_ (waitQSem esem) (signalQSem esem) $ do
cnt <- bracket_ (waitQSem csem) (signalQSem csem) $ setupContainer lang
f cnt
timer :: ContainerName -> MVar Bool -> Myriad ()
timer cnt doneRef = do
-- Given time in seconds
let t = fromIntegral $ lang ^. #timeout
logDebug ["Starting timeout of ", cs . show $ t, " seconds for container ", cs cnt]
-- Takes time in microseconds
threadDelay $ t * 1000000
done <- readMVar doneRef
if done
then do
logDebug ["Finished timeout for container ", cs cnt, ", but container already done"]
else do
logDebug ["Finished timeout for container ", cs cnt, " and killing it"]
writeMVar doneRef True
void . killContainer $ lang ^. #name
eval :: ContainerName -> Snowflake -> Myriad EvalResult
eval cnt snowflake = do
logInfo ["Running code in container ", cs cnt, ", evaluation ", cs $ show snowflake, ":\n", cs code]
exec_ ["docker exec ", cs cnt, " mkdir eval/", show snowflake]
exec_ ["docker exec ", cs cnt, " chmod 777 eval/", show snowflake]
-- User 1001 will be used for the actual execution so that they can't access `eval` itself
let limit = lang ^. #outputLimit
cmd = mconcat
[ "docker exec -i -u1001:1001 -w/tmp/eval/"
, show snowflake
, " "
, cnt
, " /bin/sh /var/run/run.sh 2>&1 | head -c "
, cs limit
]
pr = setStdin (byteStringInput $ cs code) $ shell cmd
logDebug ["Executing with stdin `", cs cmd, "`"]
output <- readProcessInterleaved_ pr
exec_ ["docker exec ", cnt, " rm -rf eval/", show snowflake]
logInfo ["Ran code in container ", cs cnt, ", evaluation ", cs $ show snowflake]
pure $ EvalOk output
newContainerName :: Language -> Myriad ContainerName
newContainerName lang = do
snowflakeGen <- gview #snowflakeGen
snowflake <- liftIO $ nextSnowflake snowflakeGen
pure $ "myriad-" <> cs (lang ^. #name) <> "-" <> show snowflake
imageName :: Language -> ImageName
imageName lang = "1computer1/myriad:" <> cs (lang ^. #name)