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

{-# OPTIONS_GHC -Wno-unused-imports #-}
{-# 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 (
    -- * TCP
    waitTcp,
    waitTcpVerbose,
    waitTcpVerboseFormat,
    waitTcpWith,

    -- * Sockets
    waitSocket,
    waitSocketVerbose,
    waitSocketVerboseFormat,
    waitSocketWith,

    -- * Utility
    recoveringWith,
    recoveringWithStatus
) where

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

import Control.Exception (throwIO)
import Control.Monad.Catch
import Control.Monad.IO.Class
import Control.Retry
-- Only needed for base < 4.11, redundant otherwise
import Data.Semigroup
import System.IO.Error
import System.Timeout

import Network.Socket

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

-- | Each individual connect attempt needs a timeout to prevent it from hanging
-- indefinitely. This policy allows us to make that timeout length adaptive,
-- based on the 'RetryStatus' of the outer retry policy.
--
-- Thus, the first attempt to connect will have a short timeout (currently 100ms),
-- and then successive attempts will get longer timeouts via "FullJitter" backoff.
-- The goals of this are twofold:
--
-- 1) If a connect call hangs during the first few attempts, it is timed out quickly
-- and re-attempted, so on a healthy network you aren't penalized too much by the hang.
-- The outer retry policy can control the time between attempts, so the user can set
-- it high enough to make this be the case.
--
-- 2) If the network is slow, we will eventually reach the maximum timeout of 3 seconds,
-- which should be long enough. Note that the popular wait-for script uses 1 second
-- timeouts, so this is extra conservative:
-- https://github.com/eficode/wait-for/blob/7586b3622f010808bb2027c19aaf367221b4ad54/wait-for#L72
connectRetryPolicy :: MonadIO m => RetryPolicyM m
connectRetryPolicy :: forall (m :: * -> *). MonadIO m => RetryPolicyM m
connectRetryPolicy = Int -> RetryPolicyM m -> RetryPolicyM m
forall (m :: * -> *).
Monad m =>
Int -> RetryPolicyM m -> RetryPolicyM m
capDelay (Int
3000000) (Int -> RetryPolicyM m
forall (m :: * -> *). MonadIO m => Int -> RetryPolicyM m
fullJitterBackoff Int
100000)

