From 48967167c439a2a5adf51f4318e5be2a2f74b05e Mon Sep 17 00:00:00 2001
From: Evgeny Poberezkin <2769109+epoberezkin@users.noreply.github.com>
Date: Sun, 3 Jan 2021 11:31:06 +0000
Subject: [PATCH] add table mutexes for insert queries (to get inserted ID
 thread-safely via single connection)

---
 src/Simplex/Messaging/Agent/Env/SQLite.hs   |  9 +---
 src/Simplex/Messaging/Agent/ServerClient.hs |  2 +-
 src/Simplex/Messaging/Agent/Store/SQLite.hs | 46 +++++++++++++++++++--
 3 files changed, 44 insertions(+), 13 deletions(-)

diff --git a/src/Simplex/Messaging/Agent/Env/SQLite.hs b/src/Simplex/Messaging/Agent/Env/SQLite.hs
index ad64aa181..e8b1d4717 100644
--- a/src/Simplex/Messaging/Agent/Env/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Env/SQLite.hs
@@ -15,7 +15,6 @@ import Numeric.Natural
 import Simplex.Messaging.Agent.ServerClient
 import Simplex.Messaging.Agent.Store
 import Simplex.Messaging.Agent.Store.SQLite
-import Simplex.Messaging.Agent.Store.SQLite.Schema
 import Simplex.Messaging.Agent.Transmission
 import Simplex.Messaging.Server.Transmission (PublicKey)
 import qualified Simplex.Messaging.Server.Transmission as SMP
@@ -64,14 +63,8 @@ newAgentClient qSize = do
   commands <- newTVar M.empty
   return AgentClient {rcvQ, sndQ, respQ, servers, commands}
 
-openDB :: MonadUnliftIO m => AgentConfig -> m SQLiteStore
-openDB AgentConfig {dbFile} = liftIO $ do
-  db <- SQLiteStore <$> DB.open dbFile
-  createSchema $ conn db
-  return db
-
 newEnv :: (MonadUnliftIO m, MonadRandom m) => AgentConfig -> m Env
 newEnv config = do
   idsDrg <- drgNew >>= newTVarIO
-  db <- openDB config
+  db <- newSQLiteStore $ dbFile config
   return Env {config, idsDrg, db}
diff --git a/src/Simplex/Messaging/Agent/ServerClient.hs b/src/Simplex/Messaging/Agent/ServerClient.hs
index bf358a5a2..44d33585d 100644
--- a/src/Simplex/Messaging/Agent/ServerClient.hs
+++ b/src/Simplex/Messaging/Agent/ServerClient.hs
@@ -35,7 +35,7 @@ newServerClient ::
 newServerClient cfg smpRcvQ host port = do
   smpSndQ <- atomically . newTBQueue $ tbqSize cfg
   let c = ServerClient {smpSndQ, smpRcvQ}
-  _srvA <- async $ runTCPClient host p (client c)
+  _srvA <- async $ runTCPClient host port (client c)
   -- TODO because exception can be thrown inside async it is not caught by newSMPServer
   -- there possibly needs to be another channel to communicate with ServerClient if it fails
   -- alternatively, there may be just timeout on sent commands -
diff --git a/src/Simplex/Messaging/Agent/Store/SQLite.hs b/src/Simplex/Messaging/Agent/Store/SQLite.hs
index e4f7b261a..78c21bd75 100644
--- a/src/Simplex/Messaging/Agent/Store/SQLite.hs
+++ b/src/Simplex/Messaging/Agent/Store/SQLite.hs
@@ -12,7 +12,10 @@ import Database.SQLite.Simple (NamedParam (..))
 import qualified Database.SQLite.Simple as DB
 import Multiline (s)
 import Simplex.Messaging.Agent.Store
+import Simplex.Messaging.Agent.Store.SQLite.Schema
 import Simplex.Messaging.Agent.Transmission
+import qualified UnliftIO.Exception as E
+import UnliftIO.STM
 
 addServerQuery :: DB.Query
 addServerQuery =
@@ -25,13 +28,48 @@ addServerQuery =
       key_hash=excluded.key_hash;
   |]
 
-newtype SQLiteStore = SQLiteStore {conn :: DB.Connection}
+data SQLiteStore = SQLiteStore
+  { conn :: DB.Connection,
+    serversLock :: TMVar (),
+    recipientQueuesLock :: TMVar (),
+    senderQueuesLock :: TMVar (),
+    connectionsLock :: TMVar (),
+    messagesLock :: TMVar ()
+  }
+
+newSQLiteStore :: MonadUnliftIO m => String -> m SQLiteStore
+newSQLiteStore dbFile = do
+  conn <- liftIO $ DB.open dbFile
+  liftIO $ createSchema conn
+  serversLock <- newTMVarIO ()
+  recipientQueuesLock <- newTMVarIO ()
+  senderQueuesLock <- newTMVarIO ()
+  connectionsLock <- newTMVarIO ()
+  messagesLock <- newTMVarIO ()
+  return
+    SQLiteStore
+      { conn,
+        serversLock,
+        recipientQueuesLock,
+        senderQueuesLock,
+        connectionsLock,
+        messagesLock
+      }
+
+withLock :: MonadUnliftIO m => SQLiteStore -> (SQLiteStore -> TMVar ()) -> (DB.Connection -> m a) -> m a
+withLock store tableLock query = do
+  let lock = tableLock store
+  E.bracket_
+    (atomically $ takeTMVar lock)
+    (atomically $ putTMVar lock ())
+    (query $ conn store)
 
 instance MonadUnliftIO m => MonadAgentStore SQLiteStore m where
   addServer :: SQLiteStore -> SMPServer -> m (Either StoreError SMPServerId)
-  addServer store SMPServer {host, port, keyHash} = liftIO $ do
-    DB.executeNamed (conn store) addServerQuery [":host_address" := host, ":port" := port, ":key_hash" := keyHash]
-    Right <$> DB.lastInsertRowId (conn store)
+  addServer store SMPServer {host, port, keyHash} =
+    withLock store serversLock $ \c -> liftIO $ do
+      DB.executeNamed c addServerQuery [":host_address" := host, ":port" := port, ":key_hash" := keyHash]
+      Right <$> DB.lastInsertRowId c
 
 --   createRcvConn :: DB.Connection -> Maybe ConnAlias -> ReceiveQueue -> m (Either StoreError (Connection CReceive))
 --   createRcvConn conn connAlias q = do