From 92d00749a35fddd72000d383d7bb6a582f87d1ed Mon Sep 17 00:00:00 2001
From: Robert Vollmert <rob@vllmrt.net>
Date: Tue, 2 Aug 2022 15:44:55 +0200
Subject: [PATCH] refactor: move pool release on shutdown out of signal handler

This changes behaviour somewhat in that:
- We now consistently release the pool on shutdown, even on non-Unix
  platforms, and including for CmdDumpConfig.
- We release the pool *after* interrupting `App.run`, which will
  rather cause more than fewer connection to be closed properly.
  (Previously any in-use connections would not have been caught by
  `releasePool`, though *maybe* the `UserInterrupt` handling in
  the web handler ends up closing the connections properly already
  anyway).

(The main aim of the change is to make it clearer when and why the
pool is released.)
---
 src/PostgREST/AppState.hs |  4 ++++
 src/PostgREST/CLI.hs      | 21 +++++++++++++--------
 src/PostgREST/Unix.hs     |  6 +-----
 3 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/src/PostgREST/AppState.hs b/src/PostgREST/AppState.hs
index 682567f14..72ad8ad3c 100644
--- a/src/PostgREST/AppState.hs
+++ b/src/PostgREST/AppState.hs
@@ -2,6 +2,7 @@
 
 module PostgREST.AppState
   ( AppState
+  , destroy
   , getConfig
   , getDbStructure
   , getIsListenerOn
@@ -89,6 +90,9 @@ initWithPool newPool conf =
     <*> myThreadId
     <*> newIORef 0
 
+destroy :: AppState -> IO ()
+destroy = releasePool
+
 initPool :: AppConfig -> IO SQL.Pool
 initPool AppConfig{..} =
   SQL.acquire (configDbPoolSize, configDbPoolTimeout, toUtf8 configDbUri)
diff --git a/src/PostgREST/CLI.hs b/src/PostgREST/CLI.hs
index cd21c349e..b9c912de0 100644
--- a/src/PostgREST/CLI.hs
+++ b/src/PostgREST/CLI.hs
@@ -34,13 +34,19 @@ main :: App.SignalHandlerInstaller -> Maybe App.SocketRunner -> CLI -> IO ()
 main installSignalHandlers runAppWithSocket CLI{cliCommand, cliPath} = do
   conf@AppConfig{..} <-
     either panic identity <$> Config.readAppConfig mempty cliPath Nothing
-  appState <- AppState.init conf
-  case cliCommand of
-    CmdDumpConfig -> do
-      when configDbConfig $ reReadConfig True appState
-      putStr . Config.toText =<< AppState.getConfig appState
-    CmdDumpSchema -> putStrLn =<< dumpSchema appState
-    CmdRun -> App.run installSignalHandlers runAppWithSocket appState
+
+  -- Per https://github.com/PostgREST/postgrest/issues/268, we want to
+  -- explicitly close the connections to PostgreSQL on shutdown.
+  -- 'AppState.destroy' takes care of that.
+  bracket
+    (AppState.init conf)
+    AppState.destroy
+    (\appState -> case cliCommand of
+      CmdDumpConfig -> do
+        when configDbConfig $ reReadConfig True appState
+        putStr . Config.toText =<< AppState.getConfig appState
+      CmdDumpSchema -> putStrLn =<< dumpSchema appState
+      CmdRun -> App.run installSignalHandlers runAppWithSocket appState)
 
 -- | Dump DbStructure schema to JSON
 dumpSchema :: AppState -> IO LBS.ByteString
@@ -54,7 +60,6 @@ dumpSchema appState = do
           (toList configDbSchemas)
           configDbExtraSearchPath
           configDbPreparedStatements
-  AppState.releasePool appState
   case result of
     Left e -> do
       hPutStrLn stderr $ "An error ocurred when loading the schema cache:\n" <> show e
diff --git a/src/PostgREST/Unix.hs b/src/PostgREST/Unix.hs
index acc1256ff..c6926121c 100644
--- a/src/PostgREST/Unix.hs
+++ b/src/PostgREST/Unix.hs
@@ -43,11 +43,7 @@ runAppWithSocket settings app socketFileMode socketFilePath =
 -- | Set signal handlers, only for systems with signals
 installSignalHandlers :: AppState.AppState -> IO ()
 installSignalHandlers appState = do
-  -- Releases the connection pool whenever the program is terminated,
-  -- see https://github.com/PostgREST/postgrest/issues/268
-  let interrupt = do
-        AppState.releasePool appState
-        throwTo (AppState.getMainThreadId appState) UserInterrupt
+  let interrupt = throwTo (AppState.getMainThreadId appState) UserInterrupt
   install Signals.sigINT interrupt
   install Signals.sigTERM interrupt