Skip to content

Commit eda9e36

Browse files
epoberezkinepoberezkin
authored
agent: track queries (#1439)
1 parent 2318975 commit eda9e36

File tree

6 files changed

+59
-44
lines changed

6 files changed

+59
-44
lines changed

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

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -69,32 +69,33 @@ data DBOpts = DBOpts
6969
{ dbFilePath :: FilePath,
7070
dbKey :: ScrubbedBytes,
7171
keepKey :: Bool,
72-
vacuum :: Bool
72+
vacuum :: Bool,
73+
track :: DB.TrackQueries
7374
}
7475

7576
createDBStore :: DBOpts -> [Migration] -> MigrationConfirmation -> IO (Either MigrationError DBStore)
76-
createDBStore DBOpts {dbFilePath, dbKey, keepKey, vacuum} migrations confirmMigrations = do
77+
createDBStore DBOpts {dbFilePath, dbKey, keepKey, track, vacuum} migrations confirmMigrations = do
7778
let dbDir = takeDirectory dbFilePath
7879
createDirectoryIfMissing True dbDir
79-
st <- connectSQLiteStore dbFilePath dbKey keepKey
80+
st <- connectSQLiteStore dbFilePath dbKey keepKey track
8081
r <- migrateSchema st migrations confirmMigrations vacuum `onException` closeDBStore st
8182
case r of
8283
Right () -> pure $ Right st
8384
Left e -> closeDBStore st $> Left e
8485

85-
connectSQLiteStore :: FilePath -> ScrubbedBytes -> Bool -> IO DBStore
86-
connectSQLiteStore dbFilePath key keepKey = do
86+
connectSQLiteStore :: FilePath -> ScrubbedBytes -> Bool -> DB.TrackQueries -> IO DBStore
87+
connectSQLiteStore dbFilePath key keepKey track = do
8788
dbNew <- not <$> doesFileExist dbFilePath
88-
dbConn <- dbBusyLoop (connectDB dbFilePath key)
89+
dbConn <- dbBusyLoop (connectDB dbFilePath key track)
8990
dbConnection <- newMVar dbConn
9091
dbKey <- newTVarIO $! storeKey key keepKey
9192
dbClosed <- newTVarIO False
9293
dbSem <- newTVarIO 0
9394
pure DBStore {dbFilePath, dbKey, dbSem, dbConnection, dbNew, dbClosed}
9495

95-
connectDB :: FilePath -> ScrubbedBytes -> IO DB.Connection
96-
connectDB path key = do
97-
db <- DB.open path
96+
connectDB :: FilePath -> ScrubbedBytes -> DB.TrackQueries -> IO DB.Connection
97+
connectDB path key track = do
98+
db <- DB.open path track
9899
prepare db `onException` DB.close db
99100
-- _printPragmas db path
100101
pure db
@@ -127,12 +128,12 @@ openSQLiteStore_ DBStore {dbConnection, dbFilePath, dbKey, dbClosed} key keepKey
127128
bracketOnError
128129
(takeMVar dbConnection)
129130
(tryPutMVar dbConnection)
130-
$ \DB.Connection {slow} -> do
131-
DB.Connection {conn} <- connectDB dbFilePath key
131+
$ \DB.Connection {slow, track} -> do
132+
DB.Connection {conn} <- connectDB dbFilePath key track
132133
atomically $ do
133134
writeTVar dbClosed False
134135
writeTVar dbKey $! storeKey key keepKey
135-
putMVar dbConnection DB.Connection {conn, slow}
136+
putMVar dbConnection DB.Connection {conn, slow, track}
136137

137138
reopenDBStore :: DBStore -> IO ()
138139
reopenDBStore st@DBStore {dbKey, dbClosed} =

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

Lines changed: 34 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -11,6 +11,7 @@ module Simplex.Messaging.Agent.Store.SQLite.DB
1111
Binary (..),
1212
Connection (..),
1313
SlowQueryStats (..),
14+
TrackQueries (..),
1415
open,
1516
close,
1617
execute,
@@ -38,7 +39,7 @@ import Database.SQLite.Simple.ToField (ToField (..))
3839
import Simplex.Messaging.Parsers (defaultJSON)
3940
import Simplex.Messaging.TMap (TMap)
4041
import qualified Simplex.Messaging.TMap as TM
41-
import Simplex.Messaging.Util (diffToMilliseconds, tshow)
42+
import Simplex.Messaging.Util (diffToMicroseconds, tshow)
4243

4344
newtype BoolInt = BI {unBI :: Bool}
4445
deriving newtype (FromField, ToField)
@@ -48,9 +49,13 @@ newtype Binary = Binary {fromBinary :: ByteString}
4849

4950
data Connection = Connection
5051
{ conn :: SQL.Connection,
52+
track :: TrackQueries,
5153
slow :: TMap Query SlowQueryStats
5254
}
5355

56+
data TrackQueries = TQAll | TQSlow Int64 | TQOff
57+
deriving (Eq)
58+
5459
data SlowQueryStats = SlowQueryStats
5560
{ count :: Int64,
5661
timeMax :: Int64,
@@ -59,22 +64,29 @@ data SlowQueryStats = SlowQueryStats
5964
}
6065
deriving (Show)
6166

62-
timeIt :: TMap Query SlowQueryStats -> Query -> IO a -> IO a
63-
timeIt slow sql a = do
64-
t <- getCurrentTime
65-
r <-
66-
a `catch` \e -> do
67-
atomically $ TM.alter (Just . updateQueryErrors e) sql slow
68-
throwIO e
69-
t' <- getCurrentTime
70-
let diff = diffToMilliseconds $ diffUTCTime t' t
71-
when (diff > 1) $ atomically $ TM.alter (updateQueryStats diff) sql slow
72-
pure r
67+
timeIt :: Connection -> Query -> IO a -> IO a
68+
timeIt Connection {slow, track} sql a
69+
| track == TQOff = makeQuery
70+
| otherwise = do
71+
t <- getCurrentTime
72+
r <- makeQuery
73+
t' <- getCurrentTime
74+
let diff = diffToMicroseconds $ diffUTCTime t' t
75+
when (trackQuery diff) $ atomically $ TM.alter (updateQueryStats diff) sql slow
76+
pure r
7377
where
78+
makeQuery =
79+
a `catch` \e -> do
80+
atomically $ TM.alter (Just . updateQueryErrors e) sql slow
81+
throwIO e
82+
trackQuery diff = case track of
83+
TQOff -> False
84+
TQSlow t -> diff > t
85+
TQAll -> True
7486
updateQueryErrors :: SomeException -> Maybe SlowQueryStats -> SlowQueryStats
7587
updateQueryErrors e Nothing = SlowQueryStats 0 0 0 $ M.singleton (tshow e) 1
76-
updateQueryErrors e (Just stats@SlowQueryStats {errs}) =
77-
stats {errs = M.alter (Just . maybe 1 (+ 1)) (tshow e) errs}
88+
updateQueryErrors e (Just st@SlowQueryStats {errs}) =
89+
st {errs = M.alter (Just . maybe 1 (+ 1)) (tshow e) errs}
7890
updateQueryStats :: Int64 -> Maybe SlowQueryStats -> Maybe SlowQueryStats
7991
updateQueryStats diff Nothing = Just $ SlowQueryStats 1 diff diff M.empty
8092
updateQueryStats diff (Just SlowQueryStats {count, timeMax, timeAvg, errs}) =
@@ -86,33 +98,33 @@ timeIt slow sql a = do
8698
errs
8799
}
88100

