feat: enable verbose await-tree by default in production (#10144) (#10537)

Signed-off-by: Bugen Zhao <[email protected]>
Co-authored-by: Bugen Zhao <[email protected]>

Author: lmatz

Cargo.lock

@@ -417,12 +417,27 @@ pub struct FileCacheConfig {
     pub unrecognized: Unrecognized<Self>,
 }
 
-#[derive(Debug, Default, Clone, ValueEnum, Serialize, Deserialize)]
+#[derive(Debug, Default, Clone, Copy, ValueEnum, Serialize, Deserialize)]
 pub enum AsyncStackTraceOption {
+    /// Disabled.
     Off,
-    #[default]
+    /// Enabled with basic instruments.
     On,
-    Verbose,
+    /// Enabled with extra verbose instruments in release build.
+    /// Behaves the same as `on` in debug build due to performance concern.
+    #[default]
+    #[clap(alias = "verbose")]
+    ReleaseVerbose,
+}
+
+impl AsyncStackTraceOption {
+    pub fn is_verbose(self) -> Option<bool> {
+        match self {
+            Self::Off => None,
+            Self::On => Some(false),
+            Self::ReleaseVerbose => Some(!cfg!(debug_assertions)),
+        }
+    }
 }
 
 serde_with::with_prefix!(streaming_prefix "stream_");
@@ -716,7 +731,7 @@ mod default {
         }
 
         pub fn async_stack_trace() -> AsyncStackTraceOption {
-            AsyncStackTraceOption::On
+            AsyncStackTraceOption::default()
         }
 
         pub fn unique_user_stream_errors() -> usize {

src/common/src/config.rs

@@ -16,7 +16,7 @@ normal = ["workspace-hack"]
 [dependencies]
 anyhow = "1"
 async-trait = "0.1"
-await-tree = "0.1.1"
+await-tree = { workspace = true }
 clap = { version = "4", features = ["derive"] }
 either = "1"
 futures = { version = "0.3", default-features = false, features = ["alloc"] }

src/compute/Cargo.toml

@@ -264,7 +264,7 @@ pub async fn compute_node_serve(
     let await_tree_config = match &config.streaming.async_stack_trace {
         AsyncStackTraceOption::Off => None,
         c => await_tree::ConfigBuilder::default()
-            .verbose(matches!(c, AsyncStackTraceOption::Verbose))
+            .verbose(c.is_verbose().unwrap())
             .build()
             .ok(),
     };

src/compute/src/server.rs

@@ -36,7 +36,7 @@ batch_chunk_size = 1024
 [streaming]
 in_flight_barrier_nums = 10000
 enable_jaeger_tracing = false
-async_stack_trace = "On"
+async_stack_trace = "ReleaseVerbose"
 unique_user_stream_errors = 10
 
 [streaming.developer]

src/config/example.toml

@@ -21,6 +21,7 @@ arrow-array = "36"
 arrow-schema = "36"
 async-trait = "0.1"
 auto_enums = "0.8"
+await-tree = { workspace = true }
 chrono = { version = "0.4", default-features = false, features = ["clock", "std"] }
 chrono-tz = { version = "0.7", features = ["case-insensitive"] }
 ctor = "0.1"

src/expr/Cargo.toml

@@ -17,6 +17,7 @@ use std::convert::TryFrom;
 use std::sync::{Arc, LazyLock, Mutex, Weak};
 
 use arrow_schema::{Field, Schema, SchemaRef};
+use await_tree::InstrumentAwait;
 use risingwave_common::array::{ArrayImpl, ArrayRef, DataChunk};
 use risingwave_common::row::OwnedRow;
 use risingwave_common::types::{DataType, Datum};
@@ -36,6 +37,7 @@ pub struct UdfExpression {
     arg_schema: SchemaRef,
     client: Arc<ArrowFlightUdfClient>,
     identifier: String,
+    span: await_tree::Span,
 }
 
 #[cfg(not(madsim))]
@@ -81,7 +83,11 @@ impl UdfExpression {
         let input =
             arrow_array::RecordBatch::try_new_with_options(self.arg_schema.clone(), columns, &opts)
                 .expect("failed to build record batch");
-        let output = self.client.call(&self.identifier, input).await?;
+        let output = self
+            .client
+            .call(&self.identifier, input)
+            .instrument_await(self.span.clone())
+            .await?;
         if output.num_rows() != vis.len() {
             bail!(
                 "UDF returned {} rows, but expected {}",
@@ -124,6 +130,7 @@ impl<'a> TryFrom<&'a ExprNode> for UdfExpression {
             arg_schema,
             client,
             identifier: udf.identifier.clone(),
+            span: format!("expr_udf_call ({})", udf.identifier).into(),
         })
     }
 }

src/expr/src/expr/expr_udf.rs

@@ -10,7 +10,7 @@ repository = { workspace = true }
 
 [dependencies]
 async-trait = "0.1"
-await-tree = "0.1.1"
+await-tree = { workspace = true }
 aws-config = { workspace = true }
 aws-sdk-s3 = { version = "0.2", package = "madsim-aws-sdk-s3" }
 aws-smithy-http = { workspace = true }

src/object_store/Cargo.toml

@@ -17,7 +17,7 @@ normal = ["workspace-hack"]
 arc-swap = "1"
 async-trait = "0.1"
 auto_enums = { version = "0.8", features = ["futures03"] }
-await-tree = "0.1.1"
+await-tree = { workspace = true }
 bytes = { version = "1", features = ["serde"] }
 crossbeam = "0.8.1"
 dashmap = { version = "5", default-features = false }

src/storage/Cargo.toml

@@ -185,7 +185,11 @@ impl<S: LocalStateStore> LocalStateStore for MonitoredStateStore<S> {
                 .may_exist_duration
                 .with_label_values(&[table_id_label.as_str()])
                 .start_timer();
-            let res = self.inner.may_exist(key_range, read_options).await;
+            let res = self
+                .inner
+                .may_exist(key_range, read_options)
+                .verbose_instrument_await("store_may_exist")
+                .await;
             timer.observe_duration();
             res
         }
@@ -217,7 +221,9 @@ impl<S: LocalStateStore> LocalStateStore for MonitoredStateStore<S> {
 
     fn flush(&mut self, delete_ranges: Vec<(Bound<Bytes>, Bound<Bytes>)>) -> Self::FlushFuture<'_> {
         // TODO: collect metrics
-        self.inner.flush(delete_ranges)
+        self.inner
+            .flush(delete_ranges)
+            .verbose_instrument_await("store_flush")
     }
 
     fn epoch(&self) -> u64 {
@@ -264,7 +270,7 @@ impl<S: StateStore> StateStore for MonitoredStateStore<S> {
             let sync_result = self
                 .inner
                 .sync(epoch)
-                .instrument_await("store_await_sync")
+                .instrument_await("store_sync")
                 .await
                 .inspect_err(|e| error!("Failed in sync: {:?}", e))?;
             timer.observe_duration();

src/storage/src/monitor/monitored_store.rs

@@ -19,7 +19,7 @@ anyhow = "1"
 async-recursion = "1"
 async-stream = "0.3"
 async-trait = "0.1"
-await-tree = "0.1.1"
+await-tree = { workspace = true }
 bytes = "1"
 dyn-clone = "1"
 educe = "0.4"

src/stream/Cargo.toml

@@ -15,7 +15,7 @@ ignored = ["workspace-hack"]
 normal = ["workspace-hack"]
 
 [dependencies]
-await-tree = "0.1.1"
+await-tree = { workspace = true }
 console = "0.15"
 console-subscriber = "0.1.8"
 futures = { version = "0.3", default-features = false, features = ["alloc"] }