add logging for query result bytes

Summary: log the number of bytes in a query result so that we can use these logs to set the `max_bytes` field in T211119652

Reviewed By: simonmar

Differential Revision: D67976583

fbshipit-source-id: 62eea04faa9df7f4c2d3d25ca16f1d7e53672a4f

glean/db/Glean/Backend/Logging.hs

@@ -321,6 +321,8 @@ logQueryStats Thrift.UserQueryStats{..} = mconcat
       userQueryStats_compile_time_ns
   , maybe mempty (Logger.SetExecuteTimeUs . fromIntegral . (`quot` 1000))
       userQueryStats_execute_time_ns
+  , maybe mempty (Logger.SetQueryResultBytes . fromIntegral)
+      userQueryStats_result_bytes
   ]
 
 runLogDerivePredicate

glean/db/Glean/Logger/Server.hs

@@ -45,6 +45,7 @@ data GleanServerLog
   | SetBytecodeSize Int
   | SetCompileTimeUs Int
   | SetExecuteTimeUs Int
+  | SetQueryResultBytes Int
   | SetClientUnixname Text
   | SetClientApplication Text
   | SetClientName Text

glean/db/Glean/Query/UserQuery.hs

@@ -197,6 +197,9 @@ data Results stats fact = Results
 
     -- | Query execution time after compilation, for logging, in nanoseconds
   , resExecutionTime :: Maybe Word64
+
+  -- | Size of query results, for logging
+  , resResultBytes :: Maybe Int64
   }
 
 class Encoding e where
@@ -482,6 +485,7 @@ userQueryFactsImpl
         , resCompileTime = Nothing
         , resCodegenTime = Nothing
         , resExecutionTime = Nothing
+        , resResultBytes = Nothing
         }
 
   return $ if Thrift.userQueryOptions_omit_results opts
@@ -889,6 +893,7 @@ runQuery
           , resCompileTime = Just compileTime
           , resCodegenTime = Just codegenTime
           , resExecutionTime = Just queryResultsElapsedNs
+          , resResultBytes = Just queryResultsResultBytes
           }
 
     return $ if Thrift.userQueryOptions_omit_results opts
@@ -1052,6 +1057,7 @@ emptyResult = Results {
   , resCompileTime = Nothing
   , resCodegenTime = Nothing
   , resExecutionTime = Nothing
+  , resResultBytes = Nothing
   }
 
 
@@ -1210,6 +1216,8 @@ withStats io = do
         , Thrift.userQueryStats_execute_time_ns =
             fromIntegral <$> resExecutionTime res
         , Thrift.userQueryStats_full_scans = statFullScans $ resStats res
+        , Thrift.userQueryStats_result_bytes =
+            fromIntegral <$> resResultBytes res
         }
   return res{ resStats = stats }
 

glean/hs/Glean/RTS/Foreign/Query.hsc

@@ -72,6 +72,7 @@ data QueryResults = QueryResults
   , queryResultsStats :: Maybe (Map Int64 Int64)
   , queryResultsElapsedNs :: Word64
   , queryResultsCont :: Maybe ByteString
+  , queryResultsResultBytes :: Int64
   }
 
 data CompiledQuery = CompiledQuery
