agent: track queries (#1439)

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

@@ -69,32 +69,33 @@ data DBOpts = DBOpts
   { dbFilePath :: FilePath,
     dbKey :: ScrubbedBytes,
     keepKey :: Bool,
-    vacuum :: Bool
+    vacuum :: Bool,
+    track :: DB.TrackQueries
   }
 
 createDBStore :: DBOpts -> [Migration] -> MigrationConfirmation -> IO (Either MigrationError DBStore)
-createDBStore DBOpts {dbFilePath, dbKey, keepKey, vacuum} migrations confirmMigrations = do
+createDBStore DBOpts {dbFilePath, dbKey, keepKey, track, vacuum} migrations confirmMigrations = do
   let dbDir = takeDirectory dbFilePath
   createDirectoryIfMissing True dbDir
-  st <- connectSQLiteStore dbFilePath dbKey keepKey
+  st <- connectSQLiteStore dbFilePath dbKey keepKey track
   r <- migrateSchema st migrations confirmMigrations vacuum `onException` closeDBStore st
   case r of
     Right () -> pure $ Right st
     Left e -> closeDBStore st $> Left e
 
-connectSQLiteStore :: FilePath -> ScrubbedBytes -> Bool -> IO DBStore
-connectSQLiteStore dbFilePath key keepKey = do
+connectSQLiteStore :: FilePath -> ScrubbedBytes -> Bool -> DB.TrackQueries -> IO DBStore
+connectSQLiteStore dbFilePath key keepKey track = do
   dbNew <- not <$> doesFileExist dbFilePath
-  dbConn <- dbBusyLoop (connectDB dbFilePath key)
+  dbConn <- dbBusyLoop (connectDB dbFilePath key track)
   dbConnection <- newMVar dbConn
   dbKey <- newTVarIO $! storeKey key keepKey
   dbClosed <- newTVarIO False
   dbSem <- newTVarIO 0
   pure DBStore {dbFilePath, dbKey, dbSem, dbConnection, dbNew, dbClosed}
 
-connectDB :: FilePath -> ScrubbedBytes -> IO DB.Connection
-connectDB path key = do
-  db <- DB.open path
+connectDB :: FilePath -> ScrubbedBytes -> DB.TrackQueries -> IO DB.Connection
+connectDB path key track = do
+  db <- DB.open path track
   prepare db `onException` DB.close db
   -- _printPragmas db path
   pure db
@@ -127,12 +128,12 @@ openSQLiteStore_ DBStore {dbConnection, dbFilePath, dbKey, dbClosed} key keepKey
   bracketOnError
     (takeMVar dbConnection)
     (tryPutMVar dbConnection)
-    $ \DB.Connection {slow} -> do
-      DB.Connection {conn} <- connectDB dbFilePath key
+    $ \DB.Connection {slow, track} -> do
+      DB.Connection {conn} <- connectDB dbFilePath key track
       atomically $ do
         writeTVar dbClosed False
         writeTVar dbKey $! storeKey key keepKey
-      putMVar dbConnection DB.Connection {conn, slow}
+      putMVar dbConnection DB.Connection {conn, slow, track}
 
 reopenDBStore :: DBStore -> IO ()
 reopenDBStore st@DBStore {dbKey, dbClosed} =

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

@@ -11,6 +11,7 @@ module Simplex.Messaging.Agent.Store.SQLite.DB
     Binary (..),
     Connection (..),
     SlowQueryStats (..),
+    TrackQueries (..),
     open,
     close,
     execute,
@@ -38,7 +39,7 @@ import Database.SQLite.Simple.ToField (ToField (..))
 import Simplex.Messaging.Parsers (defaultJSON)
 import Simplex.Messaging.TMap (TMap)
 import qualified Simplex.Messaging.TMap as TM
-import Simplex.Messaging.Util (diffToMilliseconds, tshow)
+import Simplex.Messaging.Util (diffToMicroseconds, tshow)
 
 newtype BoolInt = BI {unBI :: Bool}
   deriving newtype (FromField, ToField)
@@ -48,9 +49,13 @@ newtype Binary = Binary {fromBinary :: ByteString}
 
 data Connection = Connection
   { conn :: SQL.Connection,
+    track :: TrackQueries,
     slow :: TMap Query SlowQueryStats
   }
 
+data TrackQueries = TQAll | TQSlow Int64 | TQOff
+  deriving (Eq)
+
 data SlowQueryStats = SlowQueryStats
   { count :: Int64,
     timeMax :: Int64,
@@ -59,22 +64,29 @@ data SlowQueryStats = SlowQueryStats
   }
   deriving (Show)
 
-timeIt :: TMap Query SlowQueryStats -> Query -> IO a -> IO a
-timeIt slow sql a = do
-  t <- getCurrentTime
-  r <-
-    a `catch` \e -> do
-      atomically $ TM.alter (Just . updateQueryErrors e) sql slow
-      throwIO e
-  t' <- getCurrentTime
-  let diff = diffToMilliseconds $ diffUTCTime t' t
-  when (diff > 1) $ atomically $ TM.alter (updateQueryStats diff) sql slow
-  pure r
+timeIt :: Connection -> Query -> IO a -> IO a
+timeIt Connection {slow, track} sql a
+  | track == TQOff = makeQuery
+  | otherwise = do
+      t <- getCurrentTime
+      r <- makeQuery
+      t' <- getCurrentTime
+      let diff = diffToMicroseconds $ diffUTCTime t' t
+      when (trackQuery diff) $ atomically $ TM.alter (updateQueryStats diff) sql slow
+      pure r
   where
+    makeQuery =
+      a `catch` \e -> do
+        atomically $ TM.alter (Just . updateQueryErrors e) sql slow
+        throwIO e
+    trackQuery diff = case track of
+      TQOff -> False
+      TQSlow t -> diff > t
+      TQAll -> True
     updateQueryErrors :: SomeException -> Maybe SlowQueryStats -> SlowQueryStats
     updateQueryErrors e Nothing = SlowQueryStats 0 0 0 $ M.singleton (tshow e) 1
-    updateQueryErrors e (Just stats@SlowQueryStats {errs}) =
-      stats {errs = M.alter (Just . maybe 1 (+ 1)) (tshow e) errs}
+    updateQueryErrors e (Just st@SlowQueryStats {errs}) =
+      st {errs = M.alter (Just . maybe 1 (+ 1)) (tshow e) errs}
     updateQueryStats :: Int64 -> Maybe SlowQueryStats -> Maybe SlowQueryStats
     updateQueryStats diff Nothing = Just $ SlowQueryStats 1 diff diff M.empty
     updateQueryStats diff (Just SlowQueryStats {count, timeMax, timeAvg, errs}) =
@@ -86,33 +98,33 @@ timeIt slow sql a = do
             errs
           }
 
