Add Retry-After hint

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

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

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
 

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.