{-# LANGUAGE DerivingStrategies #-}
{-# LANGUAGE GeneralizedNewtypeDeriving #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE OverloadedStrings #-}

module TestContainers.Docker.Reaper
  ( Reaper,
    reaperLabels,

    -- * Ryuk based reaper
    ryukImageTag,
    ryukPort,
    newRyukReaper,
  )
where

import Control.Monad (replicateM)
import Control.Monad.Trans.Resource (MonadResource, allocate)
import Data.Text (Text, pack, unpack)
import Data.Text.Encoding (encodeUtf8)
import qualified Network.Socket as Socket
import qualified Network.Socket.ByteString as Socket
import qualified System.Random as Random

-- | Reaper for safe resource cleanup.
--
-- @since 0.4.0.0
data Reaper = Reaper
  { -- | @runReaper label value@ reaps Docker any Docker resource with a matching
    -- label.
    Reaper -> Text -> Text -> IO ()
runReaper :: Text -> Text -> IO (),
    -- | Additional labels to add to any Docker resource on creation. Adding the
    -- labels is necessary in order for the 'Reaper' to find resources for cleanup.
    Reaper -> [(Text, Text)]
labels :: [(Text, Text)]
  }

-- | Additional labels to add to any Docker resource on creation. Adding the
-- labels is necessary in order for the 'Reaper' to find resources for cleanup.
--
-- @since 0.4.0.0
reaperLabels :: Reaper -> [(Text, Text)]
reaperLabels :: Reaper -> [(Text, Text)]
reaperLabels Reaper {[(Text, Text)]
labels :: [(Text, Text)]
labels :: Reaper -> [(Text, Text)]
labels} =
  [(Text, Text)]
labels

-- | Ryuk based resource reaper
--
-- @since 0.4.0.0
newtype Ryuk = Ryuk {Ryuk -> Socket
ryukSocket :: Socket.Socket}

-- | Tag for the ryuk image
--
-- @since 0.4.0.0
ryukImageTag :: Text
ryukImageTag :: Text
ryukImageTag =
  Text
"docker.io/testcontainers/ryuk:0.3.4"

-- | Exposed port for the ryuk reaper.
--
-- @since 0.4.0.0
ryukPort :: (Num a) => a
ryukPort :: forall a. Num a => a
ryukPort =
  a
8080

-- | Creates a new 'Reaper' from a host and port.
--
-- @since 0.4.0.0
newRyukReaper ::
  (MonadResource m) =>
  -- | Host
  Text ->
  -- | Port
  Int ->
  m Reaper
newRyukReaper :: forall (m :: * -> *). MonadResource m => Text -> Int -> m Reaper
newRyukReaper Text
host Int
port = do
  Text
sessionId <-
    String -> Text
pack forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
16 (forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
Random.randomRIO (Char
'a', Char
'z'))

  (ReleaseKey
_releaseKey, (Socket
_socket, Reaper
ryuk)) <-
    forall (m :: * -> *) a.
MonadResource m =>
IO a -> (a -> IO ()) -> m (ReleaseKey, a)
allocate
      ( do
          let hints :: AddrInfo
hints =
                AddrInfo
Socket.defaultHints {addrSocketType :: SocketType
Socket.addrSocketType = SocketType
Socket.Stream}
          AddrInfo
address <-
            forall a. [a] -> a
head forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe AddrInfo -> Maybe String -> Maybe String -> IO [AddrInfo]
Socket.getAddrInfo (forall a. a -> Maybe a
Just AddrInfo
hints) (forall a. a -> Maybe a
Just (Text -> String
unpack Text
host)) (forall a. a -> Maybe a
Just (forall a. Show a => a -> String
show Int
port))
          Socket
socket <-
            AddrInfo -> IO Socket
Socket.openSocket AddrInfo
address
          Socket -> SockAddr -> IO ()
Socket.connect Socket
socket (AddrInfo -> SockAddr
Socket.addrAddress AddrInfo
address)
          forall (f :: * -> *) a. Applicative f => a -> f a
pure (Socket
socket, Text -> Ryuk -> Reaper
runRyuk Text
sessionId (Socket -> Ryuk
Ryuk Socket
socket))
      )
      ( \(Socket
socket, Reaper
ryuk) -> do
          Reaper -> Text -> Text -> IO ()
runReaper Reaper
ryuk Text
sessionIdLabel Text
sessionId
          Socket -> IO ()
Socket.close Socket
socket
      )

  forall (f :: * -> *) a. Applicative f => a -> f a
pure Reaper
ryuk

runRyuk ::
  -- | Session id
  Text ->
  Ryuk ->
  Reaper
runRyuk :: Text -> Ryuk -> Reaper
runRyuk Text
sessionId Ryuk
ryuk =
  Reaper
    { runReaper :: Text -> Text -> IO ()
runReaper = \Text
label Text
value -> do
        Socket -> ByteString -> IO ()
Socket.sendAll
          (Ryuk -> Socket
ryukSocket Ryuk
ryuk)
          (ByteString
"label=" forall a. Semigroup a => a -> a -> a
<> Text -> ByteString
encodeUtf8 Text
label forall a. Semigroup a => a -> a -> a
<> ByteString
"=" forall a. Semigroup a => a -> a -> a
<> Text -> ByteString
encodeUtf8 Text
value forall a. Semigroup a => a -> a -> a
<> ByteString
"\n")
        ByteString
_ <- Socket -> Int -> IO ByteString
Socket.recv (Ryuk -> Socket
ryukSocket Ryuk
ryuk) Int
2
        forall (f :: * -> *) a. Applicative f => a -> f a
pure (),
      labels :: [(Text, Text)]
labels =
        [ (Text
sessionIdLabel, Text
sessionId)
        ]
    }

sessionIdLabel :: Text
sessionIdLabel :: Text
sessionIdLabel = Text
"org.testcontainers.haskell.session"