From 53aecb5843e5ac6c13cb34c119c4c31d4ee77407 Mon Sep 17 00:00:00 2001
From: Simon Marlow <smarlow@meta.com>
Date: Mon, 9 Dec 2024 05:24:00 -0800
Subject: [PATCH] add batch sizes to logger

Reviewed By: iamirzhan

Differential Revision: D66756347

fbshipit-source-id: f152addab017e7d2e800278d3fc297bcdd03c086
---
 glean/db/Glean/Backend/Logging.hs | 20 +++++++++++++++++++-
 glean/db/Glean/Database/Writes.hs | 17 ++++++++++++++---
 glean/db/Glean/Logger/Server.hs   |  4 ++++
 3 files changed, 37 insertions(+), 4 deletions(-)

diff --git a/glean/db/Glean/Backend/Logging.hs b/glean/db/Glean/Backend/Logging.hs
index 683eb5f35..af137e0fa 100644
--- a/glean/db/Glean/Backend/Logging.hs
+++ b/glean/db/Glean/Backend/Logging.hs
@@ -29,6 +29,7 @@ import Glean.Backend.Local ()
 import Glean.Logger.Server as Logger
 import qualified Glean.Database.List as Database
 import qualified Glean.Database.Types as Database
+import Glean.Database.Writes (batchOwnedSize, batchDependenciesSize)
 import Glean.Logger
 import qualified Glean.Types as Thrift
 import Util.Time
@@ -142,7 +143,7 @@ instance Backend LoggingBackend where
       deleteDatabase env repo
   enqueueBatch (LoggingBackend env) cbatch =
     loggingAction
-      (runLogRepo "enqueueBatch" env (Thrift.computedBatch_repo cbatch))
+      (runLogEnqueueBatch "enqueueBatch" env cbatch)
       (const mempty) $
         enqueueBatch env cbatch
   enqueueJsonBatch (LoggingBackend env) repo batch =
@@ -233,6 +234,23 @@ runLogQueryBatch cmd env repo Thrift.UserQueryBatch{..} log =
     , maybe mempty logQueryClientInfo userQueryBatch_client_info
     ]
 
+runLogEnqueueBatch
+  :: Text
+  -> Database.Env
+  -> Thrift.ComputedBatch
+  -> GleanServerLog
+  -> IO ()
+runLogEnqueueBatch cmd env Thrift.ComputedBatch{..} log =
+  let !Thrift.Batch{..} = computedBatch_batch in
+  runLogRepo cmd env computedBatch_repo $ mconcat
+    [ log
+    , Logger.SetBatchFactsSize $ ByteString.length batch_facts
+    , Logger.SetBatchFactsCount $ fromIntegral $
+        Thrift.batch_count computedBatch_batch
+    , Logger.SetBatchOwnedSize $ batchOwnedSize batch_owned
+    , Logger.SetBatchDependenciesSize $ batchDependenciesSize batch_dependencies
+    ]
+
 logQueryOptions :: Thrift.UserQueryOptions -> GleanServerLog
 logQueryOptions Thrift.UserQueryOptions{..} = mconcat
   [ Logger.SetNoBase64Binary userQueryOptions_no_base64_binary
diff --git a/glean/db/Glean/Database/Writes.hs b/glean/db/Glean/Database/Writes.hs
index 83ec5fb52..420c0fc18 100644
--- a/glean/db/Glean/Database/Writes.hs
+++ b/glean/db/Glean/Database/Writes.hs
@@ -38,6 +38,8 @@ module Glean.Database.Writes
   , reapWrites
   , writerThread
   , deleteWriteQueues
+  , batchOwnedSize
+  , batchDependenciesSize
   ) where
 
 import Control.Concurrent
@@ -307,15 +309,24 @@ enqueueBatch env ComputedBatch{..} ownership = do
 batchSize :: Thrift.Batch -> Int
 batchSize Thrift.Batch{..} =
   ByteString.length batch_facts +
-  Monoid.getSum (foldMap (Monoid.Sum . storableSize) batch_owned) +
-  Monoid.getSum (foldMap (Monoid.Sum . depsSize) batch_dependencies)
+  batchOwnedSize batch_owned +
+  batchDependenciesSize batch_dependencies
+
+batchOwnedSize :: HashMap UnitName ListOfIds -> Int
+batchOwnedSize = Monoid.getSum . foldMap (Monoid.Sum . storableSize)
+
+batchDependenciesSize :: HashMap Id [FactDependencies] -> Int
+batchDependenciesSize deps =
+  Monoid.getSum (foldMap (Monoid.Sum . depsSize) deps)
   where
-  storableSize = snd . VS.unsafeToForeignPtr0
   depsSize deps = sum
     [ storableSize f + storableSize d
     | FactDependencies f d <- deps
     ]
 
+storableSize :: VS.Storable a => VS.Vector a -> Int
+storableSize = snd . VS.unsafeToForeignPtr0
+
 enqueueJsonBatch
   :: Env
   -> Repo
diff --git a/glean/db/Glean/Logger/Server.hs b/glean/db/Glean/Logger/Server.hs
index f7c5e827f..1473f4144 100644
--- a/glean/db/Glean/Logger/Server.hs
+++ b/glean/db/Glean/Logger/Server.hs
@@ -51,6 +51,10 @@ data GleanServerLog
   | SetRequestContinuationSize Int
   | SetResponseContinuationSize Int
   | SetSchemaId Text
+  | SetBatchFactsSize Int
+  | SetBatchFactsCount Int
+  | SetBatchOwnedSize Int
+  | SetBatchDependenciesSize Int
 
 instance Semigroup GleanServerLog where
   a <> b = GleanServerLogAnd a b