diff --git a/src/Network/WebSockets/Simple/Client.hs b/src/Network/WebSockets/Simple/Client.hs index 76945ac..77cb62e 100644 --- a/src/Network/WebSockets/Simple/Client.hs +++ b/src/Network/WebSockets/Simple/Client.hs @@ -8,8 +8,10 @@ module Network.WebSockets.Simple.Client ) where +import Control.Monad (when) import Data.ByteString (ByteString) import Data.ByteString.Char8 (unpack) +import Data.Maybe (isJust) import Network.WebSockets qualified as WS import Network.WebSockets.Connection.PingPong qualified as PingPong import Network.WebSockets.Simple.Session qualified as Session @@ -20,7 +22,8 @@ import Wuss qualified data Options = Options { headers :: WS.Headers, messageLimit :: Int, - staminaSettings :: Stamina.RetrySettings + staminaSettings :: Stamina.RetrySettings, + staminaRetry :: Stamina.RetryStatus -> IO () } defaultOptions :: Options @@ -28,13 +31,16 @@ defaultOptions = Options { headers = [], messageLimit = 10000, - staminaSettings = Stamina.defaults + staminaSettings = Stamina.defaults, + staminaRetry = const $ return () } run :: (Session.Codec send, Session.Codec receive) => ByteString -> Options -> Session.Session IO send receive () -> (receive -> Session.Session IO send receive ()) -> IO () run uriBS options app receiveApp = do (isSecure, host, port, path) <- Utils.parseURI uriBS - Stamina.retry (staminaSettings options) $ \retryStatus -> + Stamina.retry (staminaSettings options) $ \retryStatus -> do + when (isJust $ Stamina.lastException retryStatus) $ + staminaRetry options retryStatus if isSecure then Wuss.runSecureClientWith (unpack host) (fromIntegral port) (unpack path) connectionOptions (headers options) (go retryStatus) else WS.runClientWith (unpack host) (fromIntegral port) (unpack path) connectionOptions (headers options) (go retryStatus)