From 6c1a19e826f48ae47187ee534f2ba03b5b06d060 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Domen=20Ko=C5=BEar?= <domen@dev.si>
Date: Sat, 30 Dec 2023 19:21:57 +0000
Subject: [PATCH] use stamina for retries

---
 README.md                               |  2 +-
 src/Network/WebSockets/Simple/Client.hs | 22 ++++++++++++----------
 stack.yaml                              |  1 +
 websockets-simple.cabal                 |  3 ++-
 4 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/README.md b/README.md
index 4032be4..8f7d89b 100644
--- a/README.md
+++ b/README.md
@@ -9,7 +9,7 @@ High-level interface for websockets.
 - Uses channels for sending/receiving messages
 - URI based API supporting `ws://` and `wss://` (for the client)
 - Simple interface between client and server send/receive types
-- TODO: retrying client connection
+- Retries connection for the client automatically using [Stamina](https://github.com/cachix/stamina.hs)
 - TODO: graceful shutdown
 
 ## Example
diff --git a/src/Network/WebSockets/Simple/Client.hs b/src/Network/WebSockets/Simple/Client.hs
index 496c6b0..76945ac 100644
--- a/src/Network/WebSockets/Simple/Client.hs
+++ b/src/Network/WebSockets/Simple/Client.hs
@@ -14,33 +14,35 @@ import Network.WebSockets qualified as WS
 import Network.WebSockets.Connection.PingPong qualified as PingPong
 import Network.WebSockets.Simple.Session qualified as Session
 import Network.WebSockets.Simple.Utils qualified as Utils
+import Stamina qualified
 import Wuss qualified
 
 data Options = Options
   { headers :: WS.Headers,
-    messageLimit :: Int
+    messageLimit :: Int,
+    staminaSettings :: Stamina.RetrySettings
   }
 
 defaultOptions :: Options
 defaultOptions =
   Options
     { headers = [],
-      messageLimit = 10000
+      messageLimit = 10000,
+      staminaSettings = Stamina.defaults
     }
 
 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
-  if isSecure
-    then Wuss.runSecureClientWith (unpack host) (fromIntegral port) (unpack path) connectionOptions (headers options) go
-    else WS.runClientWith (unpack host) (fromIntegral port) (unpack path) connectionOptions (headers options) go
+  Stamina.retry (staminaSettings 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)
   where
     connectionOptions :: WS.ConnectionOptions
     connectionOptions = WS.defaultConnectionOptions
 
-    go :: WS.ClientApp ()
-    go connection = do
+    go :: Stamina.RetryStatus -> WS.ClientApp ()
+    go retryStatus connection = do
+      Stamina.resetInitial retryStatus
       PingPong.withPingPong WS.defaultPingPongOptions connection (\conn -> Session.run (messageLimit options) conn app receiveApp)
-
--- TODO: shutdown :: IO ()
--- TODO: retries
diff --git a/stack.yaml b/stack.yaml
index 20e21eb..2bbccb9 100644
--- a/stack.yaml
+++ b/stack.yaml
@@ -7,3 +7,4 @@ nix:
 extra-deps:
 - wuss-2.0.1.7
 - websockets-0.13.0.0
+- stamina-0.1.0.0
\ No newline at end of file
diff --git a/websockets-simple.cabal b/websockets-simple.cabal
index f18edfa..aa132c0 100644
--- a/websockets-simple.cabal
+++ b/websockets-simple.cabal
@@ -26,7 +26,8 @@ common common
                    wuss >= 2.0.1.7,
                    unliftio-core,
                    bytestring,
-                   exceptions
+                   exceptions,
+                   stamina
                    
     default-extensions: OverloadedStrings