Report duration and memory for queries and mutations (#41795)

Previously only for actions. GitOrigin-RevId: 0f2c3672748bc2af12c5ec554b82fa1845aec921

Author: nipunn1313

crates/application/src/function_log.rs

@@ -638,7 +638,10 @@ impl<RT: Runtime> FunctionExecutionLog<RT> {
  if was_cached {
  CallType::CachedQuery
  } else {
- CallType::UncachedQuery
+ CallType::UncachedQuery {
+ duration: execution_time,
+ memory_in_mb: outcome.memory_in_mb,
+ }
  },
  outcome.result.is_ok(),
  usage_stats,
@@ -794,6 +797,8 @@ impl<RT: Runtime> FunctionExecutionLog<RT> {
  context.execution_id,
  context.request_id.clone(),
  CallType::Mutation {
+ duration: execution_time,
+ memory_in_mb: outcome.memory_in_mb,
  occ_info: occ_info.clone(),
  },
  outcome.result.is_ok(),

crates/database/src/tests/usage_tracking.rs

@@ -117,7 +117,11 @@ async fn vector_insert_counts_usage_for_backfilling_indexes(rt: TestRuntime) ->
  test_udf_identifier(),
  ExecutionId::new(),
  RequestId::new(),
- CallType::Mutation { occ_info: None },
+ CallType::Mutation {
+ occ_info: None,
+ duration: Duration::ZERO,
+ memory_in_mb: 0,
+ },
  true,
  tx_usage.gather_user_stats(),
  )
@@ -301,7 +305,11 @@ async fn test_usage_tracking_basic_insert_and_get(rt: TestRuntime) -> anyhow::Re
  test_udf_identifier(),
  ExecutionId::new(),
  RequestId::new(),
- CallType::Mutation { occ_info: None },
+ CallType::Mutation {
+ occ_info: None,
+ duration: Duration::ZERO,
+ memory_in_mb: 0,
+ },
  true,
  tx_usage.gather_user_stats(),
  )
@@ -330,7 +338,11 @@ async fn test_usage_tracking_basic_insert_and_get(rt: TestRuntime) -> anyhow::Re
  test_udf_identifier(),
  ExecutionId::new(),
  RequestId::new(),
- CallType::Mutation { occ_info: None },
+ CallType::Mutation {
+ occ_info: None,
+ duration: Duration::ZERO,
+ memory_in_mb: 0,
+ },
  true,
  tx_usage.gather_user_stats(),
  )
@@ -376,7 +388,11 @@ async fn test_usage_tracking_insert_with_index(rt: TestRuntime) -> anyhow::Resul
  test_udf_identifier(),
  ExecutionId::new(),
  RequestId::new(),
- CallType::Mutation { occ_info: None },
+ CallType::Mutation {
+ occ_info: None,
+ duration: Duration::ZERO,
+ memory_in_mb: 0,
+ },
  true,
  tx_usage.gather_user_stats(),
  )
@@ -404,7 +420,11 @@ async fn test_usage_tracking_insert_with_index(rt: TestRuntime) -> anyhow::Resul
  test_udf_identifier(),
  ExecutionId::new(),
  RequestId::new(),
- CallType::Mutation { occ_info: None },
+ CallType::Mutation {
+ occ_info: None,
+ duration: Duration::ZERO,
+ memory_in_mb: 0,
+ },
  true,
  tx_usage.gather_user_stats(),
  )
@@ -435,7 +455,11 @@ async fn test_usage_tracking_insert_with_index(rt: TestRuntime) -> anyhow::Resul
  test_udf_identifier(),
  ExecutionId::new(),
  RequestId::new(),
- CallType::Mutation { occ_info: None },
+ CallType::Mutation {
+ occ_info: None,
+ duration: Duration::ZERO,
+ memory_in_mb: 0,
+ },
  true,
  tx_usage.gather_user_stats(),
  )

crates/isolate/src/environment/udf/async_syscall.rs

@@ -587,6 +587,7 @@ impl<RT: Runtime> AsyncSyscallProvider<RT> for DatabaseUdfEnvironment<RT> {
  udf_server_version: _,
  unix_timestamp: _,
  rng_seed: _,
+ memory_in_mb: _,
  } = outcome;
 
  log_run_udf(

crates/isolate/src/environment/udf/mod.rs

@@ -4,6 +4,7 @@ use common::{
  ResolvedComponentFunctionPath,
  },
  execution_context::ExecutionContext,
+ knobs::ISOLATE_MAX_USER_HEAP_SIZE,
 };
 use futures::{
  future::BoxFuture,
@@ -450,6 +451,9 @@ impl<RT: Runtime> DatabaseUdfEnvironment<RT> {
  ));
  },
  )?;
