feat: Add sampled trace support (#17850)

Signed-off-by: Xuanwo <[email protected]>

Author: Xuanwo

src/query/service/src/servers/mysql/mysql_interactive_worker.rs

@@ -45,6 +45,8 @@ use opensrv_mysql::InitWriter;
 use opensrv_mysql::ParamParser;
 use opensrv_mysql::QueryResultWriter;
 use opensrv_mysql::StatementMetaWriter;
+use rand::thread_rng;
+use rand::Rng as _;
 use rand::RngCore;
 use uuid::Uuid;
 
@@ -192,7 +194,9 @@ impl<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> for InteractiveWorke
         writer: QueryResultWriter<'a, W>,
     ) -> Result<()> {
         let query_id = Uuid::new_v4().to_string();
-        let root = Span::root(func_path!(), SpanContext::random())
+        let sampled =
+            thread_rng().gen_range(0..100) <= self.base.session.get_trace_sample_rate()?;
+        let root = Span::root(func_path!(), SpanContext::random().sampled(sampled))
             .with_properties(|| self.base.session.to_fastrace_properties());
 
         let mut tracking_payload = ThreadTracker::new_tracking_payload();

src/query/service/src/sessions/session.rs

@@ -388,6 +388,10 @@ impl Session {
         self.session_ctx.set_temp_tbl_mgr(temp_tbl_mgr)
     }
 
+    pub fn get_trace_sample_rate(&self) -> Result<u64> {
+        self.session_ctx.get_settings().get_trace_sample_rate()
+    }
+
     pub fn set_query_priority(&self, priority: u8) {
         if let Some(context_shared) = self.session_ctx.get_query_context_shared() {
             context_shared.set_priority(priority);

src/query/settings/src/settings_default.rs

@@ -1271,6 +1271,13 @@ impl DefaultSettings {
                     scope: SettingScope::Both,
                     range: Some(SettingRange::Numeric(0..=1)),
                 }),
+                ("trace_sample_rate", DefaultSettingValue {
+                    value: UserSettingValue::UInt64(1),
+                    desc: "Setting the trace sample rate. The value should be between '0' and '100'",
+                    mode: SettingMode::Both,
+                    scope: SettingScope::Both,
+                    range: Some(SettingRange::Numeric(0..=100)),
+                }),
             ]);
 
             Ok(Arc::new(DefaultSettings {

src/query/settings/src/settings_getter_setter.rs

@@ -936,4 +936,8 @@ impl Settings {
     pub fn get_statement_queue_ttl_in_seconds(&self) -> Result<u64> {
         self.try_get_u64("statement_queue_ttl_in_seconds")
     }
+
+    pub fn get_trace_sample_rate(&self) -> Result<u64> {
+        self.try_get_u64("trace_sample_rate")
+    }
 }