use query_planning::{correlation::QueryId, plan::GraphAlgorithm};
use serde_json::{Value, json};
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum GraphMeasuredPeakProvenance {
CacheScopeHighWater,
Unavailable,
}
impl GraphMeasuredPeakProvenance {
pub const fn as_str(self) -> &'static str {
match self {
Self::CacheScopeHighWater => "cache_scope_high_water",
Self::Unavailable => "unavailable",
}
}
}
impl std::fmt::Display for GraphMeasuredPeakProvenance {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum GraphClaimExecutionState {
Pending,
AlgorithmStarted,
}
impl GraphClaimExecutionState {
pub const fn as_str(self) -> &'static str {
match self {
Self::Pending => "pending",
Self::AlgorithmStarted => "algorithm_started",
}
}
}
impl std::fmt::Display for GraphClaimExecutionState {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str(self.as_str())
}
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct GraphMeasuredPeakRow {
pub query_id: Option<QueryId>,
pub node_id: String,
pub node_index: usize,
pub algorithm: GraphAlgorithm,
pub execution_state: GraphClaimExecutionState,
pub measured_peak_bytes: Option<u64>,
pub measured_peak_provenance: GraphMeasuredPeakProvenance,
}
impl GraphMeasuredPeakRow {
pub fn pending(
query_id: Option<QueryId>,
node_id: impl Into<String>,
node_index: usize,
algorithm: GraphAlgorithm,
) -> Self {
Self {
query_id,
node_id: node_id.into(),
node_index,
algorithm,
execution_state: GraphClaimExecutionState::Pending,
measured_peak_bytes: None,
measured_peak_provenance: GraphMeasuredPeakProvenance::Unavailable,
}
}
pub(crate) fn record_execution_state(&mut self, execution_state: GraphClaimExecutionState) {
if self.execution_state == GraphClaimExecutionState::Pending {
self.execution_state = execution_state;
}
}
#[cfg(any(test, feature = "test-hooks"))]
#[doc(hidden)]
pub fn record_execution_state_for_test(&mut self, execution_state: GraphClaimExecutionState) {
self.record_execution_state(execution_state);
}
pub fn record_measured_peak(
&mut self,
measured_peak_bytes: Option<u64>,
measured_peak_provenance: GraphMeasuredPeakProvenance,
) {
match measured_peak_provenance {
GraphMeasuredPeakProvenance::Unavailable => self.measured_peak_bytes = None,
GraphMeasuredPeakProvenance::CacheScopeHighWater => self.measured_peak_bytes = measured_peak_bytes,
}
self.measured_peak_provenance = measured_peak_provenance;
}
pub fn to_json(&self) -> Value {
json!({
"query_id": self.query_id.as_ref().map(QueryId::as_str),
"node_id": self.node_id,
"node_index": self.node_index,
"algorithm": self.algorithm.as_str(),
"execution_state": self.execution_state.as_str(),
"measured_peak_bytes": self.measured_peak_bytes,
"measured_peak_provenance": self.measured_peak_provenance.as_str(),
})
}
}