Add connection retrying on startup and SIGHUP, Fix #742 (#869)

* Add connection retrying on startup and SIGHUP, Fix #742
* Ensure that only one connection worker can run at a time
* Change ConnectionError status code to 503 and add automatic connection retrying

main/Main.hs

@@ -12,14 +12,15 @@ import           PostgREST.Config                     (AppConfig (..),
 import           PostgREST.Error                      (encodeError)
 import           PostgREST.OpenAPI                    (isMalformedProxyUri)
 import           PostgREST.DbStructure
+import           PostgREST.Types                      (DbStructure, Schema)
 
 import           Control.AutoUpdate
+import           Control.Retry
 import           Data.ByteString.Base64               (decode)
 import           Data.String                          (IsString (..))
 import           Data.Text                            (stripPrefix, pack, replace)
 import           Data.Text.Encoding                   (encodeUtf8, decodeUtf8)
 import           Data.Text.IO                         (hPutStrLn, readFile)
-import           Data.Function                        (id)
 import           Data.Time.Clock.POSIX                (getPOSIXTime)
 import qualified Hasql.Query                          as H
 import qualified Hasql.Session                        as H
@@ -43,6 +44,65 @@ isServerVersionSupported = do
     H.statement "SELECT current_setting('server_version_num')::integer"
       HE.unit (HD.singleRow $ HD.value HD.int4) False
 
+{-|
+  Background thread that does the following :
+  1. Tries to connect to pg server and will keep trying until success.
+  2. Checks if the pg version is supported and if it's not it kills the main program.
+  3. Obtains the dbStructure.
+  4. If 2 or 3 fail to give their result it means the connection is down so it goes back to 1,
+     otherwise it finishes his work successfully.
+-}
+connectionWorker :: ThreadId -> P.Pool -> Schema -> IORef (Maybe DbStructure) -> IORef Bool -> IO ()
+connectionWorker mainTid pool schema refDbStructure refIsWorkerOn = do
+  isWorkerOn <- readIORef refIsWorkerOn
+  unless isWorkerOn $ do
+    atomicWriteIORef refIsWorkerOn True
+    void $ forkIO work
+  where
+    work = do
+      atomicWriteIORef refDbStructure Nothing
+      putStrLn ("Attempting to connect to the database..." :: Text)
+      connected <- connectingSucceeded pool
+      when connected $ do
+        result <- P.use pool $ do
+          supported <- isServerVersionSupported
+          unless supported $ liftIO $ do
+            hPutStrLn stderr
+              ("Cannot run in this PostgreSQL version, PostgREST needs at least "
+              <> pgvName minimumPgVersion)
+            killThread mainTid
+          dbStructure <- getDbStructure schema
+          liftIO $ atomicWriteIORef refDbStructure $ Just dbStructure
+        case result of
+          Left e -> do
+            putStrLn ("Failed to query the database. Retrying." :: Text)
+            hPutStrLn stderr (toS $ encodeError e)
+            work
+          Right _ -> do
+            atomicWriteIORef refIsWorkerOn False
+            putStrLn ("Connection successful" :: Text)
+
+-- | Connect to pg server if it fails retry with capped exponential backoff until success
+connectingSucceeded :: P.Pool -> IO Bool
+connectingSucceeded pool =
+  retrying (capDelay 32000000 $ exponentialBackoff 1000000)
+           shouldRetry
+           (const $ P.release pool >> isConnectionSuccessful)
+  where
+    isConnectionSuccessful :: IO Bool
+    isConnectionSuccessful = do
+      testConn <- P.use pool $ H.sql "SELECT 1"
+      case testConn of
+        Left e -> hPutStrLn stderr (toS $ encodeError e) >> pure False
+        _ -> pure True
+    shouldRetry :: RetryStatus -> Bool -> IO Bool
+    shouldRetry rs isConnSucc = do
+      delay <- pure $ fromMaybe 0 (rsPreviousDelay rs) `div` 1000000
+      itShould <- pure $ not isConnSucc
+      when itShould $
+        putStrLn $ "Attempting to reconnect to the database in " <> (show delay::Text) <> " seconds..."
+      return itShould
+
 main :: IO ()
 main = do
   hSetBuffering stdout LineBuffering
@@ -67,31 +127,24 @@ main = do
 
   pool <- P.acquire (configPool conf, 10, pgSettings)
 
-  result <- P.use pool $ do
-    supported <- isServerVersionSupported
-    unless supported $ panic (
-      "Cannot run in this PostgreSQL version, PostgREST needs at least "
-      <> pgvName minimumPgVersion)
-    getDbStructure (toS $ configSchema conf)
+  refDbStructure <- newIORef Nothing
 
-  forM_ (lefts [result]) $ \e -> do
-    hPutStrLn stderr (toS $ encodeError e)
-    exitFailure
+  -- Helper ref to make sure just one connectionWorker can run at a time
+  refIsWorkerOn <- newIORef False
 
-  refDbStructure <- newIORef $ either (panic . show) id result
+  mainTid <- myThreadId
+
+  connectionWorker mainTid pool (configSchema conf) refDbStructure refIsWorkerOn
 
 #ifndef mingw32_HOST_OS
-  tid <- myThreadId
   forM_ [sigINT, sigTERM] $ \sig ->
     void $ installHandler sig (Catch $ do
         P.release pool
-        throwTo tid UserInterrupt
+        throwTo mainTid UserInterrupt
       ) Nothing
 
   void $ installHandler sigHUP (
-      Catch . void . P.use pool $ do
-        s <- getDbStructure (toS $ configSchema conf)
-        liftIO $ atomicWriteIORef refDbStructure s
+      Catch $ connectionWorker mainTid pool (configSchema conf) refDbStructure refIsWorkerOn
    ) Nothing
 #endif
 
@@ -100,6 +153,7 @@ main = do
     defaultUpdateSettings { updateAction = getPOSIXTime }
 
   runSettings appSettings $ postgrest conf refDbStructure pool getTime
+    (connectionWorker mainTid pool (configSchema conf) refDbStructure refIsWorkerOn)
 
 loadSecretFile :: AppConfig -> IO AppConfig
 loadSecretFile conf = extractAndTransform mSecret