Make transaction modifying queries non-interruptible

Author: arybczak

src/Database/PostgreSQL/PQTypes/Internal/Connection.hs

@@ -334,14 +334,16 @@ runQueryImpl
   -> IO (Int, ForeignPtr PGresult)
   -> IO (Int, ForeignPtr PGresult, ConnectionStats -> ConnectionStats)
 runQueryImpl Connection {..} sql execSql = do
-  E.mask $ \restore -> do
+  E.uninterruptibleMask $ \restore -> do
     -- While the query runs, the current thread will not be able to receive
     -- asynchronous exceptions. This prevents clients of the library from
     -- interrupting execution of the query. To remedy that we spawn a separate
     -- thread for the query execution and while we wait for its completion, we
     -- are able to receive asynchronous exceptions (assuming that threaded GHC
     -- runtime system is used) and react appropriately.
-    queryRunner <- async . restore $ do
+    queryRunner <- asyncWithUnmask $ \unmask -> unmask $ do
+      -- Uncoditionally unmask asynchronous exceptions here so that 'cancel'
+      -- operation potentially invoked below works as expected.
       t1 <- getMonotonicTime
       (paramCount, res) <- execSql
       t2 <- getMonotonicTime
@@ -370,7 +372,7 @@ runQueryImpl Connection {..} sql execSql = do
     -- for the query runner thread to terminate. It is paramount we make the
     -- exception handler uninterruptible as we can't exit from the main block
     -- until the query runner thread has terminated.
-    E.onException (restore $ wait queryRunner) . E.uninterruptibleMask_ $ do
+    E.onException (restore $ wait queryRunner) $ do
       c_PQcancel connPtr >>= \case
         -- If query cancellation request was successfully processed, there is
         -- nothing else to do apart from waiting for the runner to terminate.

src/Database/PostgreSQL/PQTypes/Internal/State.hs

@@ -47,31 +47,36 @@ initConnectionState
 initConnectionState ics = \case
   AcquireOnDemand -> pure OnDemand
   AcquireAndHold tsIsolationLevel tsPermissions -> do
-    let isolationLevel = case tsIsolationLevel of
-          DefaultLevel -> ""
-          ReadCommitted -> "ISOLATION LEVEL READ COMMITTED"
-          RepeatableRead -> "ISOLATION LEVEL REPEATABLE READ"
-          Serializable -> "ISOLATION LEVEL SERIALIZABLE"
-        permissions = case tsPermissions of
-          DefaultPermissions -> ""
-          ReadOnly -> "READ ONLY"
-          ReadWrite -> "READ WRITE"
+    let initSql =
+          smconcat
+            [ "BEGIN"
+            , case tsIsolationLevel of
+                DefaultLevel -> ""
+                ReadCommitted -> "ISOLATION LEVEL READ COMMITTED"
+                RepeatableRead -> "ISOLATION LEVEL REPEATABLE READ"
+                Serializable -> "ISOLATION LEVEL SERIALIZABLE"
+            , case tsPermissions of
+                DefaultPermissions -> ""
+                ReadOnly -> "READ ONLY"
+                ReadWrite -> "READ WRITE"
+            ]
     (conn, cdata) <- takeConnection ics
-    _ <- liftBase . runQueryIO @SQL conn $ "BEGIN" <+> isolationLevel <+> permissions
+    _ <- liftBase . uninterruptibleMask_ $ runQueryIO @SQL conn initSql
     pure $ Acquired tsIsolationLevel tsPermissions conn cdata
 
 finalizeConnectionState
-  :: (HasCallStack, MonadBase IO m, MonadMask m)
+  :: (HasCallStack, MonadBase IO m)
   => InternalConnectionSource m cdata
   -> ExitCase r
   -> ConnectionState cdata
   -> m ()
 finalizeConnectionState ics ec = \case
   OnDemand -> pure ()
   Acquired _ _ conn cdata -> do
-    _ <- liftBase . runQueryIO @SQL conn $ case ec of
-      ExitCaseSuccess _ -> "COMMIT"
-      _ -> "ROLLBACK"
+    let finalizeSql = case ec of
+          ExitCaseSuccess _ -> "COMMIT"
+          _ -> "ROLLBACK"
+    _ <- liftBase . uninterruptibleMask_ $ runQueryIO @SQL conn finalizeSql
     putConnection ics (conn, cdata) ec
   Finalized -> error "finalized connection"
 
@@ -184,7 +189,7 @@ withConnection ConnectionData {..} action = do
     Finalized -> error "finalized connection"
 
 initConnectionData
-  :: MonadBase IO m
+  :: (MonadBase IO m, MonadMask m)
   => ConnectionSourceM m
   -> ConnectionAcquisitionMode
   -> m (ConnectionData m)

src/Database/PostgreSQL/PQTypes/Transaction.hs

@@ -47,16 +47,18 @@ withSavepoint (Savepoint savepoint) m = do
 
 ----------------------------------------
 
+-- Note: below functions that modify transaction state need to not be
+-- interruptible so we don't end up in unexpected transaction state.
+
 -- | Begin transaction using given transaction settings.
-begin :: (HasCallStack, MonadDB m, MonadThrow m) => m ()
-begin = do
+begin :: (HasCallStack, MonadDB m, MonadMask m) => m ()
+begin = uninterruptibleMask_ $ do
   getConnectionAcquisitionMode >>= \case
     AcquireOnDemand -> do
       throwDB $ HPQTypesError "Can't begin a transaction in OnDemand mode"
     AcquireAndHold isolationLevel permissions -> do
       runSQL_ $
-        mintercalate
-          " "
+        smconcat
           [ "BEGIN"
           , case isolationLevel of
               DefaultLevel -> ""
@@ -70,8 +72,8 @@ begin = do
           ]
 
 -- | Commit active transaction using given transaction settings.
-commit :: (HasCallStack, MonadDB m, MonadThrow m) => m ()
-commit = do
+commit :: (HasCallStack, MonadDB m, MonadMask m) => m ()
+commit = uninterruptibleMask_ $ do
   getConnectionAcquisitionMode >>= \case
     AcquireOnDemand -> do
       throwDB $ HPQTypesError "Can't commit a transaction in OnDemand mode"
@@ -80,8 +82,8 @@ commit = do
       begin
 
 -- | Rollback active transaction using given transaction settings.
-rollback :: (HasCallStack, MonadDB m, MonadThrow m) => m ()
-rollback = do
+rollback :: (HasCallStack, MonadDB m, MonadMask m) => m ()
+rollback = uninterruptibleMask_ $ do
   getConnectionAcquisitionMode >>= \case
     AcquireOnDemand -> do
       throwDB $ HPQTypesError "Can't rollback a transaction in OnDemand mode"
@@ -96,4 +98,8 @@ unsafeWithoutTransaction
 unsafeWithoutTransaction action = do
   getConnectionAcquisitionMode >>= \case
     AcquireOnDemand -> action
-    AcquireAndHold {} -> bracket_ (runSQL_ "COMMIT") begin action
+    AcquireAndHold {} ->
+      bracket_
+        (uninterruptibleMask_ $ runSQL_ "COMMIT")
+        begin
+        action