-open :: String -> IO Connection
-open f = do
+open :: String -> TrackQueries -> IO Connection
+open f track = do
   conn <- SQL.open f
   slow <- TM.emptyIO
-  pure Connection {conn, slow}
+  pure Connection {conn, slow, track}
 
 close :: Connection -> IO ()
 close = SQL.close . conn
 
 execute :: ToRow q => Connection -> Query -> q -> IO ()
-execute Connection {conn, slow} sql = timeIt slow sql . SQL.execute conn sql
+execute c sql = timeIt c sql . SQL.execute (conn c) sql
 {-# INLINE execute #-}
 
 execute_ :: Connection -> Query -> IO ()
-execute_ Connection {conn, slow} sql = timeIt slow sql $ SQL.execute_ conn sql
+execute_ c sql = timeIt c sql $ SQL.execute_ (conn c) sql
 {-# INLINE execute_ #-}
 
 executeMany :: ToRow q => Connection -> Query -> [q] -> IO ()
-executeMany Connection {conn, slow} sql = timeIt slow sql . SQL.executeMany conn sql
+executeMany c sql = timeIt c sql . SQL.executeMany (conn c) sql
 {-# INLINE executeMany #-}
 
 query :: (ToRow q, FromRow r) => Connection -> Query -> q -> IO [r]
-query Connection {conn, slow} sql = timeIt slow sql . SQL.query conn sql
+query c sql = timeIt c sql . SQL.query (conn c) sql
 {-# INLINE query #-}
 
 query_ :: FromRow r => Connection -> Query -> IO [r]
-query_ Connection {conn, slow} sql = timeIt slow sql $ SQL.query_ conn sql
+query_ c sql = timeIt c sql $ SQL.query_ (conn c) sql
 {-# INLINE query_ #-}
 
 $(J.deriveJSON defaultJSON ''SlowQueryStats)

tests/AgentTests/FunctionalAPITests.hs

@@ -3104,7 +3104,7 @@ insertUser :: DBStore -> IO ()
 insertUser st = withTransaction st (`DB.execute_` "INSERT INTO users DEFAULT VALUES")
 #else
 createStore :: String -> IO (Either MigrationError DBStore)
-createStore dbPath = createAgentStore (DBOpts dbPath "" False True) MCError
+createStore dbPath = createAgentStore (DBOpts dbPath "" False True DB.TQOff) MCError
 
 insertUser :: DBStore -> IO ()
 insertUser st = withTransaction st (`DB.execute_` "INSERT INTO users (user_id) VALUES (1)")

tests/AgentTests/MigrationTests.hs

@@ -228,7 +228,8 @@ createStore randSuffix migrations confirmMigrations = do
           dbFilePath = testDB randSuffix,
           dbKey = "",
           keepKey = False,
-          vacuum = True
+          vacuum = True,
+          track = DB.TQOff
         }
   createDBStore dbOpts migrations confirmMigrations
 

tests/AgentTests/SQLiteTests.hs

@@ -70,7 +70,7 @@ withStore2 = before connect2 . after (removeStore . fst)
     connect2 :: IO (DBStore, DBStore)
     connect2 = do
       s1@DBStore {dbFilePath} <- createStore'
-      s2 <- connectSQLiteStore dbFilePath "" False
+      s2 <- connectSQLiteStore dbFilePath "" False DB.TQOff
       pure (s1, s2)
 
 createStore' :: IO DBStore
@@ -81,7 +81,7 @@ createEncryptedStore key keepKey = do
   -- Randomize DB file name to avoid SQLite IO errors supposedly caused by asynchronous
   -- IO operations on multiple similarly named files; error seems to be environment specific
   r <- randomIO :: IO Word32
-  Right st <- createDBStore (DBOpts (testDB <> show r) key keepKey True) Migrations.app MCError
+  Right st <- createDBStore (DBOpts (testDB <> show r) key keepKey True DB.TQOff) Migrations.app MCError
   withTransaction' st (`SQL.execute_` "INSERT INTO users (user_id) VALUES (1);")
   pure st
 

tests/AgentTests/SchemaDump.hs

@@ -12,6 +12,7 @@ import Database.SQLite.Simple (Only (..))
 import qualified Database.SQLite.Simple as SQL
 import Simplex.Messaging.Agent.Store.SQLite
 import Simplex.Messaging.Agent.Store.SQLite.Common (withTransaction')
+import Simplex.Messaging.Agent.Store.SQLite.DB (TrackQueries (..))
 import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
 import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfirmation (..), MigrationsToRun (..), toDownMigration)
 import Simplex.Messaging.Util (ifM)
@@ -49,15 +50,15 @@ testVerifySchemaDump :: IO ()
 testVerifySchemaDump = do
   savedSchema <- ifM (doesFileExist appSchema) (readFile appSchema) (pure "")
   savedSchema `deepseq` pure ()
-  void $ createDBStore (DBOpts testDB "" False True) Migrations.app MCConsole
+  void $ createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCConsole
   getSchema testDB appSchema `shouldReturn` savedSchema
   removeFile testDB
 
 testVerifyLintFKeyIndexes :: IO ()
 testVerifyLintFKeyIndexes = do
   savedLint <- ifM (doesFileExist appLint) (readFile appLint) (pure "")
   savedLint `deepseq` pure ()
-  void $ createDBStore (DBOpts testDB "" False True) Migrations.app MCConsole
+  void $ createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCConsole
   getLintFKeyIndexes testDB "tests/tmp/agent_lint.sql" `shouldReturn` savedLint
   removeFile testDB
 
@@ -70,7 +71,7 @@ withTmpFiles =
 testSchemaMigrations :: IO ()
 testSchemaMigrations = do
   let noDownMigrations = dropWhileEnd (\Migration {down} -> isJust down) Migrations.app
-  Right st <- createDBStore (DBOpts testDB "" False True) noDownMigrations MCError
+  Right st <- createDBStore (DBOpts testDB "" False True TQOff) noDownMigrations MCError
   mapM_ (testDownMigration st) $ drop (length noDownMigrations) Migrations.app
   closeDBStore st
   removeFile testDB
@@ -93,19 +94,19 @@ testSchemaMigrations = do
 
 testUsersMigrationNew :: IO ()
 testUsersMigrationNew = do
-  Right st <- createDBStore (DBOpts testDB "" False True) Migrations.app MCError
+  Right st <- createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCError
   withTransaction' st (`SQL.query_` "SELECT user_id FROM users;")
     `shouldReturn` ([] :: [Only Int])
   closeDBStore st
 
 testUsersMigrationOld :: IO ()
 testUsersMigrationOld = do
   let beforeUsers = takeWhile (("m20230110_users" /=) . name) Migrations.app
-  Right st <- createDBStore (DBOpts testDB "" False True) beforeUsers MCError
+  Right st <- createDBStore (DBOpts testDB "" False True TQOff) beforeUsers MCError
   withTransaction' st (`SQL.query_` "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'users';")
     `shouldReturn` ([] :: [Only String])
   closeDBStore st
-  Right st' <- createDBStore (DBOpts testDB "" False True) Migrations.app MCYesUp
+  Right st' <- createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCYesUp
   withTransaction' st' (`SQL.query_` "SELECT user_id FROM users;")
     `shouldReturn` ([Only (1 :: Int)])
   closeDBStore st'