{-# 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. This type is exposed to allow users to
-- create their own 'Reapers'.
--
-- @since 0.5.0.0
data Reaper = Reaper
  { -- | Registers a @label = value@ pair for reaping. Reaping happens when
    -- closing/de-allocating of the 'Reaper' through 'MonadResource'.
    Reaper -> Text -> Text -> IO ()
register ::
      -- \| Label
      Text ->
      -- \| Value
      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.5.0.0
reaperLabels :: Reaper -> [(Text, Text)]
reaperLabels :: Reaper -> [(Text, Text)]
reaperLabels Reaper {[(Text, Text)]
labels :: Reaper -> [(Text, Text)]
labels :: [(Text, Text)]
labels} =
  [(Text, Text)]
labels

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

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

-- | Exposed port for the ryuk reaper.
--
-- @since 0.5.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.5.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 (String -> Text) -> m String -> m Text
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Int -> m Char -> m String
forall (m :: * -> *) a. Applicative m => Int -> m a -> m [a]
replicateM Int
16 ((Char, Char) -> m Char
forall a (m :: * -> *). (Random a, MonadIO m) => (a, a) -> m a
Random.randomRIO (Char
'a', Char
'z'))

  (ReleaseKey
_releaseKey, (Socket
_socket, Reaper
ryuk)) <-
    IO (Socket, Reaper)
-> ((Socket, Reaper) -> IO ()) -> m (ReleaseKey, (Socket, Reaper))
forall (m :: * -> *) a.
MonadResource m =>
IO a -> (a -> IO ()) -> m (ReleaseKey, a)
allocate
      ( do
          let hints :: AddrInfo
hints =
                AddrInfo
Socket.defaultHints {Socket.addrSocketType = Socket.Stream}
          AddrInfo
address <-
            [AddrInfo] -> AddrInfo
forall a. HasCallStack => [a] -> a
head ([AddrInfo] -> AddrInfo) -> IO [AddrInfo] -> IO AddrInfo
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> Maybe AddrInfo -> Maybe String -> Maybe String -> IO [AddrInfo]
Socket.getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
hints) (String -> Maybe String
forall a. a -> Maybe a
Just (Text -> String
unpack Text
host)) (String -> Maybe String
forall a. a -> Maybe a
Just (Int -> String
forall a. Show a => a -> String
show Int
port))
          Socket
socket <-
            Family -> SocketType -> ProtocolNumber -> IO Socket
Socket.socket
              (AddrInfo -> Family
Socket.addrFamily AddrInfo
address)
              (AddrInfo -> SocketType
Socket.addrSocketType AddrInfo
address)
              (AddrInfo -> ProtocolNumber
Socket.addrProtocol AddrInfo
address)
          Socket -> SockAddr -> IO ()
Socket.connect Socket
socket (AddrInfo -> SockAddr
Socket.addrAddress AddrInfo
address)

          -- Construct the reaper and regiter the session with it.
          -- Doing it here intead of in the teardown (like we did before)
          -- guarantees the Reaper knows about our session.
          let reaper :: Reaper
reaper =
                Text -> Ryuk -> Reaper
newReaper Text
sessionId (Socket -> Ryuk
Ryuk Socket
socket)
          Reaper -> Text -> Text -> IO ()
register Reaper
reaper Text
sessionIdLabel Text
sessionId

          (Socket, Reaper) -> IO (Socket, Reaper)
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Socket
socket, Reaper
reaper)
      )
      ( \(Socket
socket, Reaper
_ryuk) -> do
          -- Tearing down the connection lets Ryuk know it can reap the
          -- running containers.
          Socket -> IO ()
Socket.close Socket
socket
      )

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

newReaper ::
  -- | Session id
  Text ->
  Ryuk ->
  Reaper
newReaper :: Text -> Ryuk -> Reaper
newReaper Text
sessionId Ryuk
ryuk =
  Reaper
    { register :: Text -> Text -> IO ()
register = \Text
label Text
value -> do
        Socket -> ByteString -> IO ()
Socket.sendAll
          (Ryuk -> Socket
ryukSocket Ryuk
ryuk)
          (ByteString
"label=" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Text -> ByteString
encodeUtf8 Text
label ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"=" ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> Text -> ByteString
encodeUtf8 Text
value ByteString -> ByteString -> ByteString
forall a. Semigroup a => a -> a -> a
<> ByteString
"\n")
        ByteString
_ <- Socket -> Int -> IO ByteString
Socket.recv (Ryuk -> Socket
ryukSocket Ryuk
ryuk) Int
2
        () -> IO ()
forall a. a -> IO a
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"