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.