From 4f495aa68df49084a24cc4d4bd9defa2238b2e07 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Domen=20Ko=C5=BEar?= <domen@dev.si>
Date: Mon, 8 Jan 2024 17:00:11 +0000
Subject: [PATCH] Add a hook when retrying

---
 src/Network/WebSockets/Simple/Client.hs | 12 +++++++++---
 1 file changed, 9 insertions(+), 3 deletions(-)

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)