From c805fce14d84345011fade88ec12ac2d4ee2933f 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 | 107 ++++++++++++++++++++-------- src/TestContainers/Monad.hs | 36 ++++++++-- src/TestContainers/Tasty.hs | 117 +++++++++++++++++++------------ test/TestContainers/TastySpec.hs | 4 +- testcontainers.cabal | 3 + 7 files changed, 189 insertions(+), 87 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..42b93a8 100644 --- a/src/TestContainers/Docker.hs +++ b/src/TestContainers/Docker.hs @@ -94,6 +94,7 @@ module TestContainers.Docker setLink, setExpose, setWaitingFor, + withDependencies, run, -- * Following logs @@ -158,7 +159,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 +171,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 +212,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, ) @@ -288,7 +292,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 +329,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 +459,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 +569,8 @@ run request = do labels, noReaper, followLogs, - workDirectory + workDirectory, + dependencies } = request config@Config {configTracer, configCreateReaper} <- @@ -599,38 +613,69 @@ run request = do ++ [[tag]] ++ [command | Just command <- [cmd]] - stdout <- docker configTracer dockerRun + actuallyRunDocker <- applyRunStrategy $ do + _ <- + liftIO $ + traverse evaluate dependencies - let id :: ContainerId - !id = - -- N.B. Force to not leak STDOUT String - strip (pack stdout) + stdout <- docker configTracer dockerRun - -- Careful, this is really meant to be lazy - ~inspectOutput = - unsafePerformIO $ - internalInspect configTracer id + let id :: ContainerId + !id = + -- N.B. Force to not leak STDOUT String + strip (pack stdout) - -- We don't issue 'ReleaseKeys' for cleanup anymore. Ryuk takes care of cleanup - -- for us once the session has been closed. - releaseKey <- register (pure ()) + -- Careful, this is really meant to be lazy + ~inspectOutput = + unsafePerformIO $ + internalInspect configTracer id - forM_ followLogs $ - dockerFollowLogs 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 ()) - let container = - Container - { id, - releaseKey, - image, - inspectOutput, - config - } + forM_ followLogs $ + dockerFollowLogs configTracer id - -- Last but not least, execute the WaitUntilReady checks - waitUntilReady container readiness + let container = + Container + { id, + releaseKey, + image, + inspectOutput, + config + } - pure container + -- Last but not least, execute the WaitUntilReady checks + waitUntilReady container readiness + + pure container + + actuallyRunDocker + +applyRunStrategy :: + TestContainer a -> + TestContainer (TestContainer a) +applyRunStrategy action = do + Config {configRunStrategy} <- + ask + case configRunStrategy of + SequentialRunStrategy -> + pure action + ConcurrentRunStrategy _limit -> do + runInIO <- + askRunInIO + + (_releaseKey, handle) <- + allocate + (Control.Concurrent.Async.async (runInIO action)) + Control.Concurrent.Async.cancel + + let returnAction = + liftIO $ unsafeInterleaveIO $ do + Control.Concurrent.Async.wait handle + + pure returnAction -- | Sets up a Ryuk 'Reaper'. -- diff --git a/src/TestContainers/Monad.hs b/src/TestContainers/Monad.hs index 86eaf11..3c3f38f 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 #-} @@ -14,11 +15,14 @@ module TestContainers.Monad runTestContainer, -- * Runtime configuration + RunStrategy (..), Config (..), ) where import Control.Applicative (liftA2) +import qualified Control.Concurrent.Async +import Control.Exception (evaluate) import Control.Monad.Catch ( MonadCatch, MonadMask, @@ -26,13 +30,27 @@ 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 System.IO.Unsafe import TestContainers.Docker.Reaper (Reaper) import TestContainers.Trace (Tracer) +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) + newtype TestContainerEnv = TestContainerEnv { config :: Config } @@ -40,7 +58,9 @@ newtype TestContainerEnv = TestContainerEnv -- | 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 +69,7 @@ newtype TestContainer a = TestContainer {unTestContainer :: ReaderT TestContaine MonadMask, MonadCatch, MonadThrow, - MonadResource, + Control.Monad.Trans.Resource.MonadResource, MonadFix ) @@ -94,7 +114,7 @@ runTestContainer config action = do liftIO (writeIORef reaperRef (Just reaper)) pure reaper - runResourceT + Control.Monad.Trans.Resource.runResourceT ( runReaderT (unTestContainer action) ( TestContainerEnv @@ -125,5 +145,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/test/TestContainers/TastySpec.hs b/test/TestContainers/TastySpec.hs index 854822f..f4270ca 100644 --- a/test/TestContainers/TastySpec.hs +++ b/test/TestContainers/TastySpec.hs @@ -27,6 +27,7 @@ import TestContainers.Tasty waitUntilMappedPortReachable, waitUntilTimeout, withContainers, + withDependencies, withFollowLogs, withNetwork, (&), @@ -60,7 +61,7 @@ containers1 = do <> waitUntilMappedPortReachable 5672 ) - _nginx <- + nginx <- run $ containerRequest (fromTag "nginx:1.23.1-alpine") & setExpose [80] @@ -75,6 +76,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]) 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: