feat(metrics): add memory usage metrics for more executor (#10351)

src/stream/src/cache/managed_lru.rs

@@ -39,24 +39,22 @@ pub struct ManagedLruCache<K, V, S = DefaultHasher, A: Clone + Allocator = Globa
  /// The heap size of keys/values
  kv_heap_size: usize,
  /// The metrics of memory usage
- memory_usage_metrics: Option<IntGauge>,
+ memory_usage_metrics: IntGauge,
  // Metrics info
- metrics_info: Option<MetricsInfo>,
+ metrics_info: MetricsInfo,
  /// The size reported last time
  last_reported_size_bytes: usize,
 }
 
 impl<K, V, S, A: Clone + Allocator> Drop for ManagedLruCache<K, V, S, A> {
  fn drop(&mut self) {
- if let Some(metrics) = &self.memory_usage_metrics {
- metrics.set(0.into());
- }
- if let Some(info) = &self.metrics_info {
- info.metrics
- .stream_memory_usage
- .remove_label_values(&[&info.table_id, &info.actor_id, &info.desc])
- .unwrap();
- }
+ let info = &self.metrics_info;
+ self.memory_usage_metrics.set(0.into());
+
+ info.metrics
+ .stream_memory_usage
+ .remove_label_values(&[&info.table_id, &info.actor_id, &info.desc])
+ .unwrap();
  }
 }
 
@@ -66,15 +64,16 @@ impl<K: Hash + Eq + EstimateSize, V: EstimateSize, S: BuildHasher, A: Clone + Al
  pub fn new_inner(
  inner: LruCache<K, V, S, A>,
  watermark_epoch: Arc<AtomicU64>,
- metrics_info: Option<MetricsInfo>,
+ metrics_info: MetricsInfo,
  ) -> Self {
- let memory_usage_metrics = metrics_info.as_ref().map(|info| {
- info.metrics.stream_memory_usage.with_label_values(&[
- &info.table_id,
- &info.actor_id,
- &info.desc,
- ])
- });
+ let memory_usage_metrics = metrics_info
+ .metrics
+ .stream_memory_usage
+ .with_label_values(&[
+ &metrics_info.table_id,
+ &metrics_info.actor_id,
+ &metrics_info.desc,
+ ]);
 
  Self {
  inner,
@@ -219,9 +218,7 @@ impl<K: Hash + Eq + EstimateSize, V: EstimateSize, S: BuildHasher, A: Clone + Al
  if self.kv_heap_size.abs_diff(self.last_reported_size_bytes)
  > REPORT_SIZE_EVERY_N_KB_CHANGE << 10
  {
- if let Some(metrics) = self.memory_usage_metrics.as_ref() {
- metrics.set(self.kv_heap_size as _);
- }
+ self.memory_usage_metrics.set(self.kv_heap_size as _);
  self.last_reported_size_bytes = self.kv_heap_size;
  true
  } else {
@@ -232,15 +229,9 @@ impl<K: Hash + Eq + EstimateSize, V: EstimateSize, S: BuildHasher, A: Clone + Al
 
 pub fn new_unbounded<K: Hash + Eq + EstimateSize, V: EstimateSize>(
  watermark_epoch: Arc<AtomicU64>,
-) -> ManagedLruCache<K, V> {
- ManagedLruCache::new_inner(LruCache::unbounded(), watermark_epoch, None)
-}
-
-pub fn new_unbounded_with_metrics<K: Hash + Eq + EstimateSize, V: EstimateSize>(
- watermark_epoch: Arc<AtomicU64>,
  metrics_info: MetricsInfo,
 ) -> ManagedLruCache<K, V> {
- ManagedLruCache::new_inner(LruCache::unbounded(), watermark_epoch, Some(metrics_info))
+ ManagedLruCache::new_inner(LruCache::unbounded(), watermark_epoch, metrics_info)
 }
 
 pub fn new_with_hasher_in<
@@ -257,7 +248,7 @@ pub fn new_with_hasher_in<
  ManagedLruCache::new_inner(
  LruCache::unbounded_with_hasher_in(hasher, alloc),
  watermark_epoch,
- Some(metrics_info),
+ metrics_info,
  )
 }
 
@@ -269,7 +260,7 @@ pub fn new_with_hasher<K: Hash + Eq + EstimateSize, V: EstimateSize, S: BuildHas
  ManagedLruCache::new_inner(
  LruCache::unbounded_with_hasher(hasher),
  watermark_epoch,
- Some(metrics_info),
+ metrics_info,
  )
 }
 
@@ -280,15 +271,15 @@ pub struct MutGuard<'a, V: EstimateSize> {
  // The total size of a collection
  total_size: &'a mut usize,
  last_reported_size_bytes: &'a mut usize,
- memory_usage_metrics: &'a mut Option<IntGauge>,
+ memory_usage_metrics: &'a mut IntGauge,
 }
 
 impl<'a, V: EstimateSize> MutGuard<'a, V> {
  pub fn new(
  inner: &'a mut V,
  total_size: &'a mut usize,
  last_reported_size_bytes: &'a mut usize,
- memory_usage_metrics: &'a mut Option<IntGauge>,
+ memory_usage_metrics: &'a mut IntGauge,
  ) -> Self {
  let original_val_size = inner.estimated_size();
  Self {
@@ -304,9 +295,7 @@ impl<'a, V: EstimateSize> MutGuard<'a, V> {
  if self.total_size.abs_diff(*self.last_reported_size_bytes)
  > REPORT_SIZE_EVERY_N_KB_CHANGE << 10
  {
- if let Some(metrics) = self.memory_usage_metrics.as_ref() {
- metrics.set(*self.total_size as _);
- }
+ self.memory_usage_metrics.set(*self.total_size as _);
  *self.last_reported_size_bytes = *self.total_size;
  true
  } else {
@@ -346,15 +335,15 @@ pub struct UnsafeMutGuard<V: EstimateSize> {
  // The total size of a collection
  total_size: NonNull<usize>,
  last_reported_size_bytes: NonNull<usize>,
- memory_usage_metrics: NonNull<Option<IntGauge>>,
+ memory_usage_metrics: NonNull<IntGauge>,
 }
 
 impl<V: EstimateSize> UnsafeMutGuard<V> {
  pub fn new(
  inner: &mut V,
  total_size: &mut usize,
  last_reported_size_bytes: &mut usize,
- memory_usage_metrics: &mut Option<IntGauge>,
+ memory_usage_metrics: &mut IntGauge,
  ) -> Self {
  let original_val_size = inner.estimated_size();
  Self {

src/stream/src/executor/aggregation/distinct.rs

@@ -26,7 +26,7 @@ use risingwave_common::util::iter_util::ZipEqFast;
 use risingwave_storage::StateStore;
 
 use super::{AggCall, GroupKey};
-use crate::cache::{new_unbounded_with_metrics, ManagedLruCache};
+use crate::cache::{new_unbounded, ManagedLruCache};
 use crate::common::metrics::MetricsInfo;
 use crate::common::table::state_table::StateTable;
 use crate::executor::monitor::StreamingMetrics;
@@ -45,7 +45,7 @@ struct ColumnDeduplicater<S: StateStore> {
 impl<S: StateStore> ColumnDeduplicater<S> {
  fn new(watermark_epoch: &Arc<AtomicU64>, metrics_info: MetricsInfo) -> Self {
  Self {
- cache: new_unbounded_with_metrics(watermark_epoch.clone(), metrics_info.clone()),
+ cache: new_unbounded(watermark_epoch.clone(), metrics_info.clone()),
  metrics_info,
  _phantom: PhantomData,
  }

src/stream/src/executor/dedup/append_only_dedup.rs

@@ -12,6 +12,8 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+use std::sync::Arc;
+
 use futures::{stream, StreamExt};
 use futures_async_stream::try_stream;
 use itertools::Itertools;
@@ -22,8 +24,10 @@ use risingwave_common::row::{OwnedRow, Row, RowExt};
 use risingwave_storage::StateStore;
 
 use super::cache::DedupCache;
+use crate::common::metrics::MetricsInfo;
 use crate::common::table::state_table::StateTable;
 use crate::executor::error::StreamExecutorError;
+use crate::executor::monitor::StreamingMetrics;
 use crate::executor::{
  expect_first_barrier, ActorContextRef, BoxedExecutor, BoxedMessageStream, Executor, Message,
  PkIndices, PkIndicesRef, StreamExecutorResult,
@@ -51,12 +55,15 @@ impl<S: StateStore> AppendOnlyDedupExecutor<S> {
  executor_id: u64,
  ctx: ActorContextRef,
  watermark_epoch: AtomicU64Ref,
+ metrics: Arc<StreamingMetrics>,
  ) -> Self {
  let schema = input.schema().clone();
+ let metrics_info =
+ MetricsInfo::new(metrics, state_table.table_id(), ctx.id, "AppendOnly Dedup");
  Self {
  input: Some(input),
  state_table,
- cache: DedupCache::new(watermark_epoch),
+ cache: DedupCache::new(watermark_epoch, metrics_info),
  pk_indices,
  identity: format!("AppendOnlyDedupExecutor {:X}", executor_id),
  schema,
@@ -259,6 +266,7 @@ mod tests {
  1,
  ActorContext::create(123),
  Arc::new(AtomicU64::new(0)),
+ Arc::new(StreamingMetrics::unused()),
  ))
  .execute();
 

src/stream/src/executor/dedup/cache.rs

@@ -17,6 +17,7 @@ use std::hash::Hash;
 use risingwave_common::estimate_size::EstimateSize;
 
 use crate::cache::{new_unbounded, ManagedLruCache};
+use crate::common::metrics::MetricsInfo;
 use crate::task::AtomicU64Ref;
 
 /// [`DedupCache`] is used for key deduplication. Currently, the cache behaves like a set that only
@@ -26,8 +27,8 @@ pub struct DedupCache<K: Hash + Eq + EstimateSize> {
 }
 
 impl<K: Hash + Eq + EstimateSize> DedupCache<K> {
- pub fn new(watermark_epoch: AtomicU64Ref) -> Self {
- let cache = new_unbounded(watermark_epoch);
+ pub fn new(watermark_epoch: AtomicU64Ref, metrics_info: MetricsInfo) -> Self {
+ let cache = new_unbounded(watermark_epoch, metrics_info);
  Self { inner: cache }
  }
 
@@ -69,10 +70,11 @@ mod tests {
  use std::sync::Arc;
 
  use super::DedupCache;
+ use crate::common::metrics::MetricsInfo;
 
  #[test]
  fn test_dedup_cache() {
- let mut cache = DedupCache::new(Arc::new(AtomicU64::new(10000)));
+ let mut cache = DedupCache::new(Arc::new(AtomicU64::new(10000)), MetricsInfo::for_test());
 
  cache.insert(10);
  assert!(cache.contains(&10));

src/stream/src/executor/lookup/cache.rs

@@ -19,6 +19,7 @@ use risingwave_common::estimate_size::{EstimateSize, KvSize, VecWithKvSize};
 use risingwave_common::row::{OwnedRow, Row, RowExt};
 
 use crate::cache::{new_unbounded, ManagedLruCache};
+use crate::common::metrics::MetricsInfo;
 use crate::task::AtomicU64Ref;
 
 /// A cache for lookup's arrangement side.
@@ -73,8 +74,8 @@ impl LookupCache {
  self.data.clear();
  }
 
- pub fn new(watermark_epoch: AtomicU64Ref) -> Self {
- let cache = new_unbounded(watermark_epoch);
+ pub fn new(watermark_epoch: AtomicU64Ref, metrics_info: MetricsInfo) -> Self {
+ let cache = new_unbounded(watermark_epoch, metrics_info);
  Self { data: cache }
  }
 }

src/stream/src/executor/lookup/impl_.rs

@@ -30,6 +30,7 @@ use risingwave_storage::StateStore;
 
 use super::sides::{stream_lookup_arrange_prev_epoch, stream_lookup_arrange_this_epoch};
 use crate::cache::cache_may_stale;
+use crate::common::metrics::MetricsInfo;
 use crate::common::StreamChunkBuilder;
 use crate::executor::error::{StreamExecutorError, StreamExecutorResult};
 use crate::executor::lookup::cache::LookupCache;
@@ -206,6 +207,13 @@ impl<S: StateStore> LookupExecutor<S> {
  "mismatched output schema"
  );
 
+ let metrics_info = MetricsInfo::new(
+ ctx.streaming_metrics.clone(),
+ storage_table.table_id().table_id(),
+ ctx.id,
+ "Lookup",
+ );
+
  Self {
  ctx,
  chunk_data_types,
@@ -230,7 +238,7 @@ impl<S: StateStore> LookupExecutor<S> {
  },
  column_mapping,
  key_indices_mapping,
- lookup_cache: LookupCache::new(watermark_epoch),
+ lookup_cache: LookupCache::new(watermark_epoch, metrics_info),
  chunk_size,
  }
  }

src/stream/src/executor/mview/materialize.rs

@@ -38,6 +38,7 @@ use risingwave_storage::mem_table::KeyOp;
 use risingwave_storage::StateStore;
 
 use crate::cache::{new_unbounded, ManagedLruCache};
+use crate::common::metrics::MetricsInfo;
 use crate::common::table::state_table::StateTableInner;
 use crate::executor::error::StreamExecutorError;
 use crate::executor::monitor::StreamingMetrics;
@@ -97,8 +98,9 @@ impl<S: StateStore, SD: ValueRowSerde> MaterializeExecutor<S, SD> {
  StateTableInner::from_table_catalog(table_catalog, store, vnodes).await
  };
 
- let actor_id = actor_context.id;
- let table_id = table_catalog.id;
+ let metrics_info =
+ MetricsInfo::new(metrics, table_catalog.id, actor_context.id, "Materialize");
+
  Self {
  input,
  state_table,
@@ -109,7 +111,7 @@ impl<S: StateStore, SD: ValueRowSerde> MaterializeExecutor<S, SD> {
  pk_indices: arrange_columns,
  identity: format!("MaterializeExecutor {:X}", executor_id),
  },
- materialize_cache: MaterializeCache::new(watermark_epoch, metrics, actor_id, table_id),
+ materialize_cache: MaterializeCache::new(watermark_epoch, metrics_info),
  conflict_behavior,
  }
  }
@@ -235,12 +237,7 @@ impl<S: StateStore> MaterializeExecutor<S, BasicSerde> {
  pk_indices: arrange_columns,
  identity: format!("MaterializeExecutor {:X}", executor_id),
  },
- materialize_cache: MaterializeCache::new(
- watermark_epoch,
- Arc::new(StreamingMetrics::unused()),
- 0,
- 0,
- ),
+ materialize_cache: MaterializeCache::new(watermark_epoch, MetricsInfo::for_test()),
  conflict_behavior,
  }
  }
@@ -436,10 +433,8 @@ impl<S: StateStore, SD: ValueRowSerde> std::fmt::Debug for MaterializeExecutor<S
 /// A cache for materialize executors.
 pub struct MaterializeCache<SD> {
  data: ManagedLruCache<Vec<u8>, CacheValue>,
+ metrics_info: MetricsInfo,
  _serde: PhantomData<SD>,
- metrics: Arc<StreamingMetrics>,
- actor_id: String,
- table_id: String,
 }
 
 #[derive(EnumAsInner, EstimateSize)]
@@ -451,19 +446,12 @@ pub enum CacheValue {
 type EmptyValue = ();
 
 impl<SD: ValueRowSerde> MaterializeCache<SD> {
- pub fn new(
- watermark_epoch: AtomicU64Ref,
- metrics: Arc<StreamingMetrics>,
- actor_id: u32,
- table_id: u32,
- ) -> Self {
- let cache = new_unbounded(watermark_epoch);
+ pub fn new(watermark_epoch: AtomicU64Ref, metrics_info: MetricsInfo) -> Self {
+ let cache = new_unbounded(watermark_epoch, metrics_info.clone());
  Self {
  data: cache,
+ metrics_info,
  _serde: PhantomData,
- metrics,
- actor_id: actor_id.to_string(),
- table_id: table_id.to_string(),
  }
  }
 
@@ -604,14 +592,16 @@ impl<SD: ValueRowSerde> MaterializeCache<SD> {
  ) -> StreamExecutorResult<()> {
  let mut futures = vec![];
  for key in keys {
- self.metrics
+ self.metrics_info
+ .metrics
  .materialize_cache_total_count
- .with_label_values(&[&self.table_id, &self.actor_id])
+ .with_label_values(&[&self.metrics_info.table_id, &self.metrics_info.actor_id])
  .inc();
  if self.data.contains(key) {
- self.metrics
+ self.metrics_info
+ .metrics
  .materialize_cache_hit_count
- .with_label_values(&[&self.table_id, &self.actor_id])
+ .with_label_values(&[&self.metrics_info.table_id, &self.metrics_info.actor_id])
  .inc();
  continue;
  }