+ let memory_in_mb = (*ISOLATE_MAX_USER_HEAP_SIZE / (1 << 20))
+ .try_into()
+ .unwrap();
  let outcome = match self.udf_type {
  UdfType::Query => FunctionOutcome::Query(UdfOutcome {
  path: self.path.for_logging(),
@@ -468,6 +472,7 @@ impl<RT: Runtime> DatabaseUdfEnvironment<RT> {
  },
  syscall_trace: self.syscall_trace,
  udf_server_version: self.udf_server_version,
+ memory_in_mb,
  }),
  // TODO: Add num_writes and write_bandwidth to UdfOutcome,
  // and use them in log_mutation.
@@ -488,6 +493,7 @@ impl<RT: Runtime> DatabaseUdfEnvironment<RT> {
  },
  syscall_trace: self.syscall_trace,
  udf_server_version: self.udf_server_version,
+ memory_in_mb,
  }),
  _ => anyhow::bail!("UdfEnvironment should only run queries and mutations"),
  };

crates/isolate/src/isolate2/runner.rs

@@ -19,6 +19,7 @@ use common::{
  },
  errors::JsError,
  execution_context::ExecutionContext,
+ knobs::ISOLATE_MAX_USER_HEAP_SIZE,
  log_lines::{
  LogLevel,
  LogLine,
@@ -553,6 +554,7 @@ async fn run_request<RT: Runtime>(
  result: Err(js_error),
  syscall_trace: SyscallTrace::new(),
  udf_server_version,
+ memory_in_mb: 0,
  };
  return Ok(outcome);
  }
@@ -696,6 +698,9 @@ async fn run_request<RT: Runtime>(
  result: result.map(JsonPackedValue::pack),
  syscall_trace: provider.syscall_trace,
  udf_server_version,
+ memory_in_mb: (*ISOLATE_MAX_USER_HEAP_SIZE / (1 << 20))
+ .try_into()
+ .unwrap(),
  };
  Ok(outcome)
 }

crates/pb/protos/outcome.proto

@@ -31,6 +31,7 @@ message UdfOutcome {
  SyscallTrace syscall_trace = 8;
 
  optional bool observed_identity = 10;
+ uint64 memory_in_mb = 11;
 }
 
 message ActionOutcome {

crates/udf/src/udf_outcome.rs

@@ -21,10 +21,6 @@ use pb::{
  },
  outcome::UdfOutcome as UdfOutcomeProto,
 };
-#[cfg(any(test, feature = "testing"))]
-use proptest::prelude::Arbitrary;
-#[cfg(any(test, feature = "testing"))]
-use proptest::prelude::Strategy;
 use rand::Rng;
 use value::{
  heap_size::HeapSize,
@@ -37,7 +33,10 @@ use crate::{
 };
 
 #[derive(Debug, Clone)]
-#[cfg_attr(any(test, feature = "testing"), derive(PartialEq))]
+#[cfg_attr(
+ any(test, feature = "testing"),
+ derive(proptest_derive::Arbitrary, PartialEq)
+)]
 pub struct UdfOutcome {
  pub path: CanonicalizedComponentFunctionPath,
  pub arguments: ConvexArray,
@@ -59,63 +58,9 @@ pub struct UdfOutcome {
 
  pub syscall_trace: SyscallTrace,
 
+ #[cfg_attr(any(test, feature = "testing"), proptest(value = "None"))]
  pub udf_server_version: Option<semver::Version>,
-}
-
-#[cfg(any(test, feature = "testing"))]
-impl Arbitrary for UdfOutcome {
- type Parameters = ();
-
- type Strategy = impl Strategy<Value = UdfOutcome>;
-
- fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
- use proptest::prelude::*;
- (
- any::<CanonicalizedComponentFunctionPath>(),
- any::<ConvexArray>(),
- any::<InertIdentity>(),
- any::<[u8; 32]>(),
- any::<bool>(),
- any::<UnixTimestamp>(),
- any::<bool>(),
- any::<bool>(),
- any::<LogLines>(),
- any::<QueryJournal>(),
- any::<Result<JsonPackedValue, JsError>>(),
- any::<SyscallTrace>(),
- )
- .prop_map(
- |(
- path,
- arguments,
- identity,
- rng_seed,
- observed_rng,
- unix_timestamp,
- observed_time,
- observed_identity,
- log_lines,
- journal,
- result,
- syscall_trace,
- )| Self {
- path,
- arguments,
- identity,
- rng_seed,
- observed_rng,
- unix_timestamp,
- observed_time,
- observed_identity,
- log_lines,
- journal,
- result,
- syscall_trace,
- // Ok to not generate semver::Version because it is not serialized anyway
- udf_server_version: None,
- },
- )
- }
+ pub memory_in_mb: u64,
 }
 
 impl HeapSize for UdfOutcome {
@@ -148,6 +93,7 @@ impl TryFrom<UdfOutcome> for UdfOutcomeProto {
  result,
  syscall_trace,
  udf_server_version: _,
+ memory_in_mb,
  }: UdfOutcome,
  ) -> anyhow::Result<Self> {
  let result = match result {
@@ -166,6 +112,7 @@ impl TryFrom<UdfOutcome> for UdfOutcomeProto {
  }),
  syscall_trace: Some(syscall_trace.try_into()?),
  observed_identity: Some(observed_identity),
+ memory_in_mb,
  })
  }
 }
