Add connection retrying on startup and SIGHUP, Fix #742

CHANGELOG.md

@@ -7,6 +7,8 @@ This project adheres to [Semantic Versioning](http://semver.org/).
 
 ### Added
 
+- #742, Add connection retrying on startup and SIGHUP - @steve-chavez
+
 ### Fixed
 
 ## [0.4.1.0] - 2017-04-25

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,59 @@ 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) -> IO ()
+connectionWorker mainTid pool schema refDbStructure = 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 _ -> 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 +121,21 @@ 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
+  mainTid <- myThreadId
 
-  refDbStructure <- newIORef $ either (panic . show) id result
+  connectionWorker mainTid pool (configSchema conf) refDbStructure
 
 #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
    ) Nothing
 #endif
 

postgrest.cabal

@@ -40,6 +40,7 @@ executable postgrest
                      , warp
                      , bytestring
                      , base64-bytestring
+                     , retry
   if !os(windows)
     build-depends:     unix
 

src/PostgREST/App.hs

@@ -45,6 +45,7 @@ import           PostgREST.DbRequestBuilder( readRequest
 import           PostgREST.Error           ( simpleError, pgError
                                            , apiRequestError
                                            , singularityError, binaryFieldError
+                                           , connectionLostError
                                            )
 import           PostgREST.RangeQuery      (allRange, rangeOffset)
 import           PostgREST.Middleware
@@ -62,28 +63,30 @@ import           Data.Function (id)
 import           Protolude              hiding (intercalate, Proxy)
 import           Safe                   (headMay)
 
-postgrest :: AppConfig -> IORef DbStructure -> P.Pool -> IO POSIXTime ->
+postgrest :: AppConfig -> IORef (Maybe DbStructure) -> P.Pool -> IO POSIXTime ->
              Application
 postgrest conf refDbStructure pool getTime =
   let middle = (if configQuiet conf then id else logStdout) . defaultMiddle in
 
   middle $ \ req respond -> do
     time <- getTime
     body <- strictRequestBody req
-    dbStructure <- readIORef refDbStructure
-
-    response <- case userApiRequest (configSchema conf) req body of
-      Left err -> return $ apiRequestError err
-      Right apiRequest -> do
-        let jwtSecret = binarySecret <$> configJwtSecret conf
-            eClaims = jwtClaims jwtSecret (iJWT apiRequest) time
-            authed = containsRole eClaims
-            handleReq = runWithClaims conf eClaims (app dbStructure conf) apiRequest
-            txMode = transactionMode dbStructure
-              (iTarget apiRequest) (iAction apiRequest)
-        response <- P.use pool $ HT.transaction HT.ReadCommitted txMode handleReq
-        return $ either (pgError authed) identity response
-    respond response
+    maybeDbStructure <- readIORef refDbStructure
+    case maybeDbStructure of
+      Nothing -> respond connectionLostError
+      Just dbStructure -> do
+        response <- case userApiRequest (configSchema conf) req body of
+          Left err -> return $ apiRequestError err
+          Right apiRequest -> do
+            let jwtSecret = binarySecret <$> configJwtSecret conf
+                eClaims = jwtClaims jwtSecret (iJWT apiRequest) time
+                authed = containsRole eClaims
+                handleReq = runWithClaims conf eClaims (app dbStructure conf) apiRequest
+                txMode = transactionMode dbStructure
+                  (iTarget apiRequest) (iAction apiRequest)
+            response <- P.use pool $ HT.transaction HT.ReadCommitted txMode handleReq
+            return $ either (pgError authed) identity response
+        respond response
 
 transactionMode :: DbStructure -> Target -> Action -> H.Mode
 transactionMode structure target action =

src/PostgREST/Error.hs

@@ -8,6 +8,7 @@ module PostgREST.Error (
 , simpleError
 , singularityError
 , binaryFieldError
+, connectionLostError
 , encodeError
 ) where
 
@@ -73,6 +74,10 @@ binaryFieldError =
   simpleError HT.status406 (toS (toMime CTOctetStream) <>
   " requested but a single column was not selected")
 
+connectionLostError :: Response
+connectionLostError =
+  simpleError HT.status503 "Database connection lost, retrying the connection."
+
 encodeError :: JSON.ToJSON a => a -> LByteString
 encodeError = JSON.encode
 

test/Main.hs

@@ -41,7 +41,7 @@ main = do
 
 
   result <- P.use pool $ getDbStructure "test"
-  refDbStructure <- newIORef $ either (panic.show) id result
+  refDbStructure <- newIORef $ Just $ either (panic.show) id result
   let withApp      = return $ postgrest (testCfg testDbConn)          refDbStructure pool getTime
       ltdApp       = return $ postgrest (testLtdRowsCfg testDbConn)   refDbStructure pool getTime
       unicodeApp   = return $ postgrest (testUnicodeCfg testDbConn)   refDbStructure pool getTime