Merge pull request #20 from codedownio/connect-timeout

Add timeout to connect attempts to fix #19

Author: mbg

src/Network/Wait.hs

@@ -32,16 +32,39 @@ module Network.Wait (
 
 -------------------------------------------------------------------------------
 
+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 = capDelay (3000000) (fullJitterBackoff 100000)
+
 -- | `waitTcp` @retryPolicy hostName serviceName@ is a variant of `waitTcpWith`
 -- which does not install any additional handlers.
 --
@@ -141,13 +164,21 @@ waitSocketWith
     => [RetryStatus -> Handler m Bool] -> RetryPolicyM m -> AddrInfo
     -> m Socket
 waitSocketWith hs policy addr =
-    recoveringWith hs policy $
+    recoveringWith hs policy $ \retryStatus ->
     -- all of the networking code runs in IO
     liftIO $
     -- we want to make sure that we close the socket after every attempt;
     -- `bracket` will re-throw any error afterwards
-    bracket initSocket close $
-        \sock -> connect sock (addrAddress addr) >> pure sock
+    bracket initSocket close $ \sock -> do
+        maybeConnectTimeoutUs <- (getRetryPolicyM connectRetryPolicy) retryStatus
+        connectTimeoutUs <- case maybeConnectTimeoutUs of
+            Nothing -> throwIO $ userError "Timeout in connect attempt"
+            Just us -> pure us
+
+        maybeResult <- timeout connectTimeoutUs (connect sock (addrAddress addr))
+        case maybeResult of
+            Nothing -> throwIO $ userError "Timeout in connect attempt"
+            Just () -> pure sock
     where
         initSocket =
             socket (addrFamily addr) (addrSocketType addr) (addrProtocol addr)
@@ -163,15 +194,13 @@ waitSocketWith hs policy addr =
 -- the standard output or a logger.
 recoveringWith
     :: (MonadIO m, MonadMask m)
-    => [RetryStatus -> Handler m Bool] -> RetryPolicyM m -> m a -> m a
+    => [RetryStatus -> Handler m Bool] -> RetryPolicyM m -> (RetryStatus -> m a) -> m a
 recoveringWith hs policy 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.
     recovering policy (skipAsyncExceptions <> hs <> [defHandler]) $
-    -- we want to make sure that we close the socket after every attempt;
-    -- `bracket` will re-throw any error afterwards
-        const action
+        action
     where
         -- our default handler, which works with any exception derived from
         -- `SomeException`, and signals that we should retry if allowed by

src/Network/Wait/PostgreSQL.hs

@@ -93,7 +93,7 @@ waitPostgreSqlWith
     => [RetryStatus -> Handler m Bool] -> RetryPolicyM m -> info
     -> m Connection
 waitPostgreSqlWith hs policy info =
-    recoveringWith hs policy $
+    recoveringWith hs policy $ \_ ->
     liftIO $
     bracket (connectDb info) close $ \con -> do
         rs <- query_ @[Int] con "SELECT 1;"

src/Network/Wait/Redis.hs

@@ -79,6 +79,6 @@ waitRedisWith
     :: (MonadIO m, MonadMask m)
     => [RetryStatus -> Handler m Bool] -> RetryPolicyM m -> ConnectInfo
     -> m Connection
-waitRedisWith hs policy = recoveringWith hs policy . liftIO . checkedConnect
+waitRedisWith hs policy = recoveringWith hs policy . const . liftIO . checkedConnect
 
 -------------------------------------------------------------------------------

test/Spec.hs

@@ -54,6 +54,13 @@ tests = testGroup "network-wait"
         res <- try @IO @SomeException $
             waitTcp testRetryPolicy "localhost" "5999"
 
+        case res of
+            Left _ -> pure ()
+            Right _ -> assertFailure "`waitTcp` did not fail"
+    , localOption (mkTimeout $ 10*1000*1000) $ testCase "Can't connect to non-routable private IP" $ do
+        res <- try @IO @SomeException $
+            waitTcp testRetryPolicy "10.255.255.1" "80"
+
         case res of
             Left _ -> pure ()
             Right _ -> assertFailure "`waitTcp` did not fail"