From d9f1ae6a46ca8d9d48d6399efbe66531697ca7df Mon Sep 17 00:00:00 2001
From: Gaurav Gautam <gautam1168@users.noreply.github.com>
Date: Wed, 18 Aug 2021 20:44:29 +0530
Subject: [PATCH] feat: Add Retry-After hint (#1916)

Add Retry-After header when response status is 503.
Its value is the connection worker delay(seconds) when
it's recovering.

This closes issue #1817.
---
 CHANGELOG.md              |  3 ++-
 src/PostgREST/App.hs      | 14 ++++++++++----
 src/PostgREST/AppState.hs | 11 +++++++++++
 src/PostgREST/Workers.hs  |  1 +
 4 files changed, 24 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index e62fa65c56..77cc7a6409 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,7 +8,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 ### Added
 
  - #1783, Include partitioned tables into the schema cache. Allows embedding, UPSERT, INSERT with Location response, OPTIONS request and OpenAPI support for partitioned tables - @laurenceisla
-
+ - #1878, Add Retry-After hint header when in recovery mode - @gautam1168
+ 
 ### Fixed
 
 ## [8.0.0] - 2021-07-25
diff --git a/src/PostgREST/App.hs b/src/PostgREST/App.hs
index 18c96c6beb..1226aabfe3 100644
--- a/src/PostgREST/App.hs
+++ b/src/PostgREST/App.hs
@@ -152,13 +152,19 @@ postgrest logLev appState connWorker =
           runExceptT $ postgrestResponse conf maybeDbStructure jsonDbS pgVer (AppState.getPool appState) time req
 
       response <- either Error.errorResponseFor identity <$> eitherResponse
-
       -- Launch the connWorker when the connection is down.  The postgrest
       -- function can respond successfully (with a stale schema cache) before
       -- the connWorker is done.
-      when (Wai.responseStatus response == HTTP.status503) connWorker
-
-      respond response
+      let isPGAway = Wai.responseStatus response == HTTP.status503
+      when isPGAway connWorker
+      resp <- addRetryHint isPGAway appState response
+      respond resp
+
+addRetryHint :: Bool -> AppState -> Wai.Response -> IO Wai.Response
+addRetryHint shouldAdd appState response = do
+  delay <- AppState.getRetryNextIn appState
+  let h = ("Retry-After", BS8.pack $ show delay)
+  return $ Wai.mapResponseHeaders (\hs -> if shouldAdd then h:hs else hs) response
 
 postgrestResponse
   :: AppConfig
diff --git a/src/PostgREST/AppState.hs b/src/PostgREST/AppState.hs
index bc9da187e3..964f9b4612 100644
--- a/src/PostgREST/AppState.hs
+++ b/src/PostgREST/AppState.hs
@@ -10,6 +10,7 @@ module PostgREST.AppState
   , getPgVersion
   , getPool
   , getTime
+  , getRetryNextIn
   , init
   , initWithPool
   , logWithZTime
@@ -18,6 +19,7 @@ module PostgREST.AppState
   , putIsWorkerOn
   , putJsonDbS
   , putPgVersion
+  , putRetryNextIn
   , releasePool
   , signalListener
   , waitListener
@@ -60,6 +62,8 @@ data AppState = AppState
   , stateGetZTime     :: IO ZonedTime
   -- | Used for killing the main thread in case a subthread fails
   , stateMainThreadId :: ThreadId
+  -- | Keeps track of when the next retry for connecting to database is scheduled
+  , stateRetryNextIn  :: IORef Int
   }
 
 init :: AppConfig -> IO AppState
@@ -79,6 +83,7 @@ initWithPool newPool conf =
     <*> mkAutoUpdate defaultUpdateSettings { updateAction = getCurrentTime }
     <*> mkAutoUpdate defaultUpdateSettings { updateAction = getZonedTime }
     <*> myThreadId
+    <*> newIORef 0
 
 initPool :: AppConfig -> IO P.Pool
 initPool AppConfig{..} =
@@ -115,6 +120,12 @@ getIsWorkerOn = readIORef . stateIsWorkerOn
 putIsWorkerOn :: AppState -> Bool -> IO ()
 putIsWorkerOn = atomicWriteIORef . stateIsWorkerOn
 
+getRetryNextIn :: AppState -> IO Int
+getRetryNextIn = readIORef . stateRetryNextIn
+
+putRetryNextIn :: AppState -> Int -> IO ()
+putRetryNextIn = atomicWriteIORef . stateRetryNextIn
+
 getConfig :: AppState -> IO AppConfig
 getConfig = readIORef . stateConf
 
diff --git a/src/PostgREST/Workers.hs b/src/PostgREST/Workers.hs
index c282bb5412..3aa19c06d5 100644
--- a/src/PostgREST/Workers.hs
+++ b/src/PostgREST/Workers.hs
@@ -144,6 +144,7 @@ connectionStatus appState =
         "Attempting to reconnect to the database in "
         <> (show delay::Text)
         <> " seconds..."
+      when itShould $ AppState.putRetryNextIn appState delay
       return itShould
 
 -- | Load the DbStructure by using a connection from the pool.