1use async_trait::async_trait;
14use hyperactor::Actor;
15use hyperactor::Context;
16use hyperactor::Endpoint as _;
17use hyperactor::HandleClient;
18use hyperactor::Handler;
19use hyperactor::OncePortRef;
20use hyperactor::RefClient;
21use serde::Deserialize;
22use serde::Serialize;
23use typeuri::Named;
24
25use crate::config::MESH_ADMIN_PYSPY_TIMEOUT;
26use crate::config::PYSPY_BIN;
27
28#[derive(
32 Debug,
33 Clone,
34 PartialEq,
35 Serialize,
36 Deserialize,
37 Named,
38 schemars::JsonSchema
39)]
40pub enum PySpyResult {
41 Ok {
43 pid: u32,
45 binary: String,
47 stack_traces: Vec<PySpyStackTrace>,
49 warnings: Vec<String>,
53 },
54 BinaryNotFound {
56 searched: Vec<String>,
58 },
59 Failed {
61 pid: u32,
63 binary: String,
65 exit_code: Option<i32>,
67 stderr: String,
69 },
70}
71wirevalue::register_type!(PySpyResult);
72
73#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
75pub struct PySpyStackTrace {
76 pub pid: i32,
78 pub thread_id: u64,
80 pub thread_name: Option<String>,
82 pub os_thread_id: Option<u64>,
84 pub active: bool,
86 pub owns_gil: bool,
88 pub frames: Vec<PySpyFrame>,
90}
91
92#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
94pub struct PySpyFrame {
95 pub name: String,
97 pub filename: String,
99 pub module: Option<String>,
101 pub short_filename: Option<String>,
103 pub line: i32,
105 pub locals: Option<Vec<PySpyLocalVariable>>,
107 pub is_entry: bool,
109}
110
111#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, schemars::JsonSchema)]
113pub struct PySpyLocalVariable {
114 pub name: String,
116 pub addr: usize,
118 pub arg: bool,
120 pub repr: Option<String>,
122}
123
124#[derive(Debug, Clone, Serialize, Deserialize)]
126pub struct PySpyOpts {
127 pub threads: bool,
129 pub native: bool,
131 pub native_all: bool,
134 pub nonblocking: bool,
137}
138
139#[derive(Debug, Clone, Serialize, Deserialize, schemars::JsonSchema)]
146pub struct PySpyProfileOpts {
147 #[schemars(range(min = 1))]
151 pub duration_s: u32,
152 #[schemars(range(min = 1, max = 1000))]
154 pub rate_hz: u32,
155 pub native: bool,
157 pub threads: bool,
159 pub nonblocking: bool,
161}
162
163#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
165pub(crate) struct ProfileDurationSecs(std::num::NonZeroU32);
166
167impl ProfileDurationSecs {
168 pub fn get(self) -> u32 {
169 self.0.get()
170 }
171}
172
173#[derive(Debug, Clone, Copy, Serialize, Deserialize)]
175pub(crate) struct SampleRateHz(std::num::NonZeroU32);
176
177impl SampleRateHz {
178 pub fn get(self) -> u32 {
179 self.0.get()
180 }
181}
182
183#[derive(Debug, Clone, Serialize, Deserialize)]
186pub(crate) struct ValidatedProfileRequest {
187 duration: ProfileDurationSecs,
189 rate: SampleRateHz,
191 native: bool,
193 threads: bool,
195 nonblocking: bool,
197 subprocess_timeout: std::time::Duration,
199 bridge_timeout: std::time::Duration,
201}
202
203impl ValidatedProfileRequest {
204 pub fn duration(&self) -> ProfileDurationSecs {
205 self.duration
206 }
207 pub fn rate(&self) -> SampleRateHz {
208 self.rate
209 }
210 pub fn native(&self) -> bool {
211 self.native
212 }
213 pub fn threads(&self) -> bool {
214 self.threads
215 }
216 pub fn nonblocking(&self) -> bool {
217 self.nonblocking
218 }
219 pub fn subprocess_timeout(&self) -> std::time::Duration {
220 self.subprocess_timeout
221 }
222 pub fn bridge_timeout(&self) -> std::time::Duration {
223 self.bridge_timeout
224 }
225
226 pub fn try_new(
227 opts: &PySpyProfileOpts,
228 max_duration: std::time::Duration,
229 ) -> Result<Self, String> {
230 let duration = std::num::NonZeroU32::new(opts.duration_s)
231 .map(ProfileDurationSecs)
232 .ok_or_else(|| "duration_s must be positive".to_string())?;
233 if std::time::Duration::from_secs(u64::from(duration.get())) > max_duration {
234 return Err(format!(
235 "duration_s {}s exceeds max {}s",
236 duration.get(),
237 max_duration.as_secs()
238 ));
239 }
240 let rate = std::num::NonZeroU32::new(opts.rate_hz)
241 .filter(|n| n.get() <= 1000)
242 .map(SampleRateHz)
243 .ok_or_else(|| format!("rate_hz must be 1..=1000, got {}", opts.rate_hz))?;
244 let subprocess_timeout = std::time::Duration::from_secs(u64::from(duration.get()) + 15);
245 let bridge_timeout = subprocess_timeout + std::time::Duration::from_secs(5);
246 Ok(Self {
247 duration,
248 rate,
249 native: opts.native,
250 threads: opts.threads,
251 nonblocking: opts.nonblocking,
252 subprocess_timeout,
253 bridge_timeout,
254 })
255 }
256}
257
258#[derive(Debug, Clone, Serialize, Deserialize, Named)]
262pub enum PySpyProfileResult {
263 Ok {
264 pid: u32,
265 binary: String,
266 svg: Vec<u8>,
267 },
268 BinaryNotFound {
269 searched: Vec<String>,
270 },
271 TimedOut {
272 pid: u32,
273 binary: String,
274 timeout_s: u64,
275 stderr: String,
276 },
277 ExitFailure {
278 pid: u32,
279 binary: String,
280 exit_code: Option<i32>,
281 stderr: String,
282 },
283 OutputMissing {
284 pid: u32,
285 binary: String,
286 },
287 OutputEmpty {
288 pid: u32,
289 binary: String,
290 },
291 OutputReadFailure {
292 pid: u32,
293 binary: String,
294 error: String,
295 },
296 WorkerSpawnFailure {
297 error: String,
298 },
299 SubprocessSpawnFailure {
300 pid: u32,
301 binary: String,
302 error: String,
303 },
304 WaitFailure {
305 pid: u32,
306 binary: String,
307 error: String,
308 },
309 TempDirFailure {
310 pid: u32,
311 binary: String,
312 error: String,
313 },
314}
315wirevalue::register_type!(PySpyProfileResult);
316
317#[derive(Debug)]
320pub(crate) enum ProfileExecOutcome {
321 Ok {
322 pid: u32,
323 binary: String,
324 svg: Vec<u8>,
325 },
326 BinaryNotFound {
327 searched: Vec<String>,
328 },
329 TimedOut {
330 pid: u32,
331 binary: String,
332 timeout: std::time::Duration,
333 stderr: String,
334 },
335 ExitFailure {
336 pid: u32,
337 binary: String,
338 exit_code: Option<i32>,
339 stderr: String,
340 },
341 OutputMissing {
342 pid: u32,
343 binary: String,
344 },
345 OutputEmpty {
346 pid: u32,
347 binary: String,
348 },
349 OutputReadFailure {
350 pid: u32,
351 binary: String,
352 error: String,
353 },
354 SubprocessSpawnFailure {
355 pid: u32,
356 binary: String,
357 error: String,
358 },
359 WaitFailure {
360 pid: u32,
361 binary: String,
362 error: String,
363 },
364 TempDirFailure {
365 pid: u32,
366 binary: String,
367 error: String,
368 },
369}
370
371impl From<ProfileExecOutcome> for PySpyProfileResult {
372 fn from(outcome: ProfileExecOutcome) -> Self {
373 match outcome {
374 ProfileExecOutcome::Ok { pid, binary, svg } => {
375 PySpyProfileResult::Ok { pid, binary, svg }
376 }
377 ProfileExecOutcome::BinaryNotFound { searched } => {
378 PySpyProfileResult::BinaryNotFound { searched }
379 }
380 ProfileExecOutcome::TimedOut {
381 pid,
382 binary,
383 timeout,
384 stderr,
385 } => PySpyProfileResult::TimedOut {
386 pid,
387 binary,
388 timeout_s: timeout.as_secs(),
389 stderr,
390 },
391 ProfileExecOutcome::ExitFailure {
392 pid,
393 binary,
394 exit_code,
395 stderr,
396 } => PySpyProfileResult::ExitFailure {
397 pid,
398 binary,
399 exit_code,
400 stderr,
401 },
402 ProfileExecOutcome::OutputMissing { pid, binary } => {
403 PySpyProfileResult::OutputMissing { pid, binary }
404 }
405 ProfileExecOutcome::OutputEmpty { pid, binary } => {
406 PySpyProfileResult::OutputEmpty { pid, binary }
407 }
408 ProfileExecOutcome::OutputReadFailure { pid, binary, error } => {
409 PySpyProfileResult::OutputReadFailure { pid, binary, error }
410 }
411 ProfileExecOutcome::SubprocessSpawnFailure { pid, binary, error } => {
412 PySpyProfileResult::SubprocessSpawnFailure { pid, binary, error }
413 }
414 ProfileExecOutcome::WaitFailure { pid, binary, error } => {
415 PySpyProfileResult::WaitFailure { pid, binary, error }
416 }
417 ProfileExecOutcome::TempDirFailure { pid, binary, error } => {
418 PySpyProfileResult::TempDirFailure { pid, binary, error }
419 }
420 }
421 }
422}
423
424#[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
432pub struct PySpyDump {
433 pub opts: PySpyOpts,
435 #[reply]
437 pub result: OncePortRef<PySpyResult>,
438}
439wirevalue::register_type!(PySpyDump);
440
441#[allow(private_interfaces)] #[derive(Debug, Serialize, Deserialize, Named, Handler, HandleClient, RefClient)]
449pub struct PySpyProfile {
450 pub request: ValidatedProfileRequest,
452 #[reply]
454 pub result: OncePortRef<PySpyProfileResult>,
455}
456wirevalue::register_type!(PySpyProfile);
457
458pub struct PySpyRunner;
462
463impl PySpyRunner {
464 pub async fn dump_self(&self, opts: &PySpyOpts) -> PySpyResult {
474 let pid = std::process::id();
475 let pyspy_bin: String = hyperactor_config::global::get_cloned(PYSPY_BIN);
476 let candidates = resolve_candidates(if pyspy_bin.is_empty() {
477 None
478 } else {
479 Some(pyspy_bin)
480 });
481 let mut searched = vec![];
482
483 for (binary, label) in &candidates {
484 searched.push(label.clone());
485 if let Some(result) = try_exec(
486 binary,
487 pid,
488 opts,
489 hyperactor_config::global::get(MESH_ADMIN_PYSPY_TIMEOUT),
490 )
491 .await
492 {
493 return result;
494 }
495 }
496
497 PySpyResult::BinaryNotFound { searched }
498 }
499
500 pub(crate) async fn profile_self(
503 &self,
504 request: &ValidatedProfileRequest,
505 ) -> ProfileExecOutcome {
506 let pid = std::process::id();
507 let pyspy_bin: String = hyperactor_config::global::get_cloned(PYSPY_BIN);
508 let candidates = resolve_candidates(if pyspy_bin.is_empty() {
509 None
510 } else {
511 Some(pyspy_bin)
512 });
513 let mut searched = vec![];
514
515 for (binary, label) in &candidates {
516 searched.push(label.clone());
517 if let Some(result) = try_profile(binary, pid, request).await {
518 return result;
519 }
520 }
521
522 ProfileExecOutcome::BinaryNotFound { searched }
523 }
524}
525
526#[derive(Debug, Serialize, Deserialize, Named)]
530pub struct RunPySpyDump {
531 pub opts: PySpyOpts,
532 pub reply_port: hyperactor::OncePortRef<PySpyResult>,
534}
535wirevalue::register_type!(RunPySpyDump);
536
537#[hyperactor::export(handlers = [RunPySpyDump])]
542pub struct PySpyWorker;
543
544impl Actor for PySpyWorker {}
545
546impl PySpyWorker {
547 pub(crate) fn spawn_and_forward(
550 cx: &impl hyperactor::context::Actor,
551 opts: PySpyOpts,
552 reply_port: hyperactor::OncePortRef<PySpyResult>,
553 ) -> Result<(), anyhow::Error> {
554 let worker = cx.spawn(Self);
555 worker.post(cx, RunPySpyDump { opts, reply_port });
556 Ok(())
557 }
558}
559
560#[async_trait]
561impl Handler<RunPySpyDump> for PySpyWorker {
562 async fn handle(
563 &mut self,
564 cx: &Context<Self>,
565 message: RunPySpyDump,
566 ) -> Result<(), anyhow::Error> {
567 let result = PySpyRunner.dump_self(&message.opts).await;
568 message.reply_port.post(cx, result);
569 cx.stop("pyspy dump complete")?;
570 Ok(())
571 }
572}
573
574#[allow(private_interfaces)] #[derive(Debug, Serialize, Deserialize, Named)]
577pub struct RunPySpyProfile {
578 pub request: ValidatedProfileRequest,
579 pub reply_port: hyperactor::OncePortRef<PySpyProfileResult>,
580}
581wirevalue::register_type!(RunPySpyProfile);
582
583#[hyperactor::export(handlers = [RunPySpyProfile])]
586pub struct PySpyProfileWorker;
587
588impl Actor for PySpyProfileWorker {}
589
590impl PySpyProfileWorker {
591 pub(crate) fn spawn_and_forward(
593 cx: &impl hyperactor::context::Actor,
594 request: ValidatedProfileRequest,
595 reply_port: hyperactor::OncePortRef<PySpyProfileResult>,
596 ) -> Result<(), anyhow::Error> {
597 let worker = cx.spawn(Self);
598 worker.post(
599 cx,
600 RunPySpyProfile {
601 request,
602 reply_port,
603 },
604 );
605 Ok(())
606 }
607}
608
609#[async_trait]
610impl Handler<RunPySpyProfile> for PySpyProfileWorker {
611 async fn handle(
612 &mut self,
613 cx: &Context<Self>,
614 message: RunPySpyProfile,
615 ) -> Result<(), anyhow::Error> {
616 let outcome = PySpyRunner.profile_self(&message.request).await;
617 message
618 .reply_port
619 .post(cx, PySpyProfileResult::from(outcome));
620 cx.stop("pyspy profile complete")?;
621 Ok(())
622 }
623}
624
625fn resolve_candidates(pyspy_bin_env: Option<String>) -> Vec<(String, String)> {
628 let mut candidates = vec![];
629 if let Some(path) = pyspy_bin_env
630 && !path.is_empty()
631 {
632 let label = format!("PYSPY_BIN={}", path);
633 candidates.push((path, label));
634 }
635 candidates.push(("py-spy".to_string(), "py-spy on PATH".to_string()));
636 candidates
637}
638
639fn build_command(binary: &str, pid: u32, opts: &PySpyOpts) -> tokio::process::Command {
641 let mut cmd = tokio::process::Command::new(binary);
642 cmd.arg("dump")
643 .arg("--pid")
644 .arg(pid.to_string())
645 .arg("--json");
646 if opts.threads {
647 cmd.arg("--threads");
648 }
649 if opts.native {
650 cmd.arg("--native");
651 }
652 if opts.native_all {
653 cmd.arg("--native-all");
654 }
655 if opts.nonblocking {
656 cmd.arg("--nonblocking");
657 }
658 cmd.stdout(std::process::Stdio::piped());
659 cmd.stderr(std::process::Stdio::piped());
660 cmd
661}
662
663fn map_output(output: std::process::Output, pid: u32, binary: &str) -> PySpyResult {
667 if output.status.success() {
668 match serde_json::from_slice::<Vec<PySpyStackTrace>>(&output.stdout) {
669 Ok(stack_traces) => PySpyResult::Ok {
670 pid,
671 binary: binary.to_string(),
672 stack_traces,
673 warnings: vec![],
674 },
675 Err(e) => PySpyResult::Failed {
676 pid,
677 binary: binary.to_string(),
678 exit_code: None,
679 stderr: format!("failed to parse py-spy JSON output: {}", e),
680 },
681 }
682 } else {
683 let stderr = String::from_utf8_lossy(&output.stderr).into_owned();
684 PySpyResult::Failed {
685 pid,
686 binary: binary.to_string(),
687 exit_code: output.status.code(),
688 stderr,
689 }
690 }
691}
692
693fn is_unsupported_native_all(result: &PySpyResult) -> bool {
697 matches!(
698 result,
699 PySpyResult::Failed {
700 exit_code: Some(2),
701 stderr,
702 ..
703 } if stderr.contains("--native-all")
704 )
705}
706
707enum ExecOnce {
709 Result(PySpyResult),
711 NotFound,
713}
714
715async fn exec_once(
720 binary: &str,
721 pid: u32,
722 opts: &PySpyOpts,
723 deadline: tokio::time::Instant,
724 timeout: std::time::Duration,
725) -> ExecOnce {
726 let remaining = deadline.saturating_duration_since(tokio::time::Instant::now());
727 if remaining.is_zero() {
728 return ExecOnce::Result(PySpyResult::Failed {
729 pid,
730 binary: binary.to_string(),
731 exit_code: None,
732 stderr: format!("py-spy subprocess timed out after {}s", timeout.as_secs()),
733 });
734 }
735 let child = match build_command(binary, pid, opts).spawn() {
736 Ok(child) => child,
737 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return ExecOnce::NotFound,
738 Err(e) => {
739 return ExecOnce::Result(PySpyResult::Failed {
740 pid,
741 binary: binary.to_string(),
742 exit_code: None,
743 stderr: format!("failed to execute: {}", e),
744 });
745 }
746 };
747 ExecOnce::Result(collect_with_timeout(child, pid, binary, remaining).await)
748}
749
750async fn try_exec(
764 binary: &str,
765 pid: u32,
766 opts: &PySpyOpts,
767 timeout: std::time::Duration,
768) -> Option<PySpyResult> {
769 let deadline = tokio::time::Instant::now() + timeout;
770 let retries = if opts.nonblocking { 3 } else { 1 };
771 let mut last_result = None;
772 let mut effective_opts = opts.clone();
773
774 for attempt in 0..retries {
775 if attempt > 0 {
776 tokio::time::sleep(std::time::Duration::from_millis(100)).await;
777 }
778 let mut result = match exec_once(binary, pid, &effective_opts, deadline, timeout).await {
779 ExecOnce::NotFound => return None,
780 ExecOnce::Result(r) => r,
781 };
782 if is_unsupported_native_all(&result) && effective_opts.native_all {
786 effective_opts.native_all = false;
789 result = match exec_once(binary, pid, &effective_opts, deadline, timeout).await {
790 ExecOnce::NotFound => return None,
791 ExecOnce::Result(r) => r,
792 };
793 if let PySpyResult::Ok { warnings, .. } = &mut result {
795 warnings.push(
796 "--native-all unsupported by this py-spy; fell back to --native".to_string(),
797 );
798 }
799 }
802 match &result {
803 PySpyResult::Ok { .. } => return Some(result),
804 _ => {
805 last_result = Some(result);
806 }
807 }
808 }
809
810 last_result
811}
812
813async fn collect_with_timeout(
822 mut child: tokio::process::Child,
823 pid: u32,
824 binary: &str,
825 timeout: std::time::Duration,
826) -> PySpyResult {
827 let mut stdout_handle = child.stdout.take();
828 let mut stderr_handle = child.stderr.take();
829
830 let collect = async {
831 let stdout_fut = async {
832 let mut buf = Vec::new();
833 if let Some(ref mut r) = stdout_handle {
834 let _ = tokio::io::AsyncReadExt::read_to_end(r, &mut buf).await;
835 }
836 buf
837 };
838 let stderr_fut = async {
839 let mut buf = Vec::new();
840 if let Some(ref mut r) = stderr_handle {
841 let _ = tokio::io::AsyncReadExt::read_to_end(r, &mut buf).await;
842 }
843 buf
844 };
845 let (stdout_bytes, stderr_bytes, status) =
846 tokio::join!(stdout_fut, stderr_fut, child.wait());
847 (stdout_bytes, stderr_bytes, status)
848 };
849
850 match tokio::time::timeout(timeout, collect).await {
851 Ok((stdout_bytes, stderr_bytes, Ok(status))) => {
852 let output = std::process::Output {
853 status,
854 stdout: stdout_bytes,
855 stderr: stderr_bytes,
856 };
857 map_output(output, pid, binary)
858 }
859 Ok((_, _, Err(e))) => PySpyResult::Failed {
860 pid,
861 binary: binary.to_string(),
862 exit_code: None,
863 stderr: format!("failed to wait for child: {}", e),
864 },
865 Err(_) => {
866 let _ = child.start_kill();
868 let _ = child.wait().await;
869 PySpyResult::Failed {
870 pid,
871 binary: binary.to_string(),
872 exit_code: None,
873 stderr: format!("py-spy subprocess timed out after {}s", timeout.as_secs()),
874 }
875 }
876 }
877}
878
879fn build_record_command(
881 binary: &str,
882 pid: u32,
883 request: &ValidatedProfileRequest,
884 output_path: &std::path::Path,
885) -> tokio::process::Command {
886 let mut cmd = tokio::process::Command::new(binary);
887 cmd.arg("record")
888 .arg("--pid")
889 .arg(pid.to_string())
890 .arg("--duration")
891 .arg(request.duration().get().to_string())
892 .arg("--rate")
893 .arg(request.rate().get().to_string())
894 .arg("--format")
895 .arg("flamegraph")
896 .arg("--output")
897 .arg(output_path);
898 if request.native() {
899 cmd.arg("--native");
900 }
901 if request.threads() {
902 cmd.arg("--threads");
903 }
904 if request.nonblocking() {
905 cmd.arg("--nonblocking");
906 }
907 cmd.stdout(std::process::Stdio::null());
910 cmd.stderr(std::process::Stdio::piped());
911 cmd
912}
913
914async fn collect_profile_with_timeout(
917 mut child: tokio::process::Child,
918 pid: u32,
919 binary: &str,
920 timeout: std::time::Duration,
921) -> Result<(std::process::ExitStatus, String), ProfileExecOutcome> {
922 let stderr_handle = child.stderr.take();
926 let stderr_task = tokio::spawn(async move {
927 let mut buf = Vec::new();
928 if let Some(mut r) = stderr_handle {
929 let _ = tokio::io::AsyncReadExt::read_to_end(&mut r, &mut buf).await;
930 }
931 buf
932 });
933
934 match tokio::time::timeout(timeout, child.wait()).await {
935 Ok(Ok(status)) => {
936 let stderr_bytes = stderr_task.await.unwrap_or_default();
937 let stderr = String::from_utf8_lossy(&stderr_bytes).into_owned();
938 Ok((status, stderr))
939 }
940 Ok(Err(e)) => {
941 stderr_task.abort();
942 Err(ProfileExecOutcome::WaitFailure {
943 pid,
944 binary: binary.to_string(),
945 error: e.to_string(),
946 })
947 }
948 Err(_) => {
949 let _ = child.start_kill();
951 let _ = child.wait().await;
952 let stderr_bytes = stderr_task.await.unwrap_or_default();
953 let stderr = String::from_utf8_lossy(&stderr_bytes).into_owned();
954 Err(ProfileExecOutcome::TimedOut {
955 pid,
956 binary: binary.to_string(),
957 timeout,
958 stderr,
959 })
960 }
961 }
962}
963
964async fn try_profile(
967 binary: &str,
968 pid: u32,
969 request: &ValidatedProfileRequest,
970) -> Option<ProfileExecOutcome> {
971 let timeout = request.subprocess_timeout();
972 let tmp_dir = match tempfile::tempdir() {
973 Ok(d) => d,
974 Err(e) => {
975 return Some(ProfileExecOutcome::TempDirFailure {
976 pid,
977 binary: binary.to_string(),
978 error: e.to_string(),
979 });
980 }
981 };
982 let svg_path = tmp_dir.path().join("profile.svg");
983
984 let child = match build_record_command(binary, pid, request, &svg_path).spawn() {
985 Ok(c) => c,
986 Err(e) if e.kind() == std::io::ErrorKind::NotFound => return None,
987 Err(e) => {
988 return Some(ProfileExecOutcome::SubprocessSpawnFailure {
989 pid,
990 binary: binary.to_string(),
991 error: e.to_string(),
992 });
993 }
994 };
995
996 let (status, stderr) = match collect_profile_with_timeout(child, pid, binary, timeout).await {
997 Ok(pair) => pair,
998 Err(outcome) => return Some(outcome),
999 };
1000
1001 if !status.success() {
1002 return Some(ProfileExecOutcome::ExitFailure {
1003 pid,
1004 binary: binary.to_string(),
1005 exit_code: status.code(),
1006 stderr,
1007 });
1008 }
1009
1010 match std::fs::read(&svg_path) {
1011 Ok(bytes) if bytes.is_empty() => Some(ProfileExecOutcome::OutputEmpty {
1012 pid,
1013 binary: binary.to_string(),
1014 }),
1015 Ok(svg) => Some(ProfileExecOutcome::Ok {
1016 pid,
1017 binary: binary.to_string(),
1018 svg,
1019 }),
1020 Err(e) if e.kind() == std::io::ErrorKind::NotFound => {
1021 Some(ProfileExecOutcome::OutputMissing {
1022 pid,
1023 binary: binary.to_string(),
1024 })
1025 }
1026 Err(e) => Some(ProfileExecOutcome::OutputReadFailure {
1027 pid,
1028 binary: binary.to_string(),
1029 error: e.to_string(),
1030 }),
1031 }
1032}
1033
1034#[cfg(test)]
1035mod tests {
1036 use std::io::Write;
1037 use std::os::unix::fs::PermissionsExt;
1038 use std::os::unix::process::ExitStatusExt;
1039 use std::time::Duration;
1040
1041 use tokio::process::Command;
1042
1043 use super::*;
1044
1045 #[test]
1046 fn pyspy_result_wirevalue_roundtrip() {
1047 let original = PySpyResult::Ok {
1052 pid: 42,
1053 binary: "py-spy".to_string(),
1054 stack_traces: vec![PySpyStackTrace {
1055 pid: 42,
1056 thread_id: 1,
1057 thread_name: Some("main".to_string()),
1058 os_thread_id: Some(100),
1059 active: true,
1060 owns_gil: true,
1061 frames: vec![PySpyFrame {
1062 name: "do_work".to_string(),
1063 filename: "test.py".to_string(),
1064 module: None,
1065 short_filename: None,
1066 line: 10,
1067 locals: None,
1068 is_entry: false,
1069 }],
1070 }],
1071 warnings: vec![],
1072 };
1073 let any: wirevalue::Any = wirevalue::Any::serialize(&original).expect("serialize");
1074 let restored: PySpyResult = any.deserialized().expect("deserialize");
1075 assert_eq!(original, restored);
1076 }
1077
1078 #[test]
1079 fn candidates_no_env() {
1080 let candidates = resolve_candidates(None);
1082 assert_eq!(candidates.len(), 1);
1083 assert_eq!(candidates[0].0, "py-spy");
1084 assert_eq!(candidates[0].1, "py-spy on PATH");
1085 }
1086
1087 #[test]
1088 fn candidates_env_first_then_path() {
1089 let candidates = resolve_candidates(Some("/custom/py-spy".to_string()));
1091 assert_eq!(candidates.len(), 2);
1092 assert_eq!(candidates[0].0, "/custom/py-spy");
1093 assert!(candidates[0].1.contains("PYSPY_BIN=/custom/py-spy"));
1094 assert_eq!(candidates[1].0, "py-spy");
1095 }
1096
1097 #[test]
1098 fn candidates_empty_env_ignored() {
1099 let candidates = resolve_candidates(Some(String::new()));
1101 assert_eq!(candidates.len(), 1);
1102 assert_eq!(candidates[0].0, "py-spy");
1103 }
1104
1105 #[test]
1106 fn output_success_parses_json() {
1107 let json = serde_json::json!([{
1109 "pid": 42,
1110 "thread_id": 1234,
1111 "thread_name": "MainThread",
1112 "os_thread_id": 5678,
1113 "active": true,
1114 "owns_gil": true,
1115 "frames": [{
1116 "name": "do_work",
1117 "filename": "foo.py",
1118 "module": null,
1119 "short_filename": null,
1120 "line": 10,
1121 "locals": null,
1122 "is_entry": false
1123 }]
1124 }]);
1125 let output = std::process::Output {
1126 status: std::process::ExitStatus::default(),
1127 stdout: serde_json::to_vec(&json).unwrap(),
1128 stderr: vec![],
1129 };
1130 let result = map_output(output, 42, "/usr/bin/py-spy");
1131 match result {
1132 PySpyResult::Ok {
1133 pid,
1134 binary,
1135 stack_traces,
1136 ..
1137 } => {
1138 assert_eq!(pid, 42);
1139 assert_eq!(binary, "/usr/bin/py-spy");
1140 assert_eq!(stack_traces.len(), 1);
1141 assert_eq!(stack_traces[0].thread_id, 1234);
1142 assert_eq!(stack_traces[0].thread_name.as_deref(), Some("MainThread"));
1143 assert!(stack_traces[0].owns_gil);
1144 assert_eq!(stack_traces[0].frames.len(), 1);
1145 assert_eq!(stack_traces[0].frames[0].name, "do_work");
1146 assert_eq!(stack_traces[0].frames[0].filename, "foo.py");
1147 assert_eq!(stack_traces[0].frames[0].line, 10);
1148 }
1149 other => panic!("expected Ok, got {:?}", other),
1150 }
1151 }
1152
1153 #[test]
1154 fn output_invalid_json_maps_to_failed() {
1155 let output = std::process::Output {
1157 status: std::process::ExitStatus::default(),
1158 stdout: b"not valid json".to_vec(),
1159 stderr: vec![],
1160 };
1161 let result = map_output(output, 42, "py-spy");
1162 match result {
1163 PySpyResult::Failed { pid, stderr, .. } => {
1164 assert_eq!(pid, 42);
1165 assert!(
1166 stderr.contains("failed to parse py-spy JSON output"),
1167 "unexpected stderr: {stderr}"
1168 );
1169 }
1170 other => panic!("expected Failed, got {:?}", other),
1171 }
1172 }
1173
1174 #[test]
1175 fn output_nonzero_exit_maps_to_failed() {
1176 let status = std::process::ExitStatus::from_raw(256); let output = std::process::Output {
1179 status,
1180 stdout: vec![],
1181 stderr: b"Permission denied".to_vec(),
1182 };
1183 let result = map_output(output, 99, "py-spy");
1184 match result {
1185 PySpyResult::Failed {
1186 pid,
1187 binary,
1188 exit_code,
1189 stderr,
1190 } => {
1191 assert_eq!(pid, 99);
1192 assert_eq!(binary, "py-spy");
1193 assert_eq!(exit_code, Some(1));
1194 assert_eq!(stderr, "Permission denied");
1195 }
1196 other => panic!("expected Failed, got {:?}", other),
1197 }
1198 }
1199
1200 #[test]
1201 fn output_preserves_caller_pid() {
1202 let json = serde_json::json!([]);
1204 let output = std::process::Output {
1205 status: std::process::ExitStatus::default(),
1206 stdout: serde_json::to_vec(&json).unwrap(),
1207 stderr: vec![],
1208 };
1209 let result = map_output(output, 12345, "bin");
1210 match result {
1211 PySpyResult::Ok { pid, .. } => assert_eq!(pid, 12345),
1212 other => panic!("expected Ok, got {:?}", other),
1213 }
1214 }
1215
1216 fn default_opts() -> PySpyOpts {
1217 PySpyOpts {
1218 threads: false,
1219 native: false,
1220 native_all: false,
1221 nonblocking: false,
1222 }
1223 }
1224
1225 #[tokio::test]
1226 async fn exec_missing_binary_returns_none() {
1227 let result = try_exec(
1229 "/definitely/not/a/real/binary",
1230 1,
1231 &default_opts(),
1232 std::time::Duration::from_secs(5),
1233 )
1234 .await;
1235 assert!(result.is_none());
1236 }
1237
1238 #[tokio::test]
1239 async fn exec_present_binary_returns_some() {
1240 let result = try_exec(
1243 "true",
1244 1,
1245 &default_opts(),
1246 std::time::Duration::from_secs(5),
1247 )
1248 .await;
1249 match result {
1250 Some(PySpyResult::Failed { stderr, .. }) => {
1251 assert!(
1252 stderr.contains("parse"),
1253 "expected JSON parse error, got: {stderr}"
1254 );
1255 }
1256 other => panic!("expected Some(Failed{{parse..}}), got: {other:?}"),
1257 }
1258 }
1259
1260 #[tokio::test]
1261 async fn collect_timeout_kills_child_and_returns_failed() {
1262 let child = Command::new("sleep")
1265 .arg("100")
1266 .stdout(std::process::Stdio::piped())
1267 .stderr(std::process::Stdio::piped())
1268 .spawn()
1269 .expect("sleep must be available");
1270
1271 let result = collect_with_timeout(
1272 child,
1273 std::process::id(),
1274 "sleep",
1275 std::time::Duration::from_millis(100),
1276 )
1277 .await;
1278
1279 match result {
1280 PySpyResult::Failed { stderr, .. } => {
1281 assert!(
1282 stderr.contains("timed out"),
1283 "expected timeout message, got: {stderr}"
1284 );
1285 }
1286 other => panic!("expected Failed, got {:?}", other),
1287 }
1288 }
1289
1290 #[tokio::test]
1291 async fn exec_failing_binary_returns_failed() {
1292 let result = try_exec(
1294 "false",
1295 42,
1296 &default_opts(),
1297 std::time::Duration::from_secs(5),
1298 )
1299 .await;
1300 assert!(result.is_some());
1301 match result.unwrap() {
1302 PySpyResult::Failed {
1303 pid,
1304 binary,
1305 exit_code,
1306 ..
1307 } => {
1308 assert_eq!(pid, 42);
1309 assert_eq!(binary, "false");
1310 assert!(exit_code.is_some());
1311 }
1312 other => panic!("expected Failed, got {:?}", other),
1313 }
1314 }
1315
1316 fn write_fake_pyspy(script_body: &str) -> tempfile::TempPath {
1324 let mut f = tempfile::NamedTempFile::new().expect("create temp file");
1325 write!(f, "#!/bin/sh\n{script_body}").expect("write script");
1326 f.as_file().sync_all().expect("sync");
1327 std::fs::set_permissions(f.path(), std::fs::Permissions::from_mode(0o755))
1328 .expect("chmod +x");
1329 f.into_temp_path()
1330 }
1331
1332 fn read_log(script_path: &std::path::Path) -> Vec<String> {
1335 let log_path = format!("{}.log", script_path.display());
1336 match std::fs::read_to_string(&log_path) {
1337 Ok(contents) => contents.lines().map(String::from).collect(),
1338 Err(_) => vec![],
1339 }
1340 }
1341
1342 #[tokio::test]
1343 async fn native_all_downgrade_succeeds() {
1344 let script = write_fake_pyspy(
1348 r#"
1349echo "$@" >> "$0.log"
1350for arg in "$@"; do
1351 if [ "$arg" = "--native-all" ]; then
1352 echo "unrecognized option --native-all" >&2
1353 exit 2
1354 fi
1355done
1356echo "[]"
1357exit 0
1358"#,
1359 );
1360 let opts = PySpyOpts {
1361 threads: false,
1362 native: true,
1363 native_all: true,
1364 nonblocking: false,
1365 };
1366 let result = try_exec(
1367 script.to_str().unwrap(),
1368 1,
1369 &opts,
1370 std::time::Duration::from_secs(5),
1371 )
1372 .await;
1373 let result = result.expect("expected Some");
1375 match &result {
1376 PySpyResult::Ok { warnings, .. } => {
1377 assert!(
1378 warnings.iter().any(|w| w.contains("fell back to --native")),
1379 "PS-11c: expected fallback warning, got: {warnings:?}"
1380 );
1381 }
1382 other => panic!("expected Ok, got: {other:?}"),
1383 }
1384 let log = read_log(&script);
1386 assert_eq!(
1387 log.len(),
1388 2,
1389 "PS-11b: expected exactly 2 invocations, got {}",
1390 log.len()
1391 );
1392 assert!(
1393 log[0].contains("--native-all"),
1394 "PS-11a: first invocation must include --native-all, got: {}",
1395 log[0]
1396 );
1397 assert!(
1398 !log[1].contains("--native-all"),
1399 "PS-11a: second invocation must NOT include --native-all, got: {}",
1400 log[1]
1401 );
1402 }
1403
1404 #[tokio::test]
1405 async fn native_all_downgrade_fails_retries_continue() {
1406 let script = write_fake_pyspy(
1409 r#"
1410echo "$@" >> "$0.log"
1411for arg in "$@"; do
1412 if [ "$arg" = "--native-all" ]; then
1413 echo "unrecognized option --native-all" >&2
1414 exit 2
1415 fi
1416done
1417echo "Permission denied" >&2
1418exit 1
1419"#,
1420 );
1421 let opts = PySpyOpts {
1422 threads: false,
1423 native: true,
1424 native_all: true,
1425 nonblocking: true, };
1427 let result = try_exec(
1428 script.to_str().unwrap(),
1429 1,
1430 &opts,
1431 std::time::Duration::from_secs(10),
1432 )
1433 .await;
1434 let result = result.expect("expected Some");
1436 match &result {
1437 PySpyResult::Failed {
1438 stderr, exit_code, ..
1439 } => {
1440 assert!(
1441 stderr.contains("Permission denied"),
1442 "PS-11d: expected generic failure, got: {stderr}"
1443 );
1444 assert_eq!(*exit_code, Some(1));
1445 }
1446 other => panic!("expected Failed, got: {other:?}"),
1447 }
1448 let log = read_log(&script);
1453 assert_eq!(log.len(), 4, "expected 4 invocations, got {}", log.len());
1454 assert!(
1455 log[0].contains("--native-all"),
1456 "PS-11a: first invocation must include --native-all, got: {}",
1457 log[0]
1458 );
1459 for (i, line) in log[1..].iter().enumerate() {
1460 assert!(
1461 !line.contains("--native-all"),
1462 "PS-11e: invocation {} must NOT include --native-all, got: {}",
1463 i + 1,
1464 line
1465 );
1466 }
1467 }
1468
1469 #[tokio::test]
1471 async fn profile_collect_timeout_returns_timed_out() {
1472 let child = Command::new("sh")
1473 .arg("-c")
1474 .arg("echo diag >&2; sleep 60")
1475 .stdout(std::process::Stdio::null())
1476 .stderr(std::process::Stdio::piped())
1477 .spawn()
1478 .expect("sh must be available");
1479
1480 let result = collect_profile_with_timeout(
1481 child,
1482 std::process::id(),
1483 "sh",
1484 std::time::Duration::from_millis(200),
1485 )
1486 .await;
1487
1488 match result {
1489 Err(ProfileExecOutcome::TimedOut { stderr, .. }) => {
1490 assert!(
1491 stderr.contains("diag"),
1492 "expected partial stderr captured after kill, got: {stderr}"
1493 );
1494 }
1495 other => panic!("expected TimedOut, got: {other:?}"),
1496 }
1497 }
1498
1499 fn test_request() -> ValidatedProfileRequest {
1500 ValidatedProfileRequest::try_new(
1501 &PySpyProfileOpts {
1502 duration_s: 1,
1503 rate_hz: 100,
1504 native: false,
1505 threads: false,
1506 nonblocking: false,
1507 },
1508 std::time::Duration::from_secs(300),
1509 )
1510 .unwrap()
1511 }
1512
1513 #[tokio::test]
1515 async fn profile_try_missing_binary_returns_none() {
1516 let result = try_profile("/definitely/not/a/real/binary", 1, &test_request()).await;
1517 assert!(result.is_none(), "missing binary must return None");
1518 }
1519
1520 #[tokio::test]
1522 async fn profile_success_exit_empty_file_returns_output_empty() {
1523 let script = write_fake_pyspy(
1524 r#"
1525output=""
1526while [ $# -gt 0 ]; do
1527 case "$1" in
1528 --output) shift; output="$1" ;;
1529 esac
1530 shift
1531done
1532touch "$output"
1533exit 0
1534"#,
1535 );
1536 let result = try_profile(script.to_str().unwrap(), 1, &test_request()).await;
1537 assert!(
1538 matches!(result, Some(ProfileExecOutcome::OutputEmpty { .. })),
1539 "PP-3: expected OutputEmpty, got: {result:?}"
1540 );
1541 }
1542
1543 #[tokio::test]
1545 async fn profile_success_exit_missing_file_returns_output_missing() {
1546 let script = write_fake_pyspy("exit 0\n");
1547 let result = try_profile(script.to_str().unwrap(), 1, &test_request()).await;
1548 assert!(
1549 matches!(result, Some(ProfileExecOutcome::OutputMissing { .. })),
1550 "PP-3: expected OutputMissing, got: {result:?}"
1551 );
1552 }
1553
1554 #[test]
1556 fn validated_request_rejects_zero_duration() {
1557 let opts = PySpyProfileOpts {
1558 duration_s: 0,
1559 rate_hz: 100,
1560 native: false,
1561 threads: false,
1562 nonblocking: false,
1563 };
1564 let err = ValidatedProfileRequest::try_new(&opts, std::time::Duration::from_secs(300));
1565 assert!(err.is_err());
1566 assert!(err.unwrap_err().contains("positive"));
1567 }
1568
1569 #[test]
1571 fn validated_request_rejects_over_max_duration() {
1572 let opts = PySpyProfileOpts {
1573 duration_s: 999,
1574 rate_hz: 100,
1575 native: false,
1576 threads: false,
1577 nonblocking: false,
1578 };
1579 let err = ValidatedProfileRequest::try_new(&opts, std::time::Duration::from_secs(300));
1580 assert!(err.is_err());
1581 assert!(err.unwrap_err().contains("exceeds max"));
1582 }
1583
1584 #[test]
1586 fn validated_request_rejects_zero_rate() {
1587 let opts = PySpyProfileOpts {
1588 duration_s: 5,
1589 rate_hz: 0,
1590 native: false,
1591 threads: false,
1592 nonblocking: false,
1593 };
1594 let err = ValidatedProfileRequest::try_new(&opts, std::time::Duration::from_secs(300));
1595 assert!(err.is_err());
1596 assert!(err.unwrap_err().contains("rate_hz"));
1597 }
1598
1599 #[test]
1601 fn validated_request_rejects_excessive_rate() {
1602 let opts = PySpyProfileOpts {
1603 duration_s: 5,
1604 rate_hz: 9999,
1605 native: false,
1606 threads: false,
1607 nonblocking: false,
1608 };
1609 let err = ValidatedProfileRequest::try_new(&opts, std::time::Duration::from_secs(300));
1610 assert!(err.is_err());
1611 assert!(err.unwrap_err().contains("rate_hz"));
1612 }
1613
1614 #[test]
1616 fn validated_request_computes_exact_timeouts() {
1617 let opts = PySpyProfileOpts {
1618 duration_s: 30,
1619 rate_hz: 100,
1620 native: true,
1621 threads: false,
1622 nonblocking: false,
1623 };
1624 let req =
1625 ValidatedProfileRequest::try_new(&opts, std::time::Duration::from_secs(300)).unwrap();
1626 assert_eq!(req.duration().get(), 30);
1627 assert_eq!(req.rate().get(), 100);
1628 assert!(req.native());
1629 assert_eq!(req.subprocess_timeout(), std::time::Duration::from_secs(45));
1630 assert_eq!(req.bridge_timeout(), std::time::Duration::from_secs(50));
1631 }
1632
1633 #[test]
1635 fn profile_exec_outcome_conversion_is_identity() {
1636 let r = PySpyProfileResult::from(ProfileExecOutcome::Ok {
1638 pid: 1,
1639 binary: "b".into(),
1640 svg: vec![1],
1641 });
1642 assert!(matches!(r, PySpyProfileResult::Ok { pid: 1, .. }));
1643
1644 let r = PySpyProfileResult::from(ProfileExecOutcome::BinaryNotFound {
1645 searched: vec!["x".into()],
1646 });
1647 assert!(matches!(r, PySpyProfileResult::BinaryNotFound { .. }));
1648
1649 let r = PySpyProfileResult::from(ProfileExecOutcome::TimedOut {
1650 pid: 1,
1651 binary: "b".into(),
1652 timeout: Duration::from_secs(10),
1653 stderr: "s".into(),
1654 });
1655 assert!(matches!(
1656 r,
1657 PySpyProfileResult::TimedOut { timeout_s: 10, .. }
1658 ));
1659
1660 let r = PySpyProfileResult::from(ProfileExecOutcome::ExitFailure {
1661 pid: 1,
1662 binary: "b".into(),
1663 exit_code: Some(2),
1664 stderr: "e".into(),
1665 });
1666 assert!(matches!(
1667 r,
1668 PySpyProfileResult::ExitFailure {
1669 exit_code: Some(2),
1670 ..
1671 }
1672 ));
1673
1674 let r = PySpyProfileResult::from(ProfileExecOutcome::OutputMissing {
1675 pid: 1,
1676 binary: "b".into(),
1677 });
1678 assert!(matches!(
1679 r,
1680 PySpyProfileResult::OutputMissing { pid: 1, .. }
1681 ));
1682
1683 let r = PySpyProfileResult::from(ProfileExecOutcome::OutputEmpty {
1684 pid: 1,
1685 binary: "b".into(),
1686 });
1687 assert!(matches!(r, PySpyProfileResult::OutputEmpty { pid: 1, .. }));
1688
1689 let r = PySpyProfileResult::from(ProfileExecOutcome::OutputReadFailure {
1690 pid: 1,
1691 binary: "b".into(),
1692 error: "permission denied".into(),
1693 });
1694 assert!(matches!(r, PySpyProfileResult::OutputReadFailure { .. }));
1695
1696 let r = PySpyProfileResult::from(ProfileExecOutcome::SubprocessSpawnFailure {
1697 pid: 1,
1698 binary: "b".into(),
1699 error: "s".into(),
1700 });
1701 assert!(matches!(
1702 r,
1703 PySpyProfileResult::SubprocessSpawnFailure { .. }
1704 ));
1705
1706 let r = PySpyProfileResult::from(ProfileExecOutcome::WaitFailure {
1707 pid: 1,
1708 binary: "b".into(),
1709 error: "w".into(),
1710 });
1711 assert!(matches!(r, PySpyProfileResult::WaitFailure { .. }));
1712
1713 let r = PySpyProfileResult::from(ProfileExecOutcome::TempDirFailure {
1714 pid: 1,
1715 binary: "b".into(),
1716 error: "t".into(),
1717 });
1718 assert!(matches!(r, PySpyProfileResult::TempDirFailure { .. }));
1719 }
1720}