agent: use MVar for DB connection for more fair connection distribution between threads (#1147)

src/Simplex/Messaging/Agent/Store/SQLite.hs

@@ -287,6 +287,7 @@ import System.FilePath (takeDirectory)
 import System.IO (hFlush, stdout)
 import UnliftIO.Exception (bracketOnError, onException)
 import qualified UnliftIO.Exception as E
+import UnliftIO.MVar
 import UnliftIO.STM
 
 -- * SQLite Store implementation
@@ -382,8 +383,8 @@ connectSQLiteStore :: FilePath -> ScrubbedBytes -> Bool -> IO SQLiteStore
 connectSQLiteStore dbFilePath key keepKey = do
   dbNew <- not <$> doesFileExist dbFilePath
   dbConn <- dbBusyLoop (connectDB dbFilePath key)
+  dbConnection <- newMVar dbConn
   atomically $ do
-    dbConnection <- newTMVar dbConn
     dbKey <- newTVar $! storeKey key keepKey
     dbClosed <- newTVar False
     pure SQLiteStore {dbFilePath, dbKey, dbConnection, dbNew, dbClosed}
@@ -421,14 +422,14 @@ openSQLiteStore st@SQLiteStore {dbClosed} key keepKey =
 openSQLiteStore_ :: SQLiteStore -> ScrubbedBytes -> Bool -> IO ()
 openSQLiteStore_ SQLiteStore {dbConnection, dbFilePath, dbKey, dbClosed} key keepKey =
   bracketOnError
-    (atomically $ takeTMVar dbConnection)
-    (atomically . tryPutTMVar dbConnection)
+    (takeMVar dbConnection)
+    (tryPutMVar dbConnection)
     $ \DB.Connection {slow} -> do
       DB.Connection {conn} <- connectDB dbFilePath key
       atomically $ do
-        putTMVar dbConnection DB.Connection {conn, slow}
         writeTVar dbClosed False
         writeTVar dbKey $! storeKey key keepKey
+      putMVar dbConnection DB.Connection {conn, slow}
 
 reopenSQLiteStore :: SQLiteStore -> IO ()
 reopenSQLiteStore st@SQLiteStore {dbKey, dbClosed} =

src/Simplex/Messaging/Agent/Store/SQLite/Common.hs

@@ -22,8 +22,8 @@ import Database.SQLite.Simple (SQLError)
 import qualified Database.SQLite.Simple as SQL
 import qualified Simplex.Messaging.Agent.Store.SQLite.DB as DB
 import Simplex.Messaging.Util (diffToMilliseconds)
-import UnliftIO.Exception (bracket)
 import qualified UnliftIO.Exception as E
+import UnliftIO.MVar
 import UnliftIO.STM
 
 storeKey :: ScrubbedBytes -> Bool -> Maybe ScrubbedBytes
@@ -32,16 +32,13 @@ storeKey key keepKey = if keepKey || BA.null key then Just key else Nothing
 data SQLiteStore = SQLiteStore
   { dbFilePath :: FilePath,
     dbKey :: TVar (Maybe ScrubbedBytes),
-    dbConnection :: TMVar DB.Connection,
+    dbConnection :: MVar DB.Connection,
     dbClosed :: TVar Bool,
     dbNew :: Bool
   }
 
 withConnection :: SQLiteStore -> (DB.Connection -> IO a) -> IO a
-withConnection SQLiteStore {dbConnection} =
-  bracket
-    (atomically $ takeTMVar dbConnection)
-    (atomically . putTMVar dbConnection)
+withConnection SQLiteStore {dbConnection} = withMVar dbConnection
 
 withConnection' :: SQLiteStore -> (SQL.Connection -> IO a) -> IO a
 withConnection' st action = withConnection st $ action . DB.conn
@@ -71,9 +68,9 @@ dbBusyLoop action = loop 500 3000000
     loop :: Int -> Int -> IO a
     loop t tLim =
       action `E.catch` \(e :: SQLError) ->
-        let se = SQL.sqlError e in
-        if tLim > t && (se == SQL.ErrorBusy || se == SQL.ErrorLocked)
-          then do
-            threadDelay t
-            loop (t * 9 `div` 8) (tLim - t)
-          else E.throwIO e
+        let se = SQL.sqlError e
+         in if tLim > t && (se == SQL.ErrorBusy || se == SQL.ErrorLocked)
+              then do
+                threadDelay t
+                loop (t * 9 `div` 8) (tLim - t)
+              else E.throwIO e