-------------------------------------------------------------------------------
-- network-wait
-- Copyright 2022 Michael B. Gale (github@michael-gale.co.uk)
-------------------------------------------------------------------------------

{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TypeApplications #-}

-- | This module contains computations which wait for some networked service
-- to become available, subject to some retry policy from "Control.Retry".
-- The `waitSocketWith` function is the most general function exported by
-- this module, but several variants exist for convenience. You may wish
-- to start out with e.g. `waitTcp` or `waitSocket` initially and move
-- on to the more feature-rich variants if you need their functionality.
module Network.Wait (
    waitTcp,
    waitTcpVerbose,
    waitTcpVerboseFormat,
    waitTcpWith,
    waitSocket,
    waitSocketVerbose,
    waitSocketVerboseFormat,
    waitSocketWith
) where

-------------------------------------------------------------------------------

import Control.Monad.Catch
import Control.Monad.IO.Class
import Control.Retry

import Network.Socket

-------------------------------------------------------------------------------

-- | `waitTcp` @retryPolicy hostName serviceName@ is a variant of `waitTcpWith`
-- which does not install any additional handlers.
--
-- > waitTcp retryPolicyDefault "localhost" "80"
waitTcp
    :: (MonadIO m, MonadMask m)
    => RetryPolicyM m -> HostName -> ServiceName -> m ()
waitTcp :: RetryPolicyM m -> HostName -> HostName -> m ()
waitTcp = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m ()
forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m ()
waitTcpWith []

-- | `waitTcpVerbose` @outputHandler retryPolicy addrInfo@ is a variant
-- of `waitTcpVerboseFormat` which catches all exceptions derived from
-- `SomeException` and formats retry attempt information using `defaultLogMsg`
-- before passing the resulting `String` to @out@.
--
-- > waitTcpVerbose putStrLn retryPolicyDefault "localhost" "80"
waitTcpVerbose
    :: (MonadIO m, MonadMask m)
    => (String -> m ()) -> RetryPolicyM m -> HostName -> ServiceName -> m ()
waitTcpVerbose :: (HostName -> m ())
-> RetryPolicyM m -> HostName -> HostName -> m ()
waitTcpVerbose HostName -> m ()
out =
    forall e (m :: * -> *).
(MonadIO m, MonadMask m, Exception e) =>
(Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> HostName -> HostName -> m ()
forall (m :: * -> *).
(MonadIO m, MonadMask m, Exception SomeException) =>
(Bool -> SomeException -> RetryStatus -> m ())
-> RetryPolicyM m -> HostName -> HostName -> m ()
waitTcpVerboseFormat @SomeException ((Bool -> SomeException -> RetryStatus -> m ())
 -> RetryPolicyM m -> HostName -> HostName -> m ())
-> (Bool -> SomeException -> RetryStatus -> m ())
-> RetryPolicyM m
-> HostName
-> HostName
-> m ()
forall a b. (a -> b) -> a -> b
$
    \Bool
b SomeException
ex RetryStatus
st -> HostName -> m ()
out (HostName -> m ()) -> HostName -> m ()
forall a b. (a -> b) -> a -> b
$ Bool -> SomeException -> RetryStatus -> HostName
forall e. Exception e => Bool -> e -> RetryStatus -> HostName
defaultLogMsg Bool
b SomeException
ex RetryStatus
st

-- | `waitTcpVerboseFormat` @outputHandler retryPolicy addrInfo@ is a
-- variant of `waitTcpWith` which installs an extra handler based on
-- `logRetries` which passes status information for each retry attempt
-- to @outputHandler@.
--
-- > waitTcpVerboseFormat @SomeException
-- >      (\b ex st -> putStrLn $ defaultLogMsg b ex st)
-- >      retryPolicyDefault "localhost" "80"
waitTcpVerboseFormat
    :: forall e m . (MonadIO m, MonadMask m, Exception e)
    => (Bool -> e -> RetryStatus -> m ())
    -> RetryPolicyM m
    -> HostName
    -> ServiceName
    -> m ()
waitTcpVerboseFormat :: (Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> HostName -> HostName -> m ()
waitTcpVerboseFormat Bool -> e -> RetryStatus -> m ()
out = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m ()
forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m ()
waitTcpWith [RetryStatus -> Handler m Bool
h]
    where h :: RetryStatus -> Handler m Bool
h = (e -> m Bool)
-> (Bool -> e -> RetryStatus -> m ())
-> RetryStatus
-> Handler m Bool
forall (m :: * -> *) e.
(Monad m, Exception e) =>
(e -> m Bool)
-> (Bool -> e -> RetryStatus -> m ())
-> RetryStatus
-> Handler m Bool
logRetries (m Bool -> e -> m Bool
forall a b. a -> b -> a
const (m Bool -> e -> m Bool) -> m Bool -> e -> m Bool
forall a b. (a -> b) -> a -> b
$ Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True) Bool -> e -> RetryStatus -> m ()
out

-- | `waitTcpWith` @extraHandlers retryPolicy hostName serviceName@ is a
-- variant of `waitSocketWith` which constructs a suitable `AddrInfo` value
-- for a TCP socket from @hostName@ and @serviceName@.
waitTcpWith
    :: (MonadIO m, MonadMask m)
    => [RetryStatus -> Handler m Bool]
    -> RetryPolicyM m -> HostName -> ServiceName -> m ()
waitTcpWith :: [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m ()
waitTcpWith [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy HostName
host HostName
port = do
    let hints :: AddrInfo
hints = AddrInfo
defaultHints { addrSocketType :: SocketType
addrSocketType = SocketType
Stream }
    AddrInfo
addr <- [AddrInfo] -> AddrInfo
forall a. [a] -> a
head ([AddrInfo] -> AddrInfo) -> m [AddrInfo] -> m AddrInfo
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> IO [AddrInfo] -> m [AddrInfo]
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Maybe AddrInfo -> Maybe HostName -> Maybe HostName -> IO [AddrInfo]
getAddrInfo (AddrInfo -> Maybe AddrInfo
forall a. a -> Maybe a
Just AddrInfo
hints) (HostName -> Maybe HostName
forall a. a -> Maybe a
Just HostName
host) (HostName -> Maybe HostName
forall a. a -> Maybe a
Just HostName
port))
    [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m ()
forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m ()
waitSocketWith [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy AddrInfo
addr

-- | `waitSocket` @retryPolicy addrInfo@ is a variant of `waitSocketWith` which
-- does not install any additional exception handlers.
waitSocket
    :: (MonadIO m, MonadMask m)
    => RetryPolicyM m -> AddrInfo -> m ()
waitSocket :: RetryPolicyM m -> AddrInfo -> m ()
waitSocket = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m ()
forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m ()
waitSocketWith []

-- | `waitSocketVerbose` @outputHandler retryPolicy addrInfo@ is a variant
-- of `waitSocketVerboseFormat` which catches all exceptions derived from
-- `SomeException` and formats retry attempt information using `defaultLogMsg`
-- before passing the resulting `String` to @out@.
waitSocketVerbose
    :: (MonadIO m, MonadMask m)
    => (String -> m ()) -> RetryPolicyM m -> AddrInfo -> m ()
waitSocketVerbose :: (HostName -> m ()) -> RetryPolicyM m -> AddrInfo -> m ()
waitSocketVerbose HostName -> m ()
out =
    forall e (m :: * -> *).
(MonadIO m, MonadMask m, Exception e) =>
(Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> AddrInfo -> m ()
forall (m :: * -> *).
(MonadIO m, MonadMask m, Exception SomeException) =>
(Bool -> SomeException -> RetryStatus -> m ())
-> RetryPolicyM m -> AddrInfo -> m ()
waitSocketVerboseFormat @SomeException ((Bool -> SomeException -> RetryStatus -> m ())
 -> RetryPolicyM m -> AddrInfo -> m ())
-> (Bool -> SomeException -> RetryStatus -> m ())
-> RetryPolicyM m
-> AddrInfo
-> m ()
forall a b. (a -> b) -> a -> b
$
    \Bool
b SomeException
ex RetryStatus
st -> HostName -> m ()
out (HostName -> m ()) -> HostName -> m ()
forall a b. (a -> b) -> a -> b
$ Bool -> SomeException -> RetryStatus -> HostName
forall e. Exception e => Bool -> e -> RetryStatus -> HostName
defaultLogMsg Bool
b SomeException
ex RetryStatus
st

-- | `waitSocketVerboseFormat` @outputHandler retryPolicy addrInfo@ is a
-- variant of `waitSocketWith` which installs an extra handler based on
-- `logRetries` which passes status information for each retry attempt
-- to @outputHandler@.
waitSocketVerboseFormat
    :: forall e m . (MonadIO m, MonadMask m, Exception e)
    => (Bool -> e -> RetryStatus -> m ())
    -> RetryPolicyM m
    -> AddrInfo
    -> m ()
waitSocketVerboseFormat :: (Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> AddrInfo -> m ()
waitSocketVerboseFormat Bool -> e -> RetryStatus -> m ()
out = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m ()
forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m ()
waitSocketWith [RetryStatus -> Handler m Bool
h]
    where h :: RetryStatus -> Handler m Bool
h = (e -> m Bool)
-> (Bool -> e -> RetryStatus -> m ())
-> RetryStatus
-> Handler m Bool
forall (m :: * -> *) e.
(Monad m, Exception e) =>
(e -> m Bool)
-> (Bool -> e -> RetryStatus -> m ())
-> RetryStatus
-> Handler m Bool
logRetries (m Bool -> e -> m Bool
forall a b. a -> b -> a
const (m Bool -> e -> m Bool) -> m Bool -> e -> m Bool
forall a b. (a -> b) -> a -> b
$ Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True) Bool -> e -> RetryStatus -> m ()
out

-- | `waitSocketWith` @extraHandlers retryPolicy addrInfo@ will attempt to
-- connect to @addrInfo@. If the connection fails, @retryPolicy@ is used
-- to determine whether (and how often) this function should attempt to
-- retry establishing the connection. By default, this function will
-- retry after all exceptions (except for those given by
-- `skipAsyncExceptions`). This behaviour may be customised with
-- @extraHandlers@ which are installed after `skipAsyncExceptions`, but
-- before the default exception handler. The @extraHandlers@ may also
-- be used to report retry attempts to e.g. the standard output or a
-- logger.
waitSocketWith
    :: (MonadIO m, MonadMask m)
    => [RetryStatus -> Handler m Bool] -> RetryPolicyM m -> AddrInfo -> m ()
waitSocketWith :: [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m ()
waitSocketWith [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy AddrInfo
addr =
    -- apply the retry policy to the following code, with the combinations of
    -- the `skipAsyncExceptions`, given, and default handlers. The order of
    -- the handlers matters as they are checked in order.
    RetryPolicyM m
-> [RetryStatus -> Handler m Bool] -> (RetryStatus -> m ()) -> m ()
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
RetryPolicyM m
-> [RetryStatus -> Handler m Bool] -> (RetryStatus -> m a) -> m a
recovering RetryPolicyM m
policy ([RetryStatus -> Handler m Bool]
forall (m :: * -> *). MonadIO m => [RetryStatus -> Handler m Bool]
skipAsyncExceptions [RetryStatus -> Handler m Bool]
-> [RetryStatus -> Handler m Bool]
-> [RetryStatus -> Handler m Bool]
forall a. Semigroup a => a -> a -> a
<> [RetryStatus -> Handler m Bool]
hs [RetryStatus -> Handler m Bool]
-> [RetryStatus -> Handler m Bool]
-> [RetryStatus -> Handler m Bool]
forall a. Semigroup a => a -> a -> a
<> [RetryStatus -> Handler m Bool
forall (m :: * -> *) p. Applicative m => p -> Handler m Bool
defHandler])((RetryStatus -> m ()) -> m ()) -> (RetryStatus -> m ()) -> m ()
forall a b. (a -> b) -> a -> b
$ \RetryStatus
_ ->
    -- all of the networking code runs in IO
    IO () -> m ()
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO () -> m ()) -> IO () -> m ()
forall a b. (a -> b) -> a -> b
$
    -- we want to make sure that we close the socket after every attempt;
    -- `bracket` will re-throw any error afterwards
    IO Socket -> (Socket -> IO ()) -> (Socket -> IO ()) -> IO ()
forall (m :: * -> *) a c b.
MonadMask m =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket IO Socket
initSocket Socket -> IO ()
close ((Socket -> IO ()) -> IO ()) -> (Socket -> IO ()) -> IO ()
forall a b. (a -> b) -> a -> b
$
        \Socket
sock -> Socket -> SockAddr -> IO ()
connect Socket
sock (AddrInfo -> SockAddr
addrAddress AddrInfo
addr)
    where
        initSocket :: IO Socket
initSocket =
            Family -> SocketType -> ProtocolNumber -> IO Socket
socket (AddrInfo -> Family
addrFamily AddrInfo
addr) (AddrInfo -> SocketType
addrSocketType AddrInfo
addr) (AddrInfo -> ProtocolNumber
addrProtocol AddrInfo
addr)
        -- our default handler, which works with any exception derived from
        -- `SomeException`, and signals that we should retry if allowed by
        -- the retry policy
        defHandler :: p -> Handler m Bool
defHandler p
_ = (SomeException -> m Bool) -> Handler m Bool
forall (m :: * -> *) a e. Exception e => (e -> m a) -> Handler m a
Handler ((SomeException -> m Bool) -> Handler m Bool)
-> (SomeException -> m Bool) -> Handler m Bool
forall a b. (a -> b) -> a -> b
$ \(SomeException
_ :: SomeException) -> Bool -> m Bool
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

-------------------------------------------------------------------------------