feat: telemetry

proto/meta.proto

@@ -13,6 +13,17 @@ import "user.proto";
 option java_package = "com.risingwave.proto";
 option optimize_for = SPEED;
 
+message GetTelemetryInfoRequest {}
+
+message TelemetryInfoResponse {
+  optional string tracking_id = 1;
+}
+
+service TelemetryInfoService {
+  // Request telemetry info from meta node
+  rpc GetTelemetryInfo(GetTelemetryInfoRequest) returns (TelemetryInfoResponse);
+}
+
 message HeartbeatRequest {
   message ExtraInfo {
     oneof info {
@@ -344,6 +355,7 @@ message SystemParams {
   optional string data_directory = 7;
   optional string backup_storage_url = 8;
   optional string backup_storage_directory = 9;
+  optional bool telemetry_enabled = 10;
 }
 
 message GetSystemParamsRequest {}

src/common/Cargo.toml

@@ -49,6 +49,7 @@ prometheus = { version = "0.13" }
 prost = "0.11"
 rand = "0.8"
 regex = "1"
+reqwest = { version = "0.11", features = ["json"] }
 risingwave_pb = { path = "../prost" }
 rust_decimal = { version = "1", features = ["db-tokio-postgres"] }
 ryu = "1.0"
@@ -64,6 +65,7 @@ toml = "0.5"
 tonic = { version = "0.2", package = "madsim-tonic" }
 tracing = "0.1"
 url = "2"
+uuid = "1.2.2"
 
 [target.'cfg(not(madsim))'.dependencies]
 workspace-hack = { path = "../workspace-hack" }

src/common/src/config.rs

@@ -215,6 +215,9 @@ pub struct ServerConfig {
     /// >0 = open metrics
     pub metrics_level: u32,
 
+    #[serde(default = "default::server::telemetry_enabled")]
+    pub telemetry_enabled: bool,
+
     #[serde(flatten)]
     pub unrecognized: HashMap<String, Value>,
 }
@@ -572,6 +575,10 @@ mod default {
         pub fn metrics_level() -> u32 {
             0
         }
+
+        pub fn telemetry_enabled() -> bool {
+            true
+        }
     }
 
     pub mod storage {

src/common/src/lib.rs

@@ -54,6 +54,8 @@ pub mod monitor;
 pub mod row;
 pub mod session_config;
 pub mod system_param;
+pub mod telemetry;
+
 #[cfg(test)]
 pub mod test_utils;
 pub mod types;

src/common/src/system_param/mod.rs

@@ -42,6 +42,7 @@ macro_rules! for_all_undeprecated_params {
             { data_directory, String, "hummock_001".to_string() },
             { backup_storage_url, String, "memory".to_string() },
             { backup_storage_directory, String, "backup".to_string() },
+            { telemetry_enabled, bool, true},
             $({ $field, $type, $default },)*
         }
     };
@@ -291,6 +292,10 @@ impl ValidateOnSet for OverrideValidateOnSet {
         // TODO
         Ok(())
     }
+
+    fn telemetry_enabled(_: &bool) -> Result<()> {
+        Ok(())
+    }
 }
 
 for_all_undeprecated_params!(impl_default_from_other_params);
@@ -315,6 +320,7 @@ mod tests {
             (DATA_DIRECTORY_KEY, "a"),
             (BACKUP_STORAGE_URL_KEY, "a"),
             (BACKUP_STORAGE_DIRECTORY_KEY, "a"),
+            (TELEMETRY_ENABLED_KEY, "false"),
         ];
 
         // To kv - missing field.

src/common/src/system_param/reader.rs

@@ -76,6 +76,10 @@ impl SystemParamsReader {
         self.prost.backup_storage_directory.as_ref().unwrap()
     }
 
+    pub fn telemetry_enabled(&self) -> bool {
+        self.prost.telemetry_enabled.unwrap()
+    }
+
     pub fn to_kv(&self) -> Vec<(String, String)> {
         system_params_to_kv(&self.prost).unwrap()
     }

src/common/src/telemetry/manager.rs

@@ -0,0 +1,163 @@
+// Copyright 2023 RisingWave Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::sync::atomic::{AtomicBool, Ordering};
+use std::sync::Arc;
+
+use parking_lot::RwLock;
+use tokio::select;
+use tokio::sync::oneshot::{self, Sender};
+use tokio::sync::watch::Receiver;
+use tokio::task::JoinHandle;
+
+use super::report::{start_telemetry_reporting, TelemetryInfoFetcher, TelemetryReportCreator};
+use crate::system_param::local_manager::SystemParamsReaderRef;
+use crate::telemetry::telemetry_env_enabled;
+
+pub struct TelemetryManager<F, I>
+where
+    F: TelemetryReportCreator + Send + Sync + 'static,
+    I: TelemetryInfoFetcher + Send + Sync + 'static,
+{
+    core: Arc<RwLock<TelemetryManagerCore<F, I>>>,
+    sys_params_change_rx: Receiver<SystemParamsReaderRef>,
+}
+
+impl<F, I> TelemetryManager<F, I>
+where
+    F: TelemetryReportCreator + Send + Sync + 'static,
+    I: TelemetryInfoFetcher + Send + Sync + 'static,
+{
+    pub fn new(
+        sys_params_change_rx: Receiver<SystemParamsReaderRef>,
+        info_fetcher: Arc<I>,
+        report_creator: Arc<F>,
+    ) -> Self {
+        Self {
+            core: Arc::new(RwLock::new(TelemetryManagerCore::new(
+                info_fetcher,
+                report_creator,
+            ))),
+            sys_params_change_rx,
+        }
+    }
+
+    pub fn start_telemetry_reporting(&self) {
+        self.core.write().start();
+    }
+
+    pub fn watch_params_change(mut self) -> (JoinHandle<()>, Sender<()>) {
+        let (shutdown_tx, mut shutdown_rx) = oneshot::channel();
+        let handle = tokio::spawn(async move {
+            loop {
+                select! {
+                    Ok(_) = self.sys_params_change_rx.changed() => {
+                                let telemetry_enabled = {
+                                    let params = self.sys_params_change_rx.borrow().load();
+                                    // check both environment variable and system params
+                                    // if either is false, then stop telemetry
+                                    params.telemetry_enabled() && telemetry_env_enabled()
+                                };
+
+                                let telemetry_running = {
+                                    let core = self.core.read();
+                                    core.telemetry_running()
+                                };
+
+                                match (telemetry_running, telemetry_enabled) {
+                                    (false, true) => {
+                                        tracing::info!("telemetry config changed to true, start reporting");
+                                        self.core.write().start();
+                                    }
+                                    (true, false) => {
+                                        tracing::info!("telemetry config changed to false, stop reporting");
+                                        self.core.write().stop();
+                                    }
+                                    _ => {}
+                                };
+                        }
+                    ,
+                    _ = &mut shutdown_rx =>{
+                        tracing::info!("Telemetry exit");
+                        return;
+                    }
+                }
+            }
+        });
+        (handle, shutdown_tx)
+    }
+}
+
+struct TelemetryManagerCore<F, I>
+where
+    F: TelemetryReportCreator + Send + Sync + 'static,
+    I: TelemetryInfoFetcher + Send + Sync + 'static,
+{
+    telemetry_handle: Option<JoinHandle<()>>,
+    telemetry_shutdown_tx: Option<Sender<()>>,
+    telemetry_running: Arc<AtomicBool>,
+    info_fetcher: Arc<I>,
+    report_creator: Arc<F>,
+}
+
+impl<F, I> TelemetryManagerCore<F, I>
+where
+    F: TelemetryReportCreator + Send + Sync + 'static,
+    I: TelemetryInfoFetcher + Send + Sync + 'static,
+{
+    fn new(info_fetcher: Arc<I>, report_creator: Arc<F>) -> Self {
+        Self {
+            telemetry_handle: None,
+            telemetry_shutdown_tx: None,
+            telemetry_running: Arc::new(AtomicBool::new(false)),
+            info_fetcher,
+            report_creator,
+        }
+    }
+
+    fn telemetry_running(&self) -> bool {
+        self.telemetry_running.load(Ordering::Relaxed)
+    }
+
+    fn start(&mut self) {
+        if self.telemetry_running() {
+            return;
+        }
+
+        let (handle, tx) =
+            start_telemetry_reporting(self.info_fetcher.clone(), self.report_creator.clone());
+        self.telemetry_handle = Some(handle);
+        self.telemetry_shutdown_tx = Some(tx);
+        self.telemetry_running.store(true, Ordering::Relaxed);
+    }
+
+    fn stop(&mut self) {
+        match (
+            self.telemetry_running.load(Ordering::Relaxed),
+            self.telemetry_shutdown_tx.take(),
+            self.telemetry_handle.take(),
+        ) {
+            (true, Some(shutdown_rx), Some(_)) => {
+                if let Err(()) = shutdown_rx.send(()) {
+                    tracing::error!("telemetry mgr failed to send stop signal");
+                } else {
+                    self.telemetry_running.store(false, Ordering::Relaxed)
+                }
+            }
+            // do nothing if telemetry is not running
+            (false, None, None) => {}
+            _ => unreachable!("impossible telemetry handler"),
+        }
+    }
+}