backend-db: Actually catch serialization errors to retry them

backend-db/Rhyolite/Backend/DB/PsqlSimple.hs

@@ -17,7 +17,9 @@ module Rhyolite.Backend.DB.PsqlSimple
   , Binary (..), (:.)(..), PGArray (..)
   , ToRow (..), FromRow (..)
   , ToField (..), FromField (..)
-  , Query (..), sql, traceQuery, traceExecute, traceExecute_
+  , Query (..)
+  , WrappedSqlError (..)
+  , sql, traceQuery, traceExecute, traceExecute_
   , liftWithConn
   , queryQ, executeQ, executeQ_, sqlQ, traceQueryQ, traceExecuteQ, traceExecuteQ_
   , fromIdRow
@@ -63,7 +65,6 @@ data WrappedSqlError = WrappedSqlError
   , _wrappedSqlError_error :: SqlError
   }
   deriving Show
-
 instance Exception WrappedSqlError
 
 rethrowWithQuery :: ToRow q => Connection -> Query -> q -> SqlError -> IO a
@@ -128,7 +129,6 @@ class PostgresRaw m where
   default returning :: (m ~ t n, ToRow q, FromRow r, PostgresRaw n, Monad n, MonadTrans t) => Query -> [q] -> m [r]
   returning psql qs = lift $ returning psql qs
 
-
 traceQuery :: (PostgresRaw m, MonadIO m, ToRow q, FromRow r) => Query -> q -> m [r]
 traceQuery p q = do
   s <- formatQuery p q

backend-db/Rhyolite/Backend/DB/Serializable.hs

@@ -10,30 +10,34 @@
 
 module Rhyolite.Backend.DB.Serializable
   ( Serializable
+  , SqlSerializationError (..)
   , runSerializable
   , toDbPersist
   , unsafeLiftDbPersist
   , unsafeMkSerializable
   , unSerializable
+  , withSqlSerializationErrorWrapping
   ) where
 
+import qualified Control.Exception as E
 import Control.Monad.Base (MonadBase (liftBase))
 import Control.Monad.Catch (MonadThrow)
+import qualified Control.Monad.Catch as MonadCatch
 import Control.Monad.IO.Class (MonadIO (liftIO))
 import Control.Monad.Reader (ReaderT, runReaderT, withReaderT)
 import Control.Monad.Logger (MonadLogger, LoggingT)
+import qualified Control.Monad.State as S
 import Data.Coerce (coerce)
 import qualified Database.Groundhog.Generic.Migration as Mig
 import Database.Groundhog.Postgresql (Postgresql (..))
 import qualified Database.PostgreSQL.Simple as Pg
 import qualified Database.PostgreSQL.Simple.Transaction as Pg
 import Data.Pool (Pool, withResource)
 import qualified Database.Groundhog.Core as Hog
+
 import qualified Rhyolite.Backend.DB.PsqlSimple as PsqlSimple
 import Rhyolite.Logging (LoggingEnv, runLoggingEnv)
 
-import qualified Control.Monad.State as S
-
 -- | A monad for database transactions with serializable isolation level.
 --
 -- Because this monad may retry execution of code automatically, it does not lawfully lift any effects other
@@ -101,6 +105,20 @@ instance Mig.SchemaAnalyzer Serializable where
   getMigrationPack i = coerce <$> unsafeLiftDbPersist (Mig.getMigrationPack i)
 
 
+data SqlSerializationError = SqlSerializationError deriving (Eq, Ord, Show)
+instance E.Exception SqlSerializationError
+
+withSqlSerializationErrorWrapping :: forall m a. (MonadCatch.MonadCatch m, MonadThrow m) => m a -> m a
+withSqlSerializationErrorWrapping = flip MonadCatch.catches
+  [ MonadCatch.Handler $ \(e :: Pg.SqlError) -> convert id e
+  , MonadCatch.Handler $ \(e :: PsqlSimple.WrappedSqlError) -> convert PsqlSimple._wrappedSqlError_error e
+  ]
+  where
+    convert :: E.Exception e => (e -> Pg.SqlError) -> e -> m a
+    convert toSqlError e = if Pg.isSerializationError (toSqlError e)
+      then MonadCatch.throwM SqlSerializationError
+      else MonadCatch.throwM e
+
 unsafeMkSerializable :: ReaderT Pg.Connection (LoggingT IO) a -> Serializable a
 unsafeMkSerializable = Serializable
 
