From a9a191c4ec21b12b595690e9b2b65f013e6a344a Mon Sep 17 00:00:00 2001
From: Finley McIlwaine <finley@well-typed.com>
Date: Fri, 30 Aug 2024 11:08:49 -0700
Subject: [PATCH] Introduce `ReconnectTo`

Reconnect policies can now specify whether they want to attempt reconnection
with the original server given to `withConnection`, the last server we attempted
connection with, or a new server specified by the policy itself.
---
 src/Network/GRPC/Client/Connection.hs  | 36 +++++++++++++++++++++-----
 test-grapesy/Test/Sanity/Disconnect.hs |  3 +++
 2 files changed, 32 insertions(+), 7 deletions(-)

diff --git a/src/Network/GRPC/Client/Connection.hs b/src/Network/GRPC/Client/Connection.hs
index a9f4ffa..dfd6b24 100644
--- a/src/Network/GRPC/Client/Connection.hs
+++ b/src/Network/GRPC/Client/Connection.hs
@@ -169,15 +169,27 @@ data ReconnectPolicy =
     -- | Reconnect to the (potentially different) server after the IO action
     -- returns
     --
-    -- If the 'Maybe' is 'Just', we'll attempt to reconnect to a server at the
-    -- new address. If 'Nothing', we'll attempt to connect to the original
-    -- server that 'withConnection' was given.
+    -- The 'ReconnectTo' can be used to implement a rudimentary redundancy
+    -- scheme. For example, you could decide to reconnect to a known fallback
+    -- server after connection to a main server fails a certain number of times.
     --
     -- This is a very general API: typically the IO action will call
     -- 'threadDelay' after some amount of time (which will typically involve
     -- some randomness), but it can be used to do things such as display a
     -- message to the user somewhere that the client is reconnecting.
-  | ReconnectAfter (Maybe Server) (IO ReconnectPolicy)
+  | ReconnectAfter ReconnectTo (IO ReconnectPolicy)
+
+-- | What server should we attempt to reconnect to?
+--
+-- * 'ReconnectToPrevious' will attempt to reconnect to the last server we
+--   attempted to connect to, whether or not that attempt was successful.
+-- * 'ReconnectToOriginal' will attempt to reconnect to the original server that
+--   'withConnection' was given.
+-- * 'ReconnectToNew' will attempt to connect to the newly specified server.
+data ReconnectTo =
+      ReconnectToPrevious
+    | ReconnectToOriginal
+    | ReconnectToNew Server
 
 -- | The default policy is 'DontReconnect'
 --
@@ -186,6 +198,9 @@ data ReconnectPolicy =
 instance Default ReconnectPolicy where
   def = DontReconnect
 
+instance Default ReconnectTo where
+  def = ReconnectToPrevious
+
 -- | Exponential backoff
 --
 -- If the exponent is @1@, the delay interval will be the same every step;
@@ -213,7 +228,7 @@ exponentialBackoff waitFor e = go
   where
     go :: (Double, Double) -> Word -> ReconnectPolicy
     go _        0 = DontReconnect
-    go (lo, hi) n = ReconnectAfter Nothing $ do
+    go (lo, hi) n = ReconnectAfter def $ do
         delay <- randomRIO (lo, hi)
         waitFor $ round $ delay * 1_000_000
         return $ go (lo * e, hi * e) (pred n)
@@ -431,9 +446,16 @@ stayConnected connParams initialServer connStateVar connOutOfScope = do
                 atomically $ writeTVar connStateVar $ ConnectionAbandoned err
               (False, DontReconnect) -> do
                 atomically $ writeTVar connStateVar $ ConnectionAbandoned err
-              (False, ReconnectAfter mNewServer f) -> do
+                atomically $ writeTVar connStateVar $ ConnectionAbandoned err
+              (False, ReconnectAfter to f) -> do
+                let
+                  nextServer =
+                    case to of
+                      ReconnectToPrevious -> server
+                      ReconnectToOriginal -> initialServer
+                      ReconnectToNew new  -> new
                 atomically $ writeTVar connStateVar $ ConnectionNotReady
-                loop (fromMaybe initialServer mNewServer) =<< f
+                loop nextServer =<< f
 
 -- | Insecure connection (no TLS)
 connectInsecure :: ConnParams -> Attempt -> Address -> IO ()
diff --git a/test-grapesy/Test/Sanity/Disconnect.hs b/test-grapesy/Test/Sanity/Disconnect.hs
index ea7c44a..26c1a1a 100644
--- a/test-grapesy/Test/Sanity/Disconnect.hs
+++ b/test-grapesy/Test/Sanity/Disconnect.hs
@@ -296,6 +296,9 @@ echoHandler disconnectCounter call = trackDisconnects disconnectCounter $ do
   Auxiliary
 -------------------------------------------------------------------------------}
 
+-- We need to use this to properly simulate the execution environment crashing
+-- in an unrecoverable way. In particular, we don't want to give the program a
+-- chance to do any of its normal exception handling/cleanup behavior.
 foreign import ccall unsafe "exit" c_exit :: CInt -> IO ()
 
 data ClientStep = KeepGoing (Maybe (IO ())) ClientStep | Done