@@ -227,6 +228,7 @@ unpackResults wantStats maybePid (Results p) = do
       else return Nothing
   elapsed_ns <- (# peek facebook::glean::rts::QueryResults, elapsed_ns) p
   cont <- (# peek facebook::glean::rts::QueryResults, continuation) p
+  result_bytes <- (# peek facebook::glean::rts::QueryResults, result_bytes) p
 
   return QueryResults
     { queryResultsFacts = resultFacts
@@ -238,6 +240,7 @@ unpackResults wantStats maybePid (Results p) = do
         if ByteString.length contBytes == 0
           then Nothing
           else Just contBytes
+    , queryResultsResultBytes = result_bytes
     }
 
 interruptRunningQueries :: IO ()

glean/if/glean.thrift

@@ -767,7 +767,9 @@ struct UserQueryStats {
   10: optional i64 codegen_time_ns;
   // time to generate bytecode for the query
   11: list<PredicateRef> full_scans = [];
-// whether the query performs full predicate scans
+  // whether the query performs full predicate scans
+  12: optional i64 result_bytes;
+// query results size in bytes
 }
 
 # Results in Glean's internal binary representation

glean/if/glean_include.hs

@@ -51,6 +51,7 @@ instance Prelude.Semigroup UserQueryStats where
       , userQueryStats_result_count = result_count1
       , userQueryStats_codegen_time_ns = codegen_time_ns1
       , userQueryStats_full_scans = full_scans1
+      , userQueryStats_result_bytes = result_bytes1
       }
     UserQueryStats
       { userQueryStats_num_facts = num_facts2
@@ -63,6 +64,7 @@ instance Prelude.Semigroup UserQueryStats where
       , userQueryStats_result_count = result_count2
       , userQueryStats_codegen_time_ns = codegen_time_ns2
       , userQueryStats_full_scans = full_scans2
+      , userQueryStats_result_bytes = result_bytes2
       }
     = UserQueryStats
       { userQueryStats_num_facts = num_facts1 Prelude.+ num_facts2
@@ -84,6 +86,8 @@ instance Prelude.Semigroup UserQueryStats where
           fMaybe (Prelude.+) codegen_time_ns1 codegen_time_ns2
       , userQueryStats_full_scans =
           List.nub (full_scans1 <> full_scans2)
+      , userQueryStats_result_bytes =
+          fMaybe (Prelude.+) result_bytes1 result_bytes2
       }
       where
       fMaybe _ Prelude.Nothing a = a

glean/rts/query.cpp

@@ -219,6 +219,7 @@ struct QueryExecutor : SetOps {
   // query stats
   folly::F14FastMap<uint64_t, uint64_t> stats;
   bool wantStats;
+  uint64_t result_bytes;
 
   // iterators
   struct Iter {
@@ -494,6 +495,7 @@ std::unique_ptr<QueryResults> QueryExecutor::finish(
   res->nested_fact_pids = std::move(nested_result_pids);
   res->nested_fact_keys = std::move(nested_result_keys);
   res->nested_fact_values = std::move(nested_result_values);
+  res->result_bytes = std::move(result_bytes);
 
   if (cont.hasValue()) {
     res->continuation = cont->string();
@@ -588,12 +590,15 @@ std::unique_ptr<QueryResults> executeQuery(
             std::copy(context_.handlers_begin(), context_.handlers_end(), args);
         *args++ = 0; // unused
         *args++ = reinterpret_cast<uint64_t>(max_results);
+        uint64_t* max_bytes_p = args;
         *args++ = reinterpret_cast<uint64_t>(max_bytes);
 
         activation.execute();
         if (activation.suspended()) {
           cont = q.queryCont(activation);
         }
+
+        q.result_bytes = max_bytes - *max_bytes_p;
       });
 
   return q.finish(std::move(cont));

glean/rts/query.h

@@ -36,6 +36,7 @@ struct QueryResults {
   HsMap<uint64_t, uint64_t> stats;
   uint64_t elapsed_ns;
   HsString continuation;
+  uint64_t result_bytes;
 };
 
 enum class Depth { ResultsOnly, ExpandRecursive, ExpandPartial };

glean/tools/gleancli/test/test-glean-query.t

@@ -11,10 +11,10 @@
   {"id":[0-9]*,"key":{"name":"\w+","line":[0-9]+}} (re)
 
   $ cat "glean_stats_output.txt"
-  {"num_facts":1,"elapsed_ns":[0-9]*,"allocated_bytes":[0-9]*,"compile_time_ns":[0-9]*,"bytecode_size":[0-9]*,"execute_time_ns":[0-9]*,"result_count":1,"codegen_time_ns":[0-9]*,"full_scans":\[]} (re)
+  {"num_facts":1,"elapsed_ns":[0-9]*,"allocated_bytes":[0-9]*,"compile_time_ns":[0-9]*,"bytecode_size":[0-9]*,"execute_time_ns":[0-9]*,"result_count":1,"codegen_time_ns":[0-9]*,"full_scans":\[],"result_bytes":[0-9]*} (re)
 
   $ glean --schema $EXAMPLE query --db-name abc 'example.Class {name="Pet"}' --stats "glean_profile_output.txt" --profile
   {"id":[0-9]*,"key":{"name":"\w+","line":[0-9]+}} (re)
 
   $ cat "glean_profile_output.txt"
-  {"num_facts":1,"elapsed_ns":[0-9]*,"allocated_bytes":[0-9]*,"facts_searched":{"[0-9]*":1},"compile_time_ns":[0-9]*,"bytecode_size":[0-9]*,"execute_time_ns":[0-9]*,"result_count":1,"codegen_time_ns":[0-9]*,"full_scans":\[]} (re)
+  {"num_facts":1,"elapsed_ns":[0-9]*,"allocated_bytes":[0-9]*,"facts_searched":{"[0-9]*":1},"compile_time_ns":[0-9]*,"bytecode_size":[0-9]*,"execute_time_ns":[0-9]*,"result_count":1,"codegen_time_ns":[0-9]*,"full_scans":\[],"result_bytes":[0-9]*} (re)