feat(rpc): support profiling task

Author: saying121

src/main.rs

@@ -32,11 +32,16 @@ use log::log_init;
 use mimalloc::MiMalloc;
 use nix::unistd::geteuid;
 use opentelemetry_otlp::ExportConfig;
-use psh_proto::HeartbeatReq;
-use runtime::{Task, TaskRuntime};
+use psh_proto::{
+    HeartbeatReq, PerfDataProto,
+    export_data_req::{Data, data::DataType},
+};
+use runtime::{TaskRuntime, WasmTask};
 use services::rpc::RpcClient;
 use tokio::{runtime::Runtime, try_join};
 
+use self::services::{rpc::WhichTask, sampling::Profiler};
+
 #[global_allocator]
 static GLOBAL: MiMalloc = mimalloc::MiMalloc;
 
@@ -86,7 +91,7 @@ fn main() -> Result<()> {
     let task_rt = TaskRuntime::new()?;
 
     if let Some(args) = wasm_with_args {
-        let task = Task {
+        let task = WasmTask {
             id: None,
             wasm_component: fs::read(&args[0])?,
             wasm_component_args: args,
@@ -143,14 +148,52 @@ async fn async_tasks(remote_cfg: RemoteConfig, mut task_rt: TaskRuntime) -> Resu
         loop {
             let idle = task_rt.is_idle();
             if idle {
-                if let Some(mut task) = client.get_task(instance_id.clone()).await? {
-                    let task_id = task
-                        .id
-                        .as_ref()
-                        .map(|it| it.to_string())
-                        .expect("No task id provided");
-                    task.wasm_component_args.insert(0, task_id);
-                    task_rt.schedule(task)?
+                if let Some(task) = client.get_task(instance_id.clone()).await? {
+                    match task {
+                        WhichTask::Wasm(mut task) => {
+                            let task_id = task
+                                .id
+                                .as_ref()
+                                .map(|it| it.to_string())
+                                .expect("No task id provided");
+                            task.wasm_component_args.insert(0, task_id);
+                            task_rt.schedule(task)?;
+                        }
+                        WhichTask::Profiling(profiling_task) => {
+                            let mut profiler = Profiler::new(
+                                profiling_task.process,
+                                profiling_task.mmap_pages as _,
+                                profiling_task.overflow_by,
+                                profiling_task.stack_depth,
+                            )?;
+                            let task_time_slice = {
+                                let delta = profiling_task.end_time.timestamp_millis()
+                                    - Utc::now().timestamp_millis();
+                                delta.max(0) as u64
+                            };
+
+                            let perf_data = tokio::task::spawn_blocking(
+                                move || -> anyhow::Result<PerfDataProto> {
+                                    profiler.enable()?;
+                                    std::thread::sleep(Duration::from_millis(task_time_slice));
+                                    profiler.disable()?;
+                                    Ok(profiler.perf_data_proto())
+                                },
+                            )
+                            .await??;
+                            if let Some(task_id) = profiling_task.id {
+                                let data = Data {
+                                    data_type: Some(DataType::PerfData(perf_data)),
+                                };
+                                client
+                                    .export_data(psh_proto::ExportDataReq {
+                                        task_id,
+                                        data: vec![data],
+                                    })
+                                    .await?;
+                            }
+                        }
+                    };
                 }
             }
 

src/runtime/data_export.rs

@@ -26,7 +26,10 @@ use profiling::data_export::{
     common::FieldValue as WitFieldValue, measurement::Point, metric::Sample,
 };
 use prost::Message;
-use psh_proto::{Data, DataType, ExportDataReq};
+use psh_proto::{
+    ExportDataReq, LineProtocolData,
+    export_data_req::{Data, data::DataType},
+};
 use wasmtime::component::Linker;
 
 use crate::{TOKIO_RUNTIME, services::rpc::RpcClient};
@@ -172,8 +175,7 @@ impl profiling::data_export::file::Host for DataExportCtx {
         };
 
         let data = Data {
-            ty: DataType::File as _,
-            bytes,
+            data_type: Some(DataType::File(psh_proto::FileData { bytes })),
         };
         ctx.exporter.schedule(data);
 
@@ -201,8 +203,7 @@ impl profiling::data_export::metric::Host for DataExportCtx {
         };
 
         let data = Data {
-            ty: DataType::LineProtocol as _,
-            bytes,
+            data_type: Some(DataType::LineProtocol(LineProtocolData { bytes })),
         };
         ctx.exporter.schedule(data);
 
@@ -238,8 +239,7 @@ impl profiling::data_export::measurement::Host for DataExportCtx {
         };
 
         let data = Data {
-            ty: DataType::LineProtocol as _,
-            bytes,
+            data_type: Some(DataType::LineProtocol(LineProtocolData { bytes })),
         };
         ctx.exporter.schedule(data);
 

src/runtime/mod.rs

@@ -39,16 +39,17 @@ pub use state::PshState;
 
 use crate::services::rpc::RpcClient;
 
-pub struct Task {
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord)]
+pub struct WasmTask {
     pub id: Option<String>,
     pub wasm_component: Vec<u8>,
     pub wasm_component_args: Vec<String>,
     pub end_time: DateTime<Utc>,
 }
 
 pub struct TaskRuntime {
-    tx: Sender<Task>,
-    rx: Option<Receiver<Task>>,
+    tx: Sender<WasmTask>,
+    rx: Option<Receiver<WasmTask>>,
     len: Arc<AtomicUsize>,
     finished_task_id: Arc<Mutex<Vec<String>>>,
 }
@@ -65,7 +66,7 @@ impl TaskRuntime {
         })
     }
 
-    pub fn schedule(&self, task: Task) -> Result<()> {
+    pub fn schedule(&self, task: WasmTask) -> Result<()> {
         self.len.fetch_add(1, Ordering::Release);
         self.tx.send(task)?;
         Ok(())

src/services/rpc.rs

@@ -13,7 +13,10 @@
 // see <https://www.gnu.org/licenses/>.
 
 use anyhow::{Result, bail};
+use chrono::DateTime;
 use chrono::{TimeZone, Utc, offset::LocalResult};
+use perf_event_rs::sampling::OverflowBy;
+use psh_proto::task::TaskType;
 use psh_proto::{
     ExportDataReq, GetTaskReq, HeartbeatReq, TaskDoneReq, Unit,
     psh_service_client::PshServiceClient,
@@ -26,7 +29,7 @@ use tonic::{
     transport::{Channel, ClientTlsConfig, Endpoint},
 };
 
-use crate::{config::RpcConfig, runtime::Task, services::host_info::new_info_req};
+use crate::{config::RpcConfig, runtime::WasmTask, services::host_info::new_info_req};
 
 #[derive(Clone)]
 pub struct RpcClient {
@@ -80,6 +83,20 @@ where
     }
 }
 
+pub enum WhichTask {
+    Wasm(WasmTask),
+    Profiling(ProfilingTask),
+}
+
+pub struct ProfilingTask {
+    pub id: Option<String>,
+    pub process: perf_event_rs::config::Process,
+    pub mmap_pages: u64,
+    pub overflow_by: OverflowBy,
+    pub stack_depth: Option<u16>,
+    pub end_time: DateTime<Utc>,
+}
+
 impl RpcClient {
     pub async fn new(config: &RpcConfig, token: String) -> Result<Self> {
         let ep = Endpoint::from_shared(config.addr.clone())?
@@ -132,7 +149,7 @@ impl RpcClient {
         Ok(())
     }
 
-    pub async fn get_task(&mut self, instance_id: String) -> Result<Option<Task>> {
+    pub async fn get_task(&mut self, instance_id: String) -> Result<Option<WhichTask>> {
         let get_task_req = GetTaskReq { instance_id };
         let token = &self.token;
 
@@ -142,20 +159,41 @@ impl RpcClient {
             self.client.get_task(req).await
         })
         .await?;
-        let task = match response.into_inner().task {
-            Some(task) => task,
-            None => return Ok(None),
+        let Some(task): Option<psh_proto::Task> = response.into_inner().task else {
+            return Ok(None);
         };
 
-        let end_time = match Utc.timestamp_millis_opt(task.end_time as _) {
-            LocalResult::Single(t) => t,
-            _ => bail!("Invalid task end time"),
+        let LocalResult::Single(end_time) = Utc.timestamp_millis_opt(task.end_time as _) else {
+            bail!("Invalid task end time")
         };
-        let task = Task {
-            id: Some(task.id),
-            wasm_component: task.wasm,
-            wasm_component_args: task.wasm_args,
-            end_time,
+        let Some(task_type) = task.task_type else {
+            return Ok(None);
+        };
+
+        let task = match task_type {
+            TaskType::Profiling(profiling_task) => {
+                let Some(process) = profiling_task.process else {
+                    return Ok(None);
+                };
+                let Some(overflow_by) = profiling_task.overflow_by else {
+                    return Ok(None);
+                };
+                let process: perf_event_rs::config::Process = process.into();
+                WhichTask::Profiling(ProfilingTask {
+                    id: task.id.into(),
+                    process,
+                    mmap_pages: profiling_task.mmap_pages,
+                    overflow_by: overflow_by.into(),
+                    stack_depth: profiling_task.stack_depth.map(|v| v as _),
+                    end_time,
+                })
+            }
+            TaskType::Wasm(wasm_task) => WhichTask::Wasm(WasmTask {
+                id: Some(task.id),
+                wasm_component: wasm_task.wasm,
+                wasm_component_args: wasm_task.wasm_args,
+                end_time,
+            }),
         };
 
         Ok(Some(task))