From c8739aa8a3349f4b1056e4599cee45285deb84e0 Mon Sep 17 00:00:00 2001
From: Robert Vollmert <rob@vllmrt.net>
Date: Mon, 1 Aug 2022 13:03:39 +0200
Subject: [PATCH] fix: fix some race conditions running connection worker

Previously, it was quite possible to have two connection workers
running, or to get into a state where a failed connection worker
is still considered running preventing new connection workers
from starting.
---
 CHANGELOG.md              |  3 +++
 src/PostgREST/AppState.hs | 18 +++++++-----------
 src/PostgREST/Workers.hs  | 12 ++++++------
 3 files changed, 16 insertions(+), 17 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0b8947ba6..87d94ab5c 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -41,6 +41,9 @@ This project adheres to [Semantic Versioning](http://semver.org/).
    + Can generate the plan for different media types using the `for` parameter: `Accept: application/vnd.pgrst.plan; for="application/vnd.pgrst.object"`
    + Different options for the plan can be used with the `options` parameter: `Accept: application/vnd.pgrst.plan; options=analyze|verbose|settings|buffers|wal`
    + The plan can be obtained in text or json by using different media type suffixes: `Accept: application/vnd.pgrst.plan+text` and `Accept: application/vnd.pgrst.plan+json`.
+   + Limited to generating the plan of a json representation(`application/json`) but can be extended later to allow other representations.
+   + The plan can be obtained in text(`Accept: application/vnd.pgrst.plan+text`) and json(`Accept: application/vnd.pgrst.plan+json` or `Accept: application/vnd.pgrst.plan`) format.
+ - #2397, Fix race conditions managing database connection helper - @robx
 
 ### Fixed
 
diff --git a/src/PostgREST/AppState.hs b/src/PostgREST/AppState.hs
index 55b6e51b3..682567f14 100644
--- a/src/PostgREST/AppState.hs
+++ b/src/PostgREST/AppState.hs
@@ -5,19 +5,18 @@ module PostgREST.AppState
   , getConfig
   , getDbStructure
   , getIsListenerOn
-  , getIsWorkerOn
   , getJsonDbS
   , getMainThreadId
   , getPgVersion
-  , getTime
   , getRetryNextIn
+  , getTime
+  , getWorkerSem
   , init
   , initWithPool
   , logWithZTime
   , putConfig
   , putDbStructure
   , putIsListenerOn
-  , putIsWorkerOn
   , putJsonDbS
   , putPgVersion
   , putRetryNextIn
@@ -52,8 +51,8 @@ data AppState = AppState
   , stateDbStructure  :: IORef (Maybe DbStructure)
   -- | Cached DbStructure in json
   , stateJsonDbS      :: IORef ByteString
-  -- | Helper ref to make sure just one connectionWorker can run at a time
-  , stateIsWorkerOn   :: IORef Bool
+  -- | Binary semaphore to make sure just one connectionWorker can run at a time
+  , stateWorkerSem    :: MVar ()
   -- | Binary semaphore used to sync the listener(NOTIFY reload) with the connectionWorker.
   , stateListener     :: MVar ()
   -- | State of the LISTEN channel, used for the admin server checks
@@ -81,7 +80,7 @@ initWithPool newPool conf =
     <$> newIORef minimumPgVersion -- assume we're in a supported version when starting, this will be corrected on a later step
     <*> newIORef Nothing
     <*> newIORef mempty
-    <*> newIORef False
+    <*> newEmptyMVar
     <*> newEmptyMVar
     <*> newIORef False
     <*> newIORef conf
@@ -118,11 +117,8 @@ getJsonDbS = readIORef . stateJsonDbS
 putJsonDbS :: AppState -> ByteString -> IO ()
 putJsonDbS appState = atomicWriteIORef (stateJsonDbS appState)
 
-getIsWorkerOn :: AppState -> IO Bool
-getIsWorkerOn = readIORef . stateIsWorkerOn
-
-putIsWorkerOn :: AppState -> Bool -> IO ()
-putIsWorkerOn = atomicWriteIORef . stateIsWorkerOn
+getWorkerSem :: AppState -> MVar ()
+getWorkerSem = stateWorkerSem
 
 getRetryNextIn :: AppState -> IO Int
 getRetryNextIn = readIORef . stateRetryNextIn
diff --git a/src/PostgREST/Workers.hs b/src/PostgREST/Workers.hs
index 3097b2a25..9deae36ea 100644
--- a/src/PostgREST/Workers.hs
+++ b/src/PostgREST/Workers.hs
@@ -48,7 +48,7 @@ data SCacheStatus
 -- up-to-date schema cache(DbStructure).  This method is meant to be called
 -- multiple times by the same thread, but does nothing if the previous
 -- invocation has not terminated. In all cases this method does not halt the
--- calling thread, the work is preformed in a separate thread.
+-- calling thread, the work is performed in a separate thread.
 --
 -- Background thread that does the following :
 --  1. Tries to connect to pg server and will keep trying until success.
@@ -57,13 +57,14 @@ data SCacheStatus
 --  3. Obtains the dbStructure. If this fails, it goes back to 1.
 connectionWorker :: AppState -> IO ()
 connectionWorker appState = do
-  isWorkerOn <- AppState.getIsWorkerOn appState
+  runExclusively (AppState.getWorkerSem appState) work
   -- Prevents multiple workers to be running at the same time. Could happen on
   -- too many SIGUSR1s.
-  unless isWorkerOn $ do
-    AppState.putIsWorkerOn appState True
-    void $ forkIO work
   where
+    runExclusively mvar action = mask_ $ do
+      success <- tryPutMVar mvar ()
+      when success $ do
+        void $ forkIO $ action `finally` takeMVar mvar
     work = do
       AppConfig{..} <- AppState.getConfig appState
       AppState.logWithZTime appState "Attempting to connect to the database..."
@@ -95,7 +96,6 @@ connectionWorker appState = do
             SCFatalFail ->
               -- die if our schema cache query has an error
               killThread $ AppState.getMainThreadId appState
-          AppState.putIsWorkerOn appState False
 
 -- | Check if a connection from the pool allows access to the PostgreSQL
 -- database.  If not, the pool connections are released and a new connection is