From 8925211a467cfbe3b1768896948bfda392e5fff5 Mon Sep 17 00:00:00 2001 From: Alex Biehl Date: Thu, 13 Jun 2024 18:10:50 +0200 Subject: [PATCH] Introduce concurrent run strategy for launching containers --- src/TestContainers.hs | 1 + src/TestContainers/Config.hs | 8 +- src/TestContainers/Docker.hs | 181 +++++++++++++++++++++---------- src/TestContainers/Monad.hs | 72 ++++++++++-- src/TestContainers/Tasty.hs | 117 ++++++++++++-------- src/TestContainers/Trace.hs | 2 + test/TestContainers/TastySpec.hs | 13 ++- testcontainers.cabal | 3 + 8 files changed, 277 insertions(+), 120 deletions(-) diff --git a/src/TestContainers.hs b/src/TestContainers.hs index 519b44e..6a63223 100644 --- a/src/TestContainers.hs +++ b/src/TestContainers.hs @@ -35,6 +35,7 @@ module TestContainers M.setExpose, M.setWaitingFor, M.withFollowLogs, + M.withDependencies, -- * Logs M.LogConsumer, diff --git a/src/TestContainers/Config.hs b/src/TestContainers/Config.hs index 3a0e580..baf4246 100644 --- a/src/TestContainers/Config.hs +++ b/src/TestContainers/Config.hs @@ -1,5 +1,6 @@ module TestContainers.Config - ( Config (..), + ( RunStrategy (..), + Config (..), defaultConfig, defaultDockerConfig, determineConfig, @@ -7,7 +8,7 @@ module TestContainers.Config where import {-# SOURCE #-} TestContainers.Docker (createRyukReaper) -import TestContainers.Monad (Config (..)) +import TestContainers.Monad (Config (..), RunStrategy (..)) -- | Default configuration. -- @@ -17,7 +18,8 @@ defaultConfig = Config { configDefaultWaitTimeout = Just 60, configTracer = mempty, - configCreateReaper = createRyukReaper + configCreateReaper = createRyukReaper, + configRunStrategy = SequentialRunStrategy } -- | Default configuration. diff --git a/src/TestContainers/Docker.hs b/src/TestContainers/Docker.hs index 85a69a5..4e1965c 100644 --- a/src/TestContainers/Docker.hs +++ b/src/TestContainers/Docker.hs @@ -8,6 +8,7 @@ {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE OverloadedStrings #-} {-# LANGUAGE RankNTypes #-} +{-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE StandaloneDeriving #-} {-# LANGUAGE TypeFamilies #-} @@ -94,6 +95,7 @@ module TestContainers.Docker setLink, setExpose, setWaitingFor, + withDependencies, run, -- * Following logs @@ -158,7 +160,8 @@ module TestContainers.Docker where import Control.Concurrent (threadDelay) -import Control.Exception (IOException, throw) +import qualified Control.Concurrent.Async +import Control.Exception (IOException, evaluate, throw) import Control.Monad (forM_, replicateM, unless) import Control.Monad.Catch ( Exception, @@ -169,11 +172,12 @@ import Control.Monad.Catch try, ) import Control.Monad.IO.Class (MonadIO (liftIO)) -import Control.Monad.IO.Unlift (MonadUnliftIO (withRunInIO)) +import Control.Monad.IO.Unlift (MonadUnliftIO (withRunInIO), askRunInIO) import Control.Monad.Reader (MonadReader (..)) import Control.Monad.Trans.Resource ( ReleaseKey, ResIO, + allocate, register, runResourceT, ) @@ -209,12 +213,13 @@ import Optics.Optic ((%), (<&>)) import System.Directory (doesFileExist) import System.Environment (lookupEnv) import System.IO (Handle, hClose) -import System.IO.Unsafe (unsafePerformIO) +import System.IO.Unsafe (unsafeInterleaveIO, unsafePerformIO) import qualified System.Process as Process import qualified System.Random as Random import System.Timeout (timeout) import TestContainers.Config ( Config (..), + RunStrategy (..), defaultDockerConfig, determineConfig, ) @@ -263,6 +268,7 @@ import TestContainers.Docker.State import TestContainers.Monad ( MonadDocker, TestContainer, + defer, ) import TestContainers.Trace (Trace (..), Tracer, newTracer, withTrace) import Prelude hiding (error, id) @@ -288,7 +294,8 @@ data ContainerRequest = ContainerRequest labels :: [(Text, Text)], noReaper :: Bool, followLogs :: Maybe LogConsumer, - workDirectory :: Maybe Text + workDirectory :: Maybe Text, + dependencies :: [Container] } instance WithoutReaper ContainerRequest where @@ -324,7 +331,8 @@ containerRequest image = labels = mempty, noReaper = False, followLogs = Nothing, - workDirectory = Nothing + workDirectory = Nothing, + dependencies = [] } -- | Set the name of a Docker container. This is equivalent to invoking @docker run@ @@ -453,6 +461,13 @@ withFollowLogs :: LogConsumer -> ContainerRequest -> ContainerRequest withFollowLogs logConsumer request = request {followLogs = Just logConsumer} +-- | +-- +-- @since x.x.x +withDependencies :: [Container] -> ContainerRequest -> ContainerRequest +withDependencies dependencies request = + request {dependencies} + -- | Defintion of a 'Port'. Allows for specifying ports using various protocols. Due to the -- 'Num' and 'IsString' instance allows for convenient Haskell literals. -- @@ -556,7 +571,8 @@ run request = do labels, noReaper, followLogs, - workDirectory + workDirectory, + dependencies } = request config@Config {configTracer, configCreateReaper} <- @@ -568,8 +584,6 @@ run request = do pure [] else reaperLabels <$> configCreateReaper - image@Image {tag} <- runToImage toImage - name <- case naming of RandomName -> return Nothing @@ -578,59 +592,110 @@ run request = do Just . (prefix <>) . ("-" <>) . pack <$> replicateM 6 (Random.randomRIO ('a', 'z')) - let dockerRun :: [Text] - dockerRun = - concat $ - [["run"]] - ++ [["--detach"]] - ++ [["--name", containerName] | Just containerName <- [name]] - ++ [["--label", label <> "=" <> value] | (label, value) <- additionalLabels ++ labels] - ++ [["--env", variable <> "=" <> value] | (variable, value) <- env] - ++ [["--publish", pack (show port) <> "/" <> protocol] | Port {port, protocol} <- exposedPorts] - ++ [["--network", networkName] | Just (Right networkName) <- [network]] - ++ [["--network", networkId dockerNetwork] | Just (Left dockerNetwork) <- [network]] - ++ [["--network-alias", alias] | Just alias <- [networkAlias]] - ++ [["--link", container] | container <- links] - ++ [["--volume", src <> ":" <> dest] | (src, dest) <- volumeMounts] - ++ [["--rm"] | rmOnExit] - ++ [["--workdir", workdir] | Just workdir <- [workDirectory]] - ++ [["--memory", value] | Just value <- [memory]] - ++ [["--cpus", value] | Just value <- [cpus]] - ++ [[tag]] - ++ [command | Just command <- [cmd]] - - stdout <- docker configTracer dockerRun - - let id :: ContainerId - !id = - -- N.B. Force to not leak STDOUT String - strip (pack stdout) - - -- Careful, this is really meant to be lazy - ~inspectOutput = - unsafePerformIO $ - internalInspect configTracer id + (actuallyRunDocker, waitOnContainer) <- applyRunStrategy $ do + image@Image {tag} <- runToImage toImage - -- We don't issue 'ReleaseKeys' for cleanup anymore. Ryuk takes care of cleanup - -- for us once the session has been closed. - releaseKey <- register (pure ()) + liftIO $ + forM_ dependencies $ + \Container {wait} -> wait + + let dockerRun :: [Text] + dockerRun = + concat $ + [["run"]] + ++ [["--detach"]] + ++ [["--name", containerName] | Just containerName <- [name]] + ++ [["--label", label <> "=" <> value] | (label, value) <- additionalLabels ++ labels] + ++ [["--env", variable <> "=" <> value] | (variable, value) <- env] + ++ [["--publish", pack (show port) <> "/" <> protocol] | Port {port, protocol} <- exposedPorts] + ++ [["--network", networkName] | Just (Right networkName) <- [network]] + ++ [["--network", networkId dockerNetwork] | Just (Left dockerNetwork) <- [network]] + ++ [["--network-alias", alias] | Just alias <- [networkAlias]] + ++ [["--link", container] | container <- links] + ++ [["--volume", src <> ":" <> dest] | (src, dest) <- volumeMounts] + ++ [["--rm"] | rmOnExit] + ++ [["--workdir", workdir] | Just workdir <- [workDirectory]] + ++ [["--memory", value] | Just value <- [memory]] + ++ [["--cpus", value] | Just value <- [cpus]] + ++ [[tag]] + ++ [command | Just command <- [cmd]] + + stdout <- docker configTracer dockerRun + + let id :: ContainerId + !id = + -- N.B. Force to not leak STDOUT String + strip (pack stdout) + + -- Careful, this is really meant to be lazy + ~inspectOutput = + unsafePerformIO $ + internalInspect configTracer id + + -- We don't issue 'ReleaseKeys' for cleanup anymore. Ryuk takes care of cleanup + -- for us once the session has been closed. + releaseKey <- register (pure ()) + + forM_ followLogs $ + dockerFollowLogs configTracer id + + let container = + Container + { id, + releaseKey, + image, + inspectOutput, + config, + wait = pure () + } + + -- Last but not least, execute the WaitUntilReady checks + waitUntilReady container readiness + + pure container + + -- We want to be very gentle and not force the container if possible so that we can + -- install the wait action. + ~container <- + actuallyRunDocker + + pure $ + let ~Container {..} = container + in Container + { wait = waitOnContainer, + .. + } - forM_ followLogs $ - dockerFollowLogs configTracer id +applyRunStrategy :: + TestContainer a -> + TestContainer (TestContainer a, IO ()) +applyRunStrategy action = do + Config {configTracer, configRunStrategy} <- + ask + case configRunStrategy of + SequentialRunStrategy -> + pure (action, pure ()) + ConcurrentRunStrategy _limit -> do + runInIO <- + askRunInIO - let container = - Container - { id, - releaseKey, - image, - inspectOutput, - config - } + (_releaseKey, handle) <- + allocate + (Control.Concurrent.Async.async (runInIO action)) + Control.Concurrent.Async.cancel + + let returnAction = + liftIO $ unsafeInterleaveIO $ do + Control.Concurrent.Async.wait handle + + waitAction = do + withTrace configTracer TraceWaitOnDependency + _ <- Control.Concurrent.Async.wait handle + pure () - -- Last but not least, execute the WaitUntilReady checks - waitUntilReady container readiness + defer waitAction - pure container + pure (returnAction, waitAction) -- | Sets up a Ryuk 'Reaper'. -- @@ -1109,7 +1174,9 @@ data Container = Container -- | Configuration used to create and run this container. config :: Config, -- | Memoized output of `docker inspect`. This is being calculated lazily. - inspectOutput :: InspectOutput + inspectOutput :: InspectOutput, + -- | Wait on the container as a dependency. + wait :: IO () } -- | Returns the id of the container. diff --git a/src/TestContainers/Monad.hs b/src/TestContainers/Monad.hs index 86eaf11..7750d67 100644 --- a/src/TestContainers/Monad.hs +++ b/src/TestContainers/Monad.hs @@ -4,6 +4,7 @@ {-# LANGUAGE GeneralizedNewtypeDeriving #-} {-# LANGUAGE MultiParamTypeClasses #-} {-# LANGUAGE NamedFieldPuns #-} +{-# LANGUAGE RankNTypes #-} {-# LANGUAGE TypeFamilies #-} {-# LANGUAGE TypeOperators #-} @@ -13,12 +14,19 @@ module TestContainers.Monad TestContainer, runTestContainer, + -- * Blocking during concurrent execution + defer, + -- * Runtime configuration + RunStrategy (..), Config (..), ) where import Control.Applicative (liftA2) +import qualified Control.Concurrent.Async +import Control.Exception (evaluate) +import Control.Monad (sequence_) import Control.Monad.Catch ( MonadCatch, MonadMask, @@ -26,21 +34,42 @@ import Control.Monad.Catch ) import Control.Monad.Fix (MonadFix) import Control.Monad.IO.Class (MonadIO, liftIO) -import Control.Monad.IO.Unlift (MonadUnliftIO (..)) +import Control.Monad.IO.Unlift (MonadUnliftIO (..), askRunInIO) import Control.Monad.Reader (MonadReader (..), ReaderT, runReaderT) -import Control.Monad.Trans.Resource (MonadResource, ResourceT, runResourceT) +import qualified Control.Monad.Trans.Resource import Data.IORef (newIORef, readIORef, writeIORef) +import qualified Data.IORef +import qualified System.IO.Unsafe import TestContainers.Docker.Reaper (Reaper) import TestContainers.Trace (Tracer) -newtype TestContainerEnv = TestContainerEnv - { config :: Config +data RunStrategy + = -- | Run containers sequentially. This is the default behaviour of the runtime. + -- + -- @since x.x.x + SequentialRunStrategy + | -- | Run and resolve graph of containers concurrently. This requires explicit + -- dependency annotations across the container graph. + -- + -- An optional limit on concurrent @docker run@ invocations can be provided. + -- + -- @since x.x.x + ConcurrentRunStrategy (Maybe Int) + +data TestContainerEnv = TestContainerEnv + { config :: Config, + -- | In the concurrent run strategy we need a way to block execution until every + -- container is done running. This barrier is used at the end of the execution and + -- blocks for them to finish. + barrierRef :: Data.IORef.IORef [IO ()] } -- | The heart and soul of the testcontainers library. -- -- @since 0.5.0.0 -newtype TestContainer a = TestContainer {unTestContainer :: ReaderT TestContainerEnv (ResourceT IO) a} +newtype TestContainer a = TestContainer + { unTestContainer :: ReaderT TestContainerEnv (Control.Monad.Trans.Resource.ResourceT IO) a + } deriving newtype ( Functor, Applicative, @@ -49,7 +78,7 @@ newtype TestContainer a = TestContainer {unTestContainer :: ReaderT TestContaine MonadMask, MonadCatch, MonadThrow, - MonadResource, + Control.Monad.Trans.Resource.MonadResource, MonadFix ) @@ -94,17 +123,34 @@ runTestContainer config action = do liftIO (writeIORef reaperRef (Just reaper)) pure reaper - runResourceT - ( runReaderT + barrierRef <- newIORef [] + + Control.Monad.Trans.Resource.runResourceT $ do + result <- + runReaderT (unTestContainer action) ( TestContainerEnv - { config = + { barrierRef, + config = config { configCreateReaper = getOrCreateReaper } } ) - ) + + -- Take the barrier and wait for execution + liftIO $ do + barrier <- readIORef barrierRef + sequence_ barrier + + pure result + +defer :: IO () -> TestContainer () +defer action = TestContainer $ do + TestContainerEnv {barrierRef} <- ask + liftIO $ do + barrier <- readIORef barrierRef + writeIORef barrierRef $! action : barrier -- | Docker related functionality is parameterized over this `Monad`. Since 0.5.0.0 this is -- just a type alias for @m ~ 'TestContainer'@. @@ -125,5 +171,9 @@ data Config = Config -- | Traces execution inside testcontainers library. configTracer :: Tracer, -- | How to obtain a 'Reaper' - configCreateReaper :: TestContainer Reaper + configCreateReaper :: TestContainer Reaper, + -- | + -- + -- @since x.x.x + configRunStrategy :: RunStrategy } diff --git a/src/TestContainers/Tasty.hs b/src/TestContainers/Tasty.hs index c1ce64d..e126ea3 100644 --- a/src/TestContainers/Tasty.hs +++ b/src/TestContainers/Tasty.hs @@ -43,7 +43,25 @@ import Test.Tasty.Options import TestContainers as Reexports hiding ( Trace, ) -import TestContainers.Monad (runTestContainer) +import TestContainers.Monad (RunStrategy (..), runTestContainer) + +newtype ConcurrentRunStrategyOption = ConcurrentRunStrategyOption Bool + +instance IsOption ConcurrentRunStrategyOption where + defaultValue = + ConcurrentRunStrategyOption False + + parseValue = + const Nothing + + optionCLParser = + mkFlagCLParser mempty (ConcurrentRunStrategyOption True) + + optionName = + pure "testcontainers-run-concurrently" + + optionHelp = + pure "Execute 'docker run' concurrently, where possible" newtype DefaultTimeout = DefaultTimeout (Maybe Int) @@ -91,7 +109,8 @@ ingredient :: Ingredient ingredient = Tasty.includingOptions [ Option (Proxy :: Proxy DefaultTimeout), - Option (Proxy :: Proxy Trace) + Option (Proxy :: Proxy Trace), + Option (Proxy :: Proxy ConcurrentRunStrategyOption) ] withContainers :: @@ -102,47 +121,53 @@ withContainers :: withContainers startContainers tests = askOption $ \(DefaultTimeout defaultTimeout) -> askOption $ \(Trace enableTrace) -> - let tracer :: Tracer - tracer - | enableTrace = newTracer $ \message -> - putStrLn (show message) - | otherwise = - mempty - - runC action = do - config <- determineConfig - - let actualConfig :: Config - actualConfig = - config - { configDefaultWaitTimeout = - defaultTimeout <|> configDefaultWaitTimeout config, - configTracer = tracer - } - - runTestContainer actualConfig action - - -- Correct resource handling is tricky here: - -- Tasty offers a bracket alike in IO. We have - -- to transfer the ReleaseMap of the ResIO safely - -- to the release function. Fortunately resourcet - -- let's us access the internal state.. - acquire :: IO (a, InternalState) - acquire = runC $ do - result <- startContainers - releaseMap <- liftResourceT getInternalState - - -- N.B. runResourceT runs the finalizers on every - -- resource. We don't want it to! We want to run - -- finalization in the release function that is - -- called by Tasty! stateAlloc increments a references - -- count to accomodate for exactly these kind of - -- cases. - liftIO $ stateAlloc releaseMap - pure (result, releaseMap) - - release :: (a, InternalState) -> IO () - release (_, internalState) = - stateCleanup ReleaseNormal internalState - in withResource acquire release $ \mk -> - tests (fmap fst mk) + askOption $ \(ConcurrentRunStrategyOption concurrentRunStrategy) -> + let tracer :: Tracer + tracer + | enableTrace = newTracer $ \message -> + putStrLn (show message) + | otherwise = + mempty + + runC action = do + config <- determineConfig + + let actualConfig :: Config + actualConfig = + config + { configDefaultWaitTimeout = + defaultTimeout <|> configDefaultWaitTimeout config, + configTracer = + tracer, + configRunStrategy = + if concurrentRunStrategy + then ConcurrentRunStrategy Nothing + else SequentialRunStrategy + } + + runTestContainer actualConfig action + + -- Correct resource handling is tricky here: + -- Tasty offers a bracket alike in IO. We have + -- to transfer the ReleaseMap of the ResIO safely + -- to the release function. Fortunately resourcet + -- let's us access the internal state.. + acquire :: IO (a, InternalState) + acquire = runC $ do + result <- startContainers + releaseMap <- liftResourceT getInternalState + + -- N.B. runResourceT runs the finalizers on every + -- resource. We don't want it to! We want to run + -- finalization in the release function that is + -- called by Tasty! stateAlloc increments a references + -- count to accomodate for exactly these kind of + -- cases. + liftIO $ stateAlloc releaseMap + pure (result, releaseMap) + + release :: (a, InternalState) -> IO () + release (_, internalState) = + stateCleanup ReleaseNormal internalState + in withResource acquire release $ \mk -> + tests (fmap fst mk) diff --git a/src/TestContainers/Trace.hs b/src/TestContainers/Trace.hs index 01cb3c8..afb4174 100644 --- a/src/TestContainers/Trace.hs +++ b/src/TestContainers/Trace.hs @@ -38,6 +38,8 @@ data Trace TraceOpenSocket Text Int (Maybe IOException) | -- | Call HTTP endpoint TraceHttpCall Text Int (Either String Int) + | -- | Waiting on a dependency + TraceWaitOnDependency deriving stock (Eq, Show) -- | Traces execution within testcontainers library. diff --git a/test/TestContainers/TastySpec.hs b/test/TestContainers/TastySpec.hs index 854822f..b0fc1d2 100644 --- a/test/TestContainers/TastySpec.hs +++ b/test/TestContainers/TastySpec.hs @@ -1,3 +1,4 @@ +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE FlexibleContexts #-} {-# LANGUAGE OverloadedStrings #-} @@ -27,6 +28,7 @@ import TestContainers.Tasty waitUntilMappedPortReachable, waitUntilTimeout, withContainers, + withDependencies, withFollowLogs, withNetwork, (&), @@ -48,7 +50,7 @@ containers1 = do waitUntilMappedPortReachable 6379 ) - _rabbitmq <- + rabbitmq <- run $ containerRequest (fromTag "rabbitmq:3.8.4") & setRm False @@ -60,7 +62,7 @@ containers1 = do <> waitUntilMappedPortReachable 5672 ) - _nginx <- + nginx <- run $ containerRequest (fromTag "nginx:1.23.1-alpine") & setExpose [80] @@ -75,6 +77,7 @@ containers1 = do containerRequest (fromTag "jaegertracing/all-in-one:1.6") & setExpose ["5775/udp", "6831/udp", "6832/udp", "5778", "16686/tcp"] & withNetwork net + & withDependencies [nginx] & setWaitingFor (waitForHttp "16686/tcp" "/" [200]) @@ -83,7 +86,11 @@ containers1 = do containerRequest (fromTag "hello-world:latest") & setWaitingFor (waitForState successfulExit) - _test <- run $ containerRequest (fromBuildContext "./test/container1" Nothing) + _test <- + run $ + containerRequest (fromBuildContext "./test/container1" Nothing) + & withDependencies [rabbitmq] + & setWaitingFor (waitForState successfulExit) pure () diff --git a/testcontainers.cabal b/testcontainers.cabal index 5bba371..06b08ed 100644 --- a/testcontainers.cabal +++ b/testcontainers.cabal @@ -41,6 +41,9 @@ library TestContainers.Tasty TestContainers.Trace + default-extensions: + DuplicateRecordFields + -- other-modules: -- other-extensions: build-depends: