OptimatistOpenSource · saying121 · Sep 25, 2025 · Sep 25, 2025 · Sep 28, 2025 · BrytonLee
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -19,7 +19,7 @@ chrono = { workspace = true }
 clap = { workspace = true, features = ["derive", "wrap_help"] }
 tonic = { workspace = true }
 prost = { workspace = true }
-tokio = { workspace = true, features = ["macros", "rt-multi-thread"] }
+tokio = { workspace = true, features = ["macros", "rt-multi-thread", "fs"] }
 nix = { workspace = true, features = ["user", "hostname"] }
 wasmtime = { workspace = true }
 wasmtime-wasi = { workspace = true }
@@ -50,6 +50,8 @@ influxdb-line-protocol = { workspace = true }
 psh-proto = { workspace = true }
 mimalloc = { workspace = true }
 nvml-wrapper = { workspace = true }
+perf-event-rs = { workspace = true }
+num_cpus = { workspace = true }
 
 [lints]
 workspace = true
@@ -58,7 +60,7 @@ workspace = true
 host-op-perf = { path = "crates/op/host-op-perf" }
 host-op-system = { path = "crates/op/host-op-system" }
 psh-system = { path = "crates/psh-system" }
-perf-event-rs = { git = "https://github.com/OptimatistOpenSource/perf-event-rs.git", rev = "423ca26f53b27193d2321028dae5fd362a9673e9" }
+perf-event-rs = { git = "https://github.com/OptimatistOpenSource/perf-event-rs.git", rev = "d6881f34b8a9cde1d70dab5fb1415271e6b0bb25" }
 tokio = "^1"
 libc = "^0.2"
 chrono = "^0.4"
@@ -87,7 +89,7 @@ local-ip-address = "^0.6"
 TinyUFO = "0.4"
 crossbeam = "0.8"
 influxdb-line-protocol = "2"
-psh-proto = { git = "https://github.com/OptimatistOpenSource/psh-proto.git", rev = "ca2919053029cb584b478611f8bf8496bf3cf7f7" }
+psh-proto = { git = "https://github.com/OptimatistOpenSource/psh-proto.git", rev = "c82356c52925ac0abf9cd93c7245c3518d78e096" }
 mimalloc = "0.1"
 nvml-wrapper = "0.10.0"
 

diff --git a/src/main.rs b/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,70 @@ 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 mut data = vec![];
+                                for ele in &perf_data.events {
+                                    let Some(event_type) = &ele.event_type else {
+                                        continue;
+                                    };
+
+                                    if let psh_proto::perf_data_proto::perf_event::EventType::MmapEvent(event)= event_type{
+                                        let Some(filename) = &event.filename else {
+                                            continue;
+                                        };
+
+                                        let data_type = DataType::ElfFile(psh_proto::ElfFile {
+                                            filename: filename.to_owned(),
+                                            build_id: event.build_id.clone(),
+                                            arch: std::env::consts::ARCH.to_string(),
+                                            bytes: tokio::fs::read(filename).await?
+                                        });
+                                        data.push(Data { data_type: Some(data_type) });
+                                    }
+                                }
+                                let dat = Data {
+                                    data_type: Some(DataType::PerfData(perf_data)),
+                                };
+                                data.push(dat);
+                                client
+                                    .export_data(psh_proto::ExportDataReq { task_id, data })
+                                    .await?;
+                            }
+                        }
+                    };
                 }
             }
 

diff --git a/src/runtime/data_export.rs b/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);
 

diff --git a/src/runtime/mod.rs b/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(())

diff --git a/src/services/mod.rs b/src/services/mod.rs
@@ -14,3 +14,4 @@
 
 pub mod host_info;
 pub mod rpc;
+pub mod sampling;
diff --git a/src/services/rpc.rs b/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))
Original file line number	Diff line number	Diff line change
Expand Up		@@ -14,3 +14,4 @@

		pub mod host_info;
		pub mod rpc;
		pub mod sampling;