@@ -111,9 +129,37 @@ toDbPersist :: forall a. Serializable a -> Hog.DbPersist Postgresql (LoggingT IO
 toDbPersist (Serializable act) = Hog.DbPersist $ withReaderT coerce act
 
 unsafeLiftDbPersist :: forall a. Hog.DbPersist Postgresql (LoggingT IO) a -> Serializable a
-unsafeLiftDbPersist (Hog.DbPersist act) = Serializable $ withReaderT coerce act
+unsafeLiftDbPersist (Hog.DbPersist act) = Serializable $ withSqlSerializationErrorWrapping $ withReaderT coerce act
 
 runSerializable :: forall a m. (MonadIO m) => Pool Pg.Connection -> LoggingEnv -> Serializable a -> m a
 runSerializable pool logger (Serializable act) = liftIO $ withResource pool $ \c ->
-  Pg.withTransactionSerializable c $
-    runLoggingEnv logger $ runReaderT act c
+  withTransactionModeRetry'
+    (Pg.TransactionMode{ Pg.isolationLevel = Pg.Serializable, Pg.readWriteMode = Pg.ReadWrite})
+    (\(_ :: SqlSerializationError) -> True)
+    c
+    (runLoggingEnv logger $ runReaderT act c)
+
+
+-- | Like 'Pg.withTransactionModeRetry' but polymorphic over exception type.
+-- Copied from https://github.com/phadej/postgresql-simple/blob/e02684f9c38acf736ac590b36b919000a2b45bc4/src/Database/PostgreSQL/Simple/Transaction.hs#L156-L174
+withTransactionModeRetry' :: forall e a. E.Exception e => Pg.TransactionMode -> (e -> Bool) -> Pg.Connection -> IO a -> IO a
+withTransactionModeRetry' mode shouldRetry conn act =
+  E.mask $ \restore ->
+    retryLoop $ E.try $ do
+      a <- restore act `E.onException` rollback_ conn
+      Pg.commit conn
+      return a
+  where
+    retryLoop :: IO (Either e a) -> IO a
+    retryLoop act' = do
+      Pg.beginMode mode conn
+      r <- act'
+      case r of
+        Left e -> case shouldRetry e of
+          True -> retryLoop act'
+          False -> E.throwIO e
+        Right a -> return a
+
+    -- | Rollback a transaction, ignoring any @IOErrors@
+    rollback_ :: Pg.Connection -> IO ()
+    rollback_ c = Pg.rollback c `E.catch` \(_ :: IOError) -> return ()

default.nix

@@ -29,9 +29,9 @@ let
     aeson-gadt-th = repos.aeson-gadt-th;
     bytestring-trie = repos.bytestring-trie;
     dependent-monoidal-map = repos.dependent-monoidal-map;
-    groundhog = repos.groundhog + /groundhog;
-    groundhog-postgresql = repos.groundhog + /groundhog-postgresql;
-    groundhog-th = repos.groundhog + /groundhog-th;
+    groundhog = repos.groundhog + "/groundhog";
+    groundhog-postgresql = repos.groundhog + "/groundhog-postgresql";
+    groundhog-th = repos.groundhog + "/groundhog-th";
     HaskellNet = repos.HaskellNet; # (super is marked as broken) unreleased fixes for newer GHC
     HaskellNet-SSL = repos.HaskellNet-SSL; # (super is marked as broken)
     postgresql-simple = repos.postgresql-simple;  # v0.5.4.0 with a fix

dep/groundhog/default.nix

@@ -1,7 +1,2 @@
 # DO NOT HAND-EDIT THIS FILE
-import ((import <nixpkgs> {}).fetchFromGitHub (
-  let json = builtins.fromJSON (builtins.readFile ./github.json);
-  in { inherit (json) owner repo rev sha256;
-       private = json.private or false;
-     }
-))
+import (import ./thunk.nix)

dep/groundhog/github.json

@@ -2,6 +2,7 @@
   "owner": "obsidiansystems",
   "repo": "groundhog",
   "branch": "develop",
-  "rev": "e08021aef70f17196214e795789d0993177231b5",
-  "sha256": "13ikpayvzv85wmfgh7wnf73ayhjy77384mnwvqcxi0nbn5m7kk56"
+  "private": false,
+  "rev": "fb16055662e81bed258bd246380bc67bb6f7afab",
+  "sha256": "0071nvdw296bz5apvxg7r6inskg4znayqp050pjbnbl72b6aq8n8"
 }

dep/groundhog/thunk.nix

@@ -0,0 +1,9 @@
+# DO NOT HAND-EDIT THIS FILE
+let fetch = { private ? false, fetchSubmodules ? false, owner, repo, rev, sha256, ... }:
+  if !fetchSubmodules && !private then builtins.fetchTarball {
+    url = "https://github.com/${owner}/${repo}/archive/${rev}.tar.gz"; inherit sha256;
+  } else (import <nixpkgs> {}).fetchFromGitHub {
+    inherit owner repo rev sha256 fetchSubmodules private;
+  };
+  json = builtins.fromJSON (builtins.readFile ./github.json);
+in fetch json