89-
open :: String -> IO Connection
90-
open f = do
101+
open :: String -> TrackQueries -> IO Connection
102+
open f track = do
91103
conn <- SQL.open f
92104
slow <- TM.emptyIO
93-
pure Connection {conn, slow}
105+
pure Connection {conn, slow, track}
94106

95107
close :: Connection -> IO ()
96108
close = SQL.close . conn
97109

98110
execute :: ToRow q => Connection -> Query -> q -> IO ()
99-
execute Connection {conn, slow} sql = timeIt slow sql . SQL.execute conn sql
111+
execute c sql = timeIt c sql . SQL.execute (conn c) sql
100112
{-# INLINE execute #-}
101113

102114
execute_ :: Connection -> Query -> IO ()
103-
execute_ Connection {conn, slow} sql = timeIt slow sql $ SQL.execute_ conn sql
115+
execute_ c sql = timeIt c sql $ SQL.execute_ (conn c) sql
104116
{-# INLINE execute_ #-}
105117

106118
executeMany :: ToRow q => Connection -> Query -> [q] -> IO ()
107-
executeMany Connection {conn, slow} sql = timeIt slow sql . SQL.executeMany conn sql
119+
executeMany c sql = timeIt c sql . SQL.executeMany (conn c) sql
108120
{-# INLINE executeMany #-}
109121

110122
query :: (ToRow q, FromRow r) => Connection -> Query -> q -> IO [r]
111-
query Connection {conn, slow} sql = timeIt slow sql . SQL.query conn sql
123+
query c sql = timeIt c sql . SQL.query (conn c) sql
112124
{-# INLINE query #-}
113125

114126
query_ :: FromRow r => Connection -> Query -> IO [r]
115-
query_ Connection {conn, slow} sql = timeIt slow sql $ SQL.query_ conn sql
127+
query_ c sql = timeIt c sql $ SQL.query_ (conn c) sql
116128
{-# INLINE query_ #-}
117129

118130
$(J.deriveJSON defaultJSON ''SlowQueryStats)

‎tests/AgentTests/FunctionalAPITests.hs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3104,7 +3104,7 @@ insertUser :: DBStore -> IO ()
31043104
insertUser st = withTransaction st (`DB.execute_` "INSERT INTO users DEFAULT VALUES")
31053105
#else
31063106
createStore :: String -> IO (Either MigrationError DBStore)
3107-
createStore dbPath = createAgentStore (DBOpts dbPath "" False True) MCError
3107+
createStore dbPath = createAgentStore (DBOpts dbPath "" False True DB.TQOff) MCError
31083108

31093109
insertUser :: DBStore -> IO ()
31103110
insertUser st = withTransaction st (`DB.execute_` "INSERT INTO users (user_id) VALUES (1)")

‎tests/AgentTests/MigrationTests.hs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -228,7 +228,8 @@ createStore randSuffix migrations confirmMigrations = do
228228
dbFilePath = testDB randSuffix,
229229
dbKey = "",
230230
keepKey = False,
231-
vacuum = True
231+
vacuum = True,
232+
track = DB.TQOff
232233
}
233234
createDBStore dbOpts migrations confirmMigrations
234235

‎tests/AgentTests/SQLiteTests.hs‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -70,7 +70,7 @@ withStore2 = before connect2 . after (removeStore . fst)
7070
connect2 :: IO (DBStore, DBStore)
7171
connect2 = do
7272
s1@DBStore {dbFilePath} <- createStore'
73-
s2 <- connectSQLiteStore dbFilePath "" False
73+
s2 <- connectSQLiteStore dbFilePath "" False DB.TQOff
7474
pure (s1, s2)
7575

7676
createStore' :: IO DBStore
@@ -81,7 +81,7 @@ createEncryptedStore key keepKey = do
8181
-- Randomize DB file name to avoid SQLite IO errors supposedly caused by asynchronous
8282
-- IO operations on multiple similarly named files; error seems to be environment specific
8383
r <- randomIO :: IO Word32
84-
Right st <- createDBStore (DBOpts (testDB <> show r) key keepKey True) Migrations.app MCError
84+
Right st <- createDBStore (DBOpts (testDB <> show r) key keepKey True DB.TQOff) Migrations.app MCError
8585
withTransaction' st (`SQL.execute_` "INSERT INTO users (user_id) VALUES (1);")
8686
pure st
8787

‎tests/AgentTests/SchemaDump.hs‎

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import Database.SQLite.Simple (Only (..))
1212
import qualified Database.SQLite.Simple as SQL
1313
import Simplex.Messaging.Agent.Store.SQLite
1414
import Simplex.Messaging.Agent.Store.SQLite.Common (withTransaction')
15+
import Simplex.Messaging.Agent.Store.SQLite.DB (TrackQueries (..))
1516
import qualified Simplex.Messaging.Agent.Store.SQLite.Migrations as Migrations
1617
import Simplex.Messaging.Agent.Store.Shared (Migration (..), MigrationConfirmation (..), MigrationsToRun (..), toDownMigration)
1718
import Simplex.Messaging.Util (ifM)
@@ -49,15 +50,15 @@ testVerifySchemaDump :: IO ()
4950
testVerifySchemaDump = do
5051
savedSchema <- ifM (doesFileExist appSchema) (readFile appSchema) (pure "")
5152
savedSchema `deepseq` pure ()
52-
void $ createDBStore (DBOpts testDB "" False True) Migrations.app MCConsole
53+
void $ createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCConsole
5354
getSchema testDB appSchema `shouldReturn` savedSchema
5455
removeFile testDB
5556

5657
testVerifyLintFKeyIndexes :: IO ()
5758
testVerifyLintFKeyIndexes = do
5859
savedLint <- ifM (doesFileExist appLint) (readFile appLint) (pure "")
5960
savedLint `deepseq` pure ()
60-
void $ createDBStore (DBOpts testDB "" False True) Migrations.app MCConsole
61+
void $ createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCConsole
6162
getLintFKeyIndexes testDB "tests/tmp/agent_lint.sql" `shouldReturn` savedLint
6263
removeFile testDB
6364

@@ -70,7 +71,7 @@ withTmpFiles =
7071
testSchemaMigrations :: IO ()
7172
testSchemaMigrations = do
7273
let noDownMigrations = dropWhileEnd (\Migration {down} -> isJust down) Migrations.app
73-
Right st <- createDBStore (DBOpts testDB "" False True) noDownMigrations MCError
74+
Right st <- createDBStore (DBOpts testDB "" False True TQOff) noDownMigrations MCError
7475
mapM_ (testDownMigration st) $ drop (length noDownMigrations) Migrations.app
7576
closeDBStore st
7677
removeFile testDB
@@ -93,19 +94,19 @@ testSchemaMigrations = do
9394

9495
testUsersMigrationNew :: IO ()
9596
testUsersMigrationNew = do
96-
Right st <- createDBStore (DBOpts testDB "" False True) Migrations.app MCError
97+
Right st <- createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCError
9798
withTransaction' st (`SQL.query_` "SELECT user_id FROM users;")
9899
`shouldReturn` ([] :: [Only Int])
99100
closeDBStore st
100101

101102
testUsersMigrationOld :: IO ()
102103
testUsersMigrationOld = do
103104
let beforeUsers = takeWhile (("m20230110_users" /=) . name) Migrations.app
104-
Right st <- createDBStore (DBOpts testDB "" False True) beforeUsers MCError
105+
Right st <- createDBStore (DBOpts testDB "" False True TQOff) beforeUsers MCError
105106
withTransaction' st (`SQL.query_` "SELECT name FROM sqlite_master WHERE type = 'table' AND name = 'users';")
106107
`shouldReturn` ([] :: [Only String])
107108
closeDBStore st
108-
Right st' <- createDBStore (DBOpts testDB "" False True) Migrations.app MCYesUp
109+
Right st' <- createDBStore (DBOpts testDB "" False True TQOff) Migrations.app MCYesUp
109110
withTransaction' st' (`SQL.query_` "SELECT user_id FROM users;")
110111
`shouldReturn` ([Only (1 :: Int)])
111112
closeDBStore st'

0 commit comments

Comments
 (0)