@@ -195,6 +142,7 @@ impl UdfOutcome {
  syscall_trace: SyscallTrace::new(),
  udf_server_version,
  observed_identity: false,
+ memory_in_mb: 0,
  })
  }
 
@@ -209,6 +157,7 @@ impl UdfOutcome {
  result,
  syscall_trace,
  observed_identity,
+ memory_in_mb,
  }: UdfOutcomeProto,
  path_and_args: ValidatedPathAndArgs,
  identity: InertIdentity,
@@ -249,6 +198,7 @@ impl UdfOutcome {
  udf_server_version,
  // TODO(lee): Remove the default once we've pushed all services.
  observed_identity: observed_identity.unwrap_or(true),
+ memory_in_mb,
  })
  }
 }

crates/udf/src/validation.rs

@@ -806,6 +806,7 @@ pub struct ValidatedUdfOutcome {
 
  pub udf_server_version: Option<semver::Version>,
  pub mutation_queue_length: Option<usize>,
+ pub memory_in_mb: u64,
 }
 
 impl HeapSize for ValidatedUdfOutcome {
@@ -845,6 +846,7 @@ impl ValidatedUdfOutcome {
  syscall_trace: SyscallTrace::new(),
  udf_server_version,
  mutation_queue_length: None,
+ memory_in_mb: 0,
  })
  }
 
@@ -868,6 +870,7 @@ impl ValidatedUdfOutcome {
  syscall_trace: outcome.syscall_trace,
  udf_server_version: outcome.udf_server_version,
  mutation_queue_length,
+ memory_in_mb: outcome.memory_in_mb,
  };
 
  // TODO(CX-6318) Don't pack json value until it's been validated.

crates/usage_tracking/src/lib.rs

@@ -112,8 +112,13 @@ pub enum CallType {
  },
  Export,
  CachedQuery,
- UncachedQuery,
+ UncachedQuery {
+ duration: Duration,
+ memory_in_mb: u64,
+ },
  Mutation {
+ duration: Duration,
+ memory_in_mb: u64,
  occ_info: Option<OccInfo>,
  },
  Import,
@@ -127,7 +132,7 @@ impl CallType {
  Self::Action { .. } => "action",
  Self::Export => "export",
  Self::CachedQuery => "cached_query",
- Self::UncachedQuery => "uncached_query",
+ Self::UncachedQuery { .. } => "uncached_query",
  Self::Mutation { .. } => "mutation",
  Self::HttpAction { .. } => "http_action",
  Self::Import => "import",
@@ -145,7 +150,7 @@ impl CallType {
 
  fn occ_document_id(&self) -> Option<String> {
  match self {
- Self::Mutation { occ_info } => {
+ Self::Mutation { occ_info, .. } => {
  occ_info.as_ref().and_then(|info| info.document_id.clone())
  },
  _ => None,
@@ -179,19 +184,21 @@ impl CallType {
 
  fn memory_megabytes(&self) -> u64 {
  match self {
- CallType::Action { memory_in_mb, .. } | CallType::HttpAction { memory_in_mb, .. } => {
- *memory_in_mb
- },
+ CallType::UncachedQuery { memory_in_mb, .. }
+ | CallType::Mutation { memory_in_mb, .. }
+ | CallType::Action { memory_in_mb, .. }
+ | CallType::HttpAction { memory_in_mb, .. } => *memory_in_mb,
  _ => 0,
  }
  }
 
  fn duration_millis(&self) -> u64 {
  match self {
- CallType::Action { duration, .. } | CallType::HttpAction { duration, .. } => {
- u64::try_from(duration.as_millis())
- .expect("Action was running for over 584 billion years??")
- },
+ CallType::UncachedQuery { duration, .. }
+ | CallType::Mutation { duration, .. }
+ | CallType::Action { duration, .. }
+ | CallType::HttpAction { duration, .. } => u64::try_from(duration.as_millis())
+ .expect("Function was running for over 584 billion years??"),
  _ => 0,
  }
  }