-- | `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 Socket
waitTcp :: forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
RetryPolicyM m -> HostName -> HostName -> m Socket
waitTcp = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m Socket
forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m Socket
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 Socket
waitTcpVerbose :: forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
(HostName -> m ())
-> RetryPolicyM m -> HostName -> HostName -> m Socket
waitTcpVerbose HostName -> m ()
out =
    forall e (m :: * -> *).
(MonadIO m, MonadMask m, Exception e) =>
(Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> HostName -> HostName -> m Socket
waitTcpVerboseFormat @SomeException ((Bool -> SomeException -> RetryStatus -> m ())
 -> RetryPolicyM m -> HostName -> HostName -> m Socket)
-> (Bool -> SomeException -> RetryStatus -> m ())
-> RetryPolicyM m
-> HostName
-> HostName
-> m Socket
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 Socket
waitTcpVerboseFormat :: forall e (m :: * -> *).
(MonadIO m, MonadMask m, Exception e) =>
(Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> HostName -> HostName -> m Socket
waitTcpVerboseFormat Bool -> e -> RetryStatus -> m ()
out = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m Socket
forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m Socket
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 a. a -> m a
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 Socket
waitTcpWith :: forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> HostName -> HostName -> m Socket
waitTcpWith [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy HostName
host HostName
port = do
    let hints :: AddrInfo
hints = AddrInfo
defaultHints { addrSocketType = Stream }
    addr <- [AddrInfo] -> AddrInfo
forall a. HasCallStack => [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 a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Maybe AddrInfo -> Maybe HostName -> Maybe HostName -> IO [AddrInfo]
forall (t :: * -> *).
GetAddrInfo t =>
Maybe AddrInfo
-> Maybe HostName -> Maybe HostName -> IO (t 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))
    waitSocketWith hs policy 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 Socket
waitSocket :: forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
RetryPolicyM m -> AddrInfo -> m Socket
waitSocket = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m Socket
forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m Socket
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 Socket
waitSocketVerbose :: forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
(HostName -> m ()) -> RetryPolicyM m -> AddrInfo -> m Socket
waitSocketVerbose HostName -> m ()
out =
    forall e (m :: * -> *).
(MonadIO m, MonadMask m, Exception e) =>
(Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> AddrInfo -> m Socket
waitSocketVerboseFormat @SomeException ((Bool -> SomeException -> RetryStatus -> m ())
 -> RetryPolicyM m -> AddrInfo -> m Socket)
-> (Bool -> SomeException -> RetryStatus -> m ())
-> RetryPolicyM m
-> AddrInfo
-> m Socket
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 Socket
waitSocketVerboseFormat :: forall e (m :: * -> *).
(MonadIO m, MonadMask m, Exception e) =>
(Bool -> e -> RetryStatus -> m ())
-> RetryPolicyM m -> AddrInfo -> m Socket
waitSocketVerboseFormat Bool -> e -> RetryStatus -> m ()
out = [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m Socket
forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m Socket
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 a. a -> m a
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 Socket
waitSocketWith :: forall (m :: * -> *).
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> AddrInfo -> m Socket
waitSocketWith [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy AddrInfo
addr =
    [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> (RetryStatus -> m Socket) -> m Socket
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> (RetryStatus -> m a) -> m a
recoveringWithStatus [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy ((RetryStatus -> m Socket) -> m Socket)
-> (RetryStatus -> m Socket) -> m Socket
forall a b. (a -> b) -> a -> b
$ \RetryStatus
retryStatus ->
    -- all of the networking code runs in IO
    IO Socket -> m Socket
forall a. IO a -> m a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (IO Socket -> m Socket) -> IO Socket -> m Socket
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 Socket) -> IO Socket
forall (m :: * -> *) a c b.
(HasCallStack, MonadMask m) =>
m a -> (a -> m c) -> (a -> m b) -> m b
bracket IO Socket
initSocket Socket -> IO ()
close ((Socket -> IO Socket) -> IO Socket)
-> (Socket -> IO Socket) -> IO Socket
forall a b. (a -> b) -> a -> b
$ \Socket
sock -> do
        maybeConnectTimeoutUs <- (RetryPolicyM IO -> RetryStatus -> IO (Maybe Int)
forall (m :: * -> *).
RetryPolicyM m -> RetryStatus -> m (Maybe Int)
getRetryPolicyM RetryPolicyM IO
forall (m :: * -> *). MonadIO m => RetryPolicyM m
connectRetryPolicy) RetryStatus
retryStatus
        connectTimeoutUs <- case maybeConnectTimeoutUs of
            Maybe Int
Nothing -> IOError -> IO Int
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (IOError -> IO Int) -> IOError -> IO Int
forall a b. (a -> b) -> a -> b
$ HostName -> IOError
userError HostName
"Timeout in connect attempt"
            Just Int
us -> Int -> IO Int
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Int
us

        maybeResult <- timeout connectTimeoutUs (connect sock (addrAddress addr))
        case maybeResult of
            Maybe ()
Nothing -> IOError -> IO Socket
forall e a. (HasCallStack, Exception e) => e -> IO a
throwIO (IOError -> IO Socket) -> IOError -> IO Socket
forall a b. (a -> b) -> a -> b
$ HostName -> IOError
userError HostName
"Timeout in connect attempt"
            Just () -> Socket -> IO Socket
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Socket
sock
    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)

-- | `recoveringWith` @extraHandlers retryPolicy action@ will attempt to
-- run @action@. If the @action@ fails, @retryPolicy@ is used
-- to determine whether (and how often) this function should attempt to
-- retry @action@. 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.
recoveringWith
    :: (MonadIO m, MonadMask m)
    => [RetryStatus -> Handler m Bool] -> RetryPolicyM m -> m a -> m a
recoveringWith :: forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool] -> RetryPolicyM m -> m a -> m a
recoveringWith [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy =
    [RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> (RetryStatus -> m a) -> m a
forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> (RetryStatus -> m a) -> m a
recoveringWithStatus [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy ((RetryStatus -> m a) -> m a)
-> (m a -> RetryStatus -> m a) -> m a -> m a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. m a -> RetryStatus -> m a
forall a b. a -> b -> a
const


-- | `recoveringWithStatus` @extraHandlers retryPolicy action@ will attempt to
-- run @action@. If the @action@ fails, @retryPolicy@ is used
-- to determine whether (and how often) this function should attempt to
-- retry @action@. The `RetryStatus` is given to @action@ as argument.
-- 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.
recoveringWithStatus
    :: (MonadIO m, MonadMask m)
    => [RetryStatus -> Handler m Bool]
    -> RetryPolicyM m
    -> (RetryStatus -> m a)
    -> m a
recoveringWithStatus :: forall (m :: * -> *) a.
(MonadIO m, MonadMask m) =>
[RetryStatus -> Handler m Bool]
-> RetryPolicyM m -> (RetryStatus -> m a) -> m a
recoveringWithStatus [RetryStatus -> Handler m Bool]
hs RetryPolicyM m
policy RetryStatus -> m a
action =
    -- 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 a) -> m a
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 a) -> m a) -> (RetryStatus -> m a) -> m a
forall a b. (a -> b) -> a -> b
$
        RetryStatus -> m a
action
    where
        -- 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 a. a -> m a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Bool
True

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