1#![allow(unused_assignments)]
16
17use std::collections::HashMap;
18use std::time::Duration;
19
20use async_trait::async_trait;
21use hyperactor::Actor;
22use hyperactor::ActorAddr;
23use hyperactor::ActorHandle;
24use hyperactor::Addr;
25use hyperactor::Client;
26use hyperactor::Context;
27use hyperactor::Data;
28use hyperactor::Endpoint as _;
29use hyperactor::Handler;
30use hyperactor::Instance;
31use hyperactor::PortAddr;
32use hyperactor::PortHandle;
33use hyperactor::PortRef;
34use hyperactor::RemoteEndpoint as _;
35use hyperactor::actor::handle_undeliverable_message;
36use hyperactor::actor::remote::Remote;
37use hyperactor::id::Label;
38use hyperactor::id::Uid;
39use hyperactor::mailbox::MessageEnvelope;
40use hyperactor::mailbox::Undeliverable;
41use hyperactor::mailbox::UndeliverableReason;
42use hyperactor::proc::Proc;
43use hyperactor::supervision::ActorSupervisionEvent;
44use hyperactor_cast::cast_actor::CAST_ACTOR_NAME;
45use hyperactor_config::CONFIG;
46use hyperactor_config::ConfigAttr;
47use hyperactor_config::Flattrs;
48use hyperactor_config::attrs::declare_attrs;
49use serde::Deserialize;
50use serde::Serialize;
51use typeuri::Named;
52
53use crate::config_dump::ConfigDump;
54use crate::config_dump::ConfigDumpResult;
55use crate::introspect::ProcessMemoryStats;
56use crate::mesh_id::ResourceId;
57use crate::pyspy::PySpyDump;
58use crate::pyspy::PySpyProfile;
59use crate::pyspy::PySpyProfileWorker;
60use crate::pyspy::PySpyWorker;
61use crate::resource;
62
63pub const PROC_AGENT_ACTOR_NAME: &str = "proc_agent";
65
66declare_attrs! {
67 @meta(CONFIG = ConfigAttr::new(
71 Some("HYPERACTOR_MESH_ORPHAN_TIMEOUT".to_string()),
72 Some("mesh_orphan_timeout".to_string()),
73 ))
74 pub attr MESH_ORPHAN_TIMEOUT: Option<Duration> = Some(Duration::from_secs(60));
75
76 @meta(CONFIG = ConfigAttr::new(
85 Some("HYPERACTOR_PROCESS_MEMORY_METRIC_INTERVAL".to_string()),
86 Some("process_memory_metric_interval".to_string()),
87 ))
88 pub attr PROCESS_MEMORY_METRIC_INTERVAL: Duration = Duration::from_secs(300);
89
90 pub(crate) attr STREAM_STATE_SUBSCRIBER: bool;
94}
95
96#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
116struct RepublishIntrospect {
117 emit_memory_metrics: bool,
118}
119wirevalue::register_type!(RepublishIntrospect);
120
121fn collect_live_children(
127 proc: &hyperactor::Proc,
128) -> (
129 Vec<hyperactor::introspect::IntrospectRef>,
130 Vec<crate::introspect::NodeRef>,
131) {
132 let all_keys = proc.all_instance_keys();
133 let mut children = Vec::with_capacity(all_keys.len());
134 let mut system_children = Vec::new();
135 for id in all_keys {
136 if let Some(cell) = proc.get_instance_by_id(&id) {
137 let actor_addr = cell.actor_addr().clone();
138 if cell.is_system() {
139 system_children.push(crate::introspect::NodeRef::Actor(actor_addr.clone()));
140 }
141 children.push(hyperactor::introspect::IntrospectRef::Actor(actor_addr));
142 }
143 }
144 (children, system_children)
145}
146
147#[derive(Debug)]
149struct ActorInstanceState {
150 create_rank: usize,
151 spawn: Result<ActorAddr, anyhow::Error>,
152 stop_initiated: bool,
156 supervision_event: Option<ActorSupervisionEvent>,
159 subscribers: Vec<PortRef<resource::State<ActorState>>>,
162 expiry_time: Option<std::time::SystemTime>,
165 generation: u64,
169 pending_wait_status: Vec<(resource::Status, PortRef<crate::StatusOverlay>)>,
173}
174
175impl ActorInstanceState {
176 fn status(&self) -> resource::Status {
179 match &self.spawn {
180 Err(e) => resource::Status::Failed(e.to_string()),
181 Ok(_) => match &self.supervision_event {
182 Some(event) if event.is_error() => resource::Status::Failed(format!("{}", event)),
183 Some(_) => resource::Status::Stopped,
184 None if self.stop_initiated => resource::Status::Stopping,
185 None => resource::Status::Running,
186 },
187 }
188 }
189
190 fn is_terminal(&self) -> bool {
193 match &self.spawn {
194 Err(_) => true,
195 Ok(_) => self.supervision_event.is_some(),
196 }
197 }
198
199 fn has_errors(&self) -> bool {
201 self.supervision_event
202 .as_ref()
203 .is_some_and(|e| e.is_error())
204 }
205
206 fn to_state(&self, id: &ResourceId) -> resource::State<ActorState> {
209 let status = self.status();
210 let actor_state = self.spawn.as_ref().ok().map(|actor_id| ActorState {
211 actor_id: actor_id.clone(),
212 create_rank: self.create_rank,
213 supervision_events: self.supervision_event.clone().into_iter().collect(),
214 });
215 resource::State {
216 id: id.clone(),
217 status,
218 state: actor_state,
219 generation: self.generation,
220 timestamp: std::time::SystemTime::now(),
221 }
222 }
223
224 fn notify_status_changed(&mut self, cx: &impl hyperactor::context::Actor, id: &ResourceId) {
229 let state = self.to_state(id);
231 for subscriber in &self.subscribers {
232 let mut headers = Flattrs::new();
233 headers.set(STREAM_STATE_SUBSCRIBER, true);
234 subscriber.post_with_headers(cx, headers, state.clone());
235 }
236
237 let status = self.status();
239 self.pending_wait_status.retain(|(min_status, reply)| {
240 if status >= *min_status {
241 let rank = self.create_rank;
242 let overlay =
243 crate::StatusOverlay::try_from_runs(vec![(rank..(rank + 1), status.clone())])
244 .expect("valid single-run overlay");
245 let _ = reply.post(cx, overlay);
246 false
247 } else {
248 true
249 }
250 });
251 }
252}
253
254#[derive(Clone, Debug, Default, PartialEq, Serialize, Deserialize, Named)]
255pub(crate) struct SelfCheck {}
256
257#[hyperactor::export(
278 handlers=[
279 ActorSupervisionEvent,
280 resource::CreateOrUpdate<ActorSpec>,
281 resource::Stop,
282 resource::StopAll,
283 resource::GetState<ActorState>,
284 resource::StreamState<ActorState>,
285 resource::KeepaliveGetState<ActorState>,
286 resource::GetRankStatus,
287 resource::WaitRankStatus,
288 RepublishIntrospect,
289 PySpyDump,
290 PySpyProfile,
291 ConfigDump,
292 ]
293)]
294pub struct ProcAgent {
295 proc: Proc,
296 remote: Remote,
297 actor_states: HashMap<ResourceId, ActorInstanceState>,
299 record_supervision_events: bool,
302 introspect_dirty: bool,
305 shutdown_tx: Option<tokio::sync::oneshot::Sender<i32>>,
309 stopping_all: bool,
313 mesh_orphan_timeout: Option<Duration>,
315}
316
317impl ProcAgent {
318 pub(crate) fn boot_v1(
319 proc: Proc,
320 shutdown_tx: Option<tokio::sync::oneshot::Sender<i32>>,
321 ) -> Result<ActorHandle<Self>, anyhow::Error> {
322 let cast_handle = proc.spawn_with_uid(
323 Uid::singleton(Label::strip(CAST_ACTOR_NAME)),
324 hyperactor_cast::cast_actor::CastActor::default(),
325 )?;
326 cast_handle.bind::<hyperactor_cast::cast_actor::CastActor>();
327
328 let orphan_timeout = hyperactor_config::global::get(MESH_ORPHAN_TIMEOUT);
329 let agent = ProcAgent {
330 proc: proc.clone(),
331 remote: Remote::collect(),
332 actor_states: HashMap::new(),
333 record_supervision_events: true,
334 introspect_dirty: false,
335 shutdown_tx,
336 stopping_all: false,
337 mesh_orphan_timeout: orphan_timeout,
338 };
339 proc.spawn_with_uid::<Self>(
340 Uid::singleton(Label::new(PROC_AGENT_ACTOR_NAME).unwrap()),
341 agent,
342 )
343 }
344
345 fn all_actors_terminal(&self) -> bool {
349 self.actor_states.values().all(|state| state.is_terminal())
350 }
351
352 async fn shutdown(&mut self) {
356 let has_errors = self.actor_states.values().any(|state| state.has_errors());
357 let exit_code = if has_errors { 1 } else { 0 };
358
359 let flush_timeout =
360 hyperactor_config::global::get(hyperactor::config::FORWARDER_FLUSH_TIMEOUT);
361 match tokio::time::timeout(flush_timeout, self.proc.flush()).await {
362 Ok(Err(err)) => {
363 tracing::warn!("forwarder flush failed during shutdown: {}", err);
364 }
365 Err(_elapsed) => {
366 tracing::warn!("forwarder flush timed out during shutdown");
367 }
368 Ok(Ok(())) => {}
369 }
370
371 self.proc.join_mailbox_server().await;
375
376 tracing::info!(
377 "shutting down process after all actors reached terminal state (exit_code={})",
378 exit_code,
379 );
380
381 if let Some(tx) = self.shutdown_tx.take() {
382 let _ = tx.send(exit_code);
383 return;
384 }
385 std::process::exit(exit_code);
386 }
387
388 fn stop_actor_by_id(&self, actor_id: &ActorAddr, reason: &str) {
391 tracing::info!(
392 name = "StopActor",
393 %actor_id,
394 actor_name = actor_id.log_name(),
395 %reason,
396 );
397 self.proc.stop_actor(actor_id.id(), reason.to_string());
398 }
399
400 fn publish_introspect_properties(
403 &self,
404 cx: &impl hyperactor::context::Actor,
405 ) -> ProcessMemoryStats {
406 let (mut children, mut system_children) = collect_live_children(&self.proc);
407
408 let mut stopped_children: Vec<crate::introspect::NodeRef> = Vec::new();
412 for id in self.proc.all_terminated_actor_ids() {
413 let child_ref = hyperactor::introspect::IntrospectRef::Actor(id.clone());
414 let node_ref = crate::introspect::NodeRef::Actor(id.clone());
415 stopped_children.push(node_ref.clone());
416 if let Some(snapshot) = self.proc.terminated_snapshot(&id) {
417 let snapshot_attrs: hyperactor_config::Attrs =
418 serde_json::from_str(&snapshot.attrs).unwrap_or_default();
419 if snapshot_attrs
420 .get(hyperactor::introspect::IS_SYSTEM)
421 .copied()
422 .unwrap_or(false)
423 {
424 system_children.push(node_ref);
425 }
426 }
427 if !children.contains(&child_ref) {
428 children.push(child_ref);
429 }
430 }
431
432 let stopped_retention_cap =
433 hyperactor_config::global::get(hyperactor::config::TERMINATED_SNAPSHOT_RETENTION);
434
435 let failed_actor_count = self
437 .actor_states
438 .values()
439 .filter(|s| s.has_errors())
440 .count();
441
442 let num_live = children.len();
444 let mut attrs = hyperactor_config::Attrs::new();
445 attrs.set(crate::introspect::NODE_TYPE, "proc".to_string());
446 attrs.set(
447 crate::introspect::PROC_NAME,
448 self.proc
449 .proc_addr()
450 .label()
451 .map(|l| l.as_str().to_string())
452 .unwrap_or_else(|| self.proc.proc_addr().id().to_string()),
453 );
454 attrs.set(crate::introspect::NUM_ACTORS, num_live);
455 attrs.set(hyperactor::introspect::CHILDREN, children);
456 attrs.set(crate::introspect::SYSTEM_CHILDREN, system_children);
457 attrs.set(crate::introspect::STOPPED_CHILDREN, stopped_children);
458 attrs.set(
459 crate::introspect::STOPPED_RETENTION_CAP,
460 stopped_retention_cap,
461 );
462 attrs.set(crate::introspect::IS_POISONED, failed_actor_count > 0);
463 attrs.set(crate::introspect::FAILED_ACTOR_COUNT, failed_actor_count);
464
465 let memory = crate::introspect::ProcessMemoryStats::read_from_procfs();
470 memory.to_attrs(&mut attrs);
471
472 let queue_total = self.proc.queue_depth_total();
474 attrs.set(crate::introspect::ACTOR_WORK_QUEUE_DEPTH_TOTAL, queue_total);
475
476 let mut queue_max: u64 = 0;
478 for actor_id in self.proc.all_instance_keys() {
479 if let Some(cell) = self.proc.get_instance_by_id(&actor_id) {
480 queue_max = queue_max.max(cell.queue_depth());
481 }
482 }
483 attrs.set(crate::introspect::ACTOR_WORK_QUEUE_DEPTH_MAX, queue_max);
484
485 attrs.set(
487 crate::introspect::ACTOR_WORK_QUEUE_DEPTH_HIGH_WATER_MARK,
488 self.proc.queue_depth_high_water_mark(),
489 );
490 attrs.set(
491 crate::introspect::LAST_NONZERO_QUEUE_DEPTH_AGE_MS,
492 self.proc.last_nonzero_queue_depth_age_ms(),
493 );
494
495 cx.instance().publish_attrs(attrs);
496
497 memory
498 }
499}
500
501#[async_trait]
502impl Actor for ProcAgent {
503 async fn init(&mut self, this: &Instance<Self>) -> Result<(), anyhow::Error> {
504 this.set_system();
505 self.proc.set_supervision_coordinator(this.port().bind())?;
506 let _ = self.publish_introspect_properties(this);
507
508 let proc = self.proc.clone();
511 let self_id = this.self_addr().clone();
512 this.set_query_child_handler(move |child_ref| {
513 use hyperactor::introspect::IntrospectResult;
514
515 if let Addr::Actor(actor_ref) = child_ref
516 && let Some(snapshot) = proc.terminated_snapshot(actor_ref)
517 {
518 return snapshot;
519 }
520
521 if let Addr::Proc(proc_ref) = child_ref
529 && *proc_ref == proc.proc_addr()
530 {
531 let (mut children, mut system_children) = collect_live_children(&proc);
532
533 let mut stopped_children: Vec<crate::introspect::NodeRef> = Vec::new();
534 for id in proc.all_terminated_actor_ids() {
535 let child_ref = hyperactor::introspect::IntrospectRef::Actor(id.clone());
536 let node_ref = crate::introspect::NodeRef::Actor(id.clone());
537 stopped_children.push(node_ref.clone());
538 if let Some(snapshot) = proc.terminated_snapshot(&id) {
539 let snapshot_attrs: hyperactor_config::Attrs =
540 serde_json::from_str(&snapshot.attrs).unwrap_or_default();
541 if snapshot_attrs
542 .get(hyperactor::introspect::IS_SYSTEM)
543 .copied()
544 .unwrap_or(false)
545 {
546 system_children.push(node_ref);
547 }
548 }
549 if !children.contains(&child_ref) {
550 children.push(child_ref);
551 }
552 }
553
554 let stopped_retention_cap = hyperactor_config::global::get(
555 hyperactor::config::TERMINATED_SNAPSHOT_RETENTION,
556 );
557
558 let (is_poisoned, failed_actor_count) = proc
559 .get_instance(&self_id)
560 .and_then(|cell| cell.published_attrs())
561 .map(|attrs| {
562 let is_poisoned = attrs
563 .get(crate::introspect::IS_POISONED)
564 .copied()
565 .unwrap_or(false);
566 let failed_actor_count = attrs
567 .get(crate::introspect::FAILED_ACTOR_COUNT)
568 .copied()
569 .unwrap_or(0);
570 (is_poisoned, failed_actor_count)
571 })
572 .unwrap_or((false, 0));
573
574 let num_live = children.len();
576 let mut attrs = hyperactor_config::Attrs::new();
577 attrs.set(crate::introspect::NODE_TYPE, "proc".to_string());
578 attrs.set(
579 crate::introspect::PROC_NAME,
580 proc_ref
581 .label()
582 .map(|l| l.as_str().to_string())
583 .unwrap_or_else(|| proc_ref.id().to_string()),
584 );
585 attrs.set(crate::introspect::NUM_ACTORS, num_live);
586 attrs.set(crate::introspect::SYSTEM_CHILDREN, system_children);
587 attrs.set(crate::introspect::STOPPED_CHILDREN, stopped_children);
588 attrs.set(
589 crate::introspect::STOPPED_RETENTION_CAP,
590 stopped_retention_cap,
591 );
592 attrs.set(crate::introspect::IS_POISONED, is_poisoned);
593 attrs.set(crate::introspect::FAILED_ACTOR_COUNT, failed_actor_count);
594
595 let memory = crate::introspect::ProcessMemoryStats::read_from_procfs();
598 memory.to_attrs(&mut attrs);
599 attrs.set(
600 crate::introspect::ACTOR_WORK_QUEUE_DEPTH_TOTAL,
601 proc.queue_depth_total(),
602 );
603 let mut queue_max: u64 = 0;
604 for aid in proc.all_instance_keys() {
605 if let Some(cell) = proc.get_instance_by_id(&aid) {
606 queue_max = queue_max.max(cell.queue_depth());
607 }
608 }
609 attrs.set(crate::introspect::ACTOR_WORK_QUEUE_DEPTH_MAX, queue_max);
610 attrs.set(
611 crate::introspect::ACTOR_WORK_QUEUE_DEPTH_HIGH_WATER_MARK,
612 proc.queue_depth_high_water_mark(),
613 );
614 attrs.set(
615 crate::introspect::LAST_NONZERO_QUEUE_DEPTH_AGE_MS,
616 proc.last_nonzero_queue_depth_age_ms(),
617 );
618
619 let attrs_json = serde_json::to_string(&attrs).unwrap_or_else(|_| "{}".to_string());
620
621 return IntrospectResult {
622 identity: hyperactor::introspect::IntrospectRef::Proc(proc_ref.clone()),
623 attrs: attrs_json,
624 children,
625 parent: None,
626 as_of: std::time::SystemTime::now(),
627 };
628 }
629
630 {
631 let mut error_attrs = hyperactor_config::Attrs::new();
632 error_attrs.set(hyperactor::introspect::ERROR_CODE, "not_found".to_string());
633 error_attrs.set(
634 hyperactor::introspect::ERROR_MESSAGE,
635 format!("child {} not found", child_ref),
636 );
637 let identity = match child_ref {
638 Addr::Proc(p) => hyperactor::introspect::IntrospectRef::Proc(p.clone()),
639 Addr::Actor(a) => hyperactor::introspect::IntrospectRef::Actor(a.clone()),
640 Addr::Port(p) => hyperactor::introspect::IntrospectRef::Actor(p.actor_addr()),
641 };
642 IntrospectResult {
643 identity,
644 attrs: serde_json::to_string(&error_attrs).unwrap_or_else(|_| "{}".to_string()),
645 children: Vec::new(),
646 parent: None,
647 as_of: std::time::SystemTime::now(),
648 }
649 }
650 });
651
652 if let Some(delay) = &self.mesh_orphan_timeout {
653 this.post_after(this, SelfCheck::default(), *delay);
654 }
655 if cfg!(target_os = "linux") {
656 let interval = hyperactor_config::global::get(PROCESS_MEMORY_METRIC_INTERVAL);
657 if !interval.is_zero() {
658 this.post_after(
659 this,
660 RepublishIntrospect {
661 emit_memory_metrics: true,
662 },
663 interval,
664 );
665 }
666 }
667 Ok(())
668 }
669
670 async fn handle_undeliverable_message(
671 &mut self,
672 cx: &Instance<Self>,
673 reason: UndeliverableReason,
674 envelope: Undeliverable<MessageEnvelope>,
675 ) -> Result<(), anyhow::Error> {
676 let Some(returned) = envelope.as_message() else {
677 return handle_undeliverable_message(cx, reason, envelope);
678 };
679 if let Some(true) = returned.headers().get(STREAM_STATE_SUBSCRIBER) {
680 let dest_port_id: PortAddr = returned.dest().clone();
681 let port = PortRef::<resource::State<ActorState>>::attest(dest_port_id);
682 for instance in self.actor_states.values_mut() {
684 instance.subscribers.retain(|s| s != &port);
685 }
686 Ok(())
687 } else {
688 handle_undeliverable_message(cx, reason, envelope)
689 }
690 }
691
692 async fn handle_invalid_reference(
693 &mut self,
694 cx: &Instance<Self>,
695 invalid: hyperactor::mailbox::InvalidReference,
696 envelope: Undeliverable<MessageEnvelope>,
697 ) -> Result<(), anyhow::Error> {
698 let Some(returned) = envelope.as_message() else {
699 return hyperactor::actor::handle_invalid_reference(cx, invalid, envelope);
700 };
701 if let Some(true) = returned.headers().get(STREAM_STATE_SUBSCRIBER) {
702 let dest_port_id: PortAddr = returned.dest().clone();
703 let port = PortRef::<resource::State<ActorState>>::attest(dest_port_id);
704 for instance in self.actor_states.values_mut() {
705 instance.subscribers.retain(|s| s != &port);
706 }
707 Ok(())
708 } else {
709 hyperactor::actor::handle_invalid_reference(cx, invalid, envelope)
710 }
711 }
712}
713
714#[async_trait]
715impl Handler<ActorSupervisionEvent> for ProcAgent {
716 async fn handle(
717 &mut self,
718 cx: &Context<Self>,
719 event: ActorSupervisionEvent,
720 ) -> anyhow::Result<()> {
721 if self.record_supervision_events {
722 if event.is_error() {
723 tracing::warn!(
724 name = "SupervisionEvent",
725 proc_id = %self.proc.proc_addr(),
726 %event,
727 "recording supervision error",
728 );
729 } else {
730 tracing::debug!(
731 name = "SupervisionEvent",
732 proc_id = %self.proc.proc_addr(),
733 %event,
734 "recording non-error supervision event",
735 );
736 }
737 if let Some((id, instance)) = self.actor_states.iter_mut().find(|(_, s)| {
739 s.spawn
740 .as_ref()
741 .ok()
742 .is_some_and(|actor_id| actor_id.id() == event.actor_id.id())
743 }) {
744 instance.supervision_event = Some(event.clone());
745 instance.generation += 1;
746 let id = id.clone();
747 instance.notify_status_changed(cx, &id);
748 }
749 if !self.introspect_dirty {
753 self.introspect_dirty = true;
754 cx.post_after(
755 cx,
756 RepublishIntrospect {
757 emit_memory_metrics: false,
758 },
759 std::time::Duration::from_millis(100),
760 );
761 }
762
763 if self.stopping_all && self.all_actors_terminal() {
766 self.shutdown().await;
767 }
768 }
769 if !self.record_supervision_events && event.is_error() {
770 tracing::error!(
773 name = "supervision_event_transmit_failed",
774 proc_id = %cx.self_addr().proc_addr(),
775 %event,
776 "could not propagate supervision event, crashing",
777 );
778
779 std::process::exit(1);
782 }
783 Ok(())
784 }
785}
786
787#[async_trait]
788impl Handler<RepublishIntrospect> for ProcAgent {
789 async fn handle(&mut self, cx: &Context<Self>, msg: RepublishIntrospect) -> anyhow::Result<()> {
790 self.introspect_dirty = false;
791 let memory = self.publish_introspect_properties(cx);
792 if msg.emit_memory_metrics {
793 let proc_id = self.proc.proc_addr().to_string();
794 let pid = std::process::id() as i64;
795 if let Some(rss) = memory.process_rss_bytes {
796 crate::metrics::PROCESS_RSS_BYTES.record(
797 rss as f64,
798 hyperactor_telemetry::kv_pairs!(
799 "proc_id" => proc_id.clone(),
800 "pid" => pid,
801 ),
802 );
803 }
804 if let Some(vm) = memory.process_vm_size_bytes {
805 crate::metrics::PROCESS_VM_SIZE_BYTES.record(
806 vm as f64,
807 hyperactor_telemetry::kv_pairs!(
808 "proc_id" => proc_id,
809 "pid" => pid,
810 ),
811 );
812 }
813 let interval = hyperactor_config::global::get(PROCESS_MEMORY_METRIC_INTERVAL);
814 if !interval.is_zero() {
815 cx.post_after(
816 cx,
817 RepublishIntrospect {
818 emit_memory_metrics: true,
819 },
820 interval,
821 );
822 }
823 }
824 Ok(())
825 }
826}
827
828#[async_trait]
829impl Handler<PySpyDump> for ProcAgent {
830 async fn handle(
831 &mut self,
832 cx: &Context<Self>,
833 message: PySpyDump,
834 ) -> Result<(), anyhow::Error> {
835 PySpyWorker::spawn_and_forward(cx, message.opts, message.result)
836 }
837}
838
839#[async_trait]
840impl Handler<PySpyProfile> for ProcAgent {
841 async fn handle(
842 &mut self,
843 cx: &Context<Self>,
844 message: PySpyProfile,
845 ) -> Result<(), anyhow::Error> {
846 PySpyProfileWorker::spawn_and_forward(cx, message.request, message.result)
847 }
848}
849
850#[async_trait]
851impl Handler<ConfigDump> for ProcAgent {
852 async fn handle(
853 &mut self,
854 cx: &Context<Self>,
855 message: ConfigDump,
856 ) -> Result<(), anyhow::Error> {
857 let entries = hyperactor_config::global::config_entries();
858 let _ = message.result.post(cx, ConfigDumpResult { entries });
861 Ok(())
862 }
863}
864
865#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
869pub struct ActorSpec {
870 pub actor_type: String,
872 pub params_data: Data,
874}
875wirevalue::register_type!(ActorSpec);
876
877#[derive(Debug, Clone, PartialEq, Serialize, Deserialize, Named)]
879pub struct ActorState {
880 pub actor_id: ActorAddr,
882 pub create_rank: usize,
884 pub supervision_events: Vec<ActorSupervisionEvent>,
886}
887wirevalue::register_type!(ActorState);
888
889#[async_trait]
890impl Handler<resource::CreateOrUpdate<ActorSpec>> for ProcAgent {
891 async fn handle(
892 &mut self,
893 cx: &Context<Self>,
894 create_or_update: resource::CreateOrUpdate<ActorSpec>,
895 ) -> anyhow::Result<()> {
896 if self.actor_states.contains_key(&create_or_update.id) {
897 return Ok(());
899 }
900 let create_rank = create_or_update.rank.unwrap();
901 if self.actor_states.values().any(|s| s.has_errors()) {
905 self.actor_states.insert(
906 create_or_update.id.clone(),
907 ActorInstanceState {
908 spawn: Err(anyhow::anyhow!(
909 "Cannot spawn new actors on mesh with supervision events"
910 )),
911 create_rank,
912 stop_initiated: false,
913 supervision_event: None,
914 subscribers: Vec::new(),
915 expiry_time: None,
916 generation: 1,
917 pending_wait_status: Vec::new(),
918 },
919 );
920 return Ok(());
921 }
922
923 let ActorSpec {
924 actor_type,
925 params_data,
926 } = create_or_update.spec;
927 self.actor_states.insert(
928 create_or_update.id.clone(),
929 ActorInstanceState {
930 create_rank,
931 spawn: self
932 .remote
933 .gspawn(
934 &self.proc,
935 &actor_type,
936 create_or_update.id.uid().clone(),
937 params_data,
938 cx.headers().clone(),
939 )
940 .await,
941 stop_initiated: false,
942 supervision_event: None,
943 subscribers: Vec::new(),
944 expiry_time: None,
945 generation: 1,
946 pending_wait_status: Vec::new(),
947 },
948 );
949
950 let _ = self.publish_introspect_properties(cx);
951 Ok(())
952 }
953}
954
955#[async_trait]
956impl Handler<resource::Stop> for ProcAgent {
957 async fn handle(&mut self, cx: &Context<Self>, message: resource::Stop) -> anyhow::Result<()> {
958 let actor_id = match self.actor_states.get_mut(&message.id) {
959 Some(actor_state) => {
960 let id = actor_state.spawn.as_ref().ok().cloned();
961 if id.is_some() && !actor_state.stop_initiated {
962 actor_state.stop_initiated = true;
963 actor_state.generation += 1;
964 actor_state.notify_status_changed(cx, &message.id);
965 id
966 } else {
967 None
968 }
969 }
970 None => None,
971 };
972 if let Some(actor_id) = actor_id {
973 self.stop_actor_by_id(&actor_id, &message.reason);
974 }
975
976 Ok(())
977 }
978}
979
980#[async_trait]
984impl Handler<resource::StopAll> for ProcAgent {
985 async fn handle(
986 &mut self,
987 _cx: &Context<Self>,
988 message: resource::StopAll,
989 ) -> anyhow::Result<()> {
990 self.stopping_all = true;
991
992 let to_stop: Vec<ActorAddr> = self
994 .actor_states
995 .values_mut()
996 .filter_map(|state| {
997 if state.stop_initiated {
998 return None;
999 }
1000 state.stop_initiated = true;
1001 state.spawn.as_ref().ok().cloned()
1002 })
1003 .collect();
1004
1005 for actor_id in &to_stop {
1006 self.stop_actor_by_id(actor_id, &message.reason);
1007 }
1008
1009 if self.all_actors_terminal() {
1011 self.shutdown().await;
1012 }
1013
1014 Ok(())
1015 }
1016}
1017
1018#[async_trait]
1019impl Handler<resource::GetRankStatus> for ProcAgent {
1020 async fn handle(
1021 &mut self,
1022 cx: &Context<Self>,
1023 get_rank_status: resource::GetRankStatus,
1024 ) -> anyhow::Result<()> {
1025 use crate::StatusOverlay;
1026 use crate::resource::Status;
1027
1028 let (rank, status) = match self.actor_states.get(&get_rank_status.id) {
1029 Some(state) => (state.create_rank, state.status()),
1030 None => (usize::MAX, Status::NotExist),
1031 };
1032
1033 let overlay = if rank == usize::MAX {
1036 StatusOverlay::new()
1037 } else {
1038 StatusOverlay::try_from_runs(vec![(rank..(rank + 1), status)])
1039 .expect("valid single-run overlay")
1040 };
1041 get_rank_status.reply.post(cx, overlay);
1042 Ok(())
1043 }
1044}
1045
1046#[async_trait]
1047impl Handler<resource::WaitRankStatus> for ProcAgent {
1048 async fn handle(
1049 &mut self,
1050 cx: &Context<Self>,
1051 msg: resource::WaitRankStatus,
1052 ) -> anyhow::Result<()> {
1053 use crate::StatusOverlay;
1054 use crate::resource::Status;
1055
1056 let (rank, status) = match self.actor_states.get(&msg.id) {
1057 Some(state) => (state.create_rank, state.status()),
1058 None => (usize::MAX, Status::NotExist),
1059 };
1060
1061 if status >= msg.min_status || rank == usize::MAX {
1063 let overlay = if rank == usize::MAX {
1064 StatusOverlay::new()
1065 } else {
1066 StatusOverlay::try_from_runs(vec![(rank..(rank + 1), status)])
1067 .expect("valid single-run overlay")
1068 };
1069 let _ = msg.reply.post(cx, overlay);
1070 return Ok(());
1071 }
1072
1073 if let Some(state) = self.actor_states.get_mut(&msg.id) {
1076 state.pending_wait_status.push((msg.min_status, msg.reply));
1077 }
1078 Ok(())
1079 }
1080}
1081
1082#[async_trait]
1083impl Handler<resource::GetState<ActorState>> for ProcAgent {
1084 async fn handle(
1085 &mut self,
1086 cx: &Context<Self>,
1087 get_state: resource::GetState<ActorState>,
1088 ) -> anyhow::Result<()> {
1089 let state = match self.actor_states.get(&get_state.id) {
1090 Some(instance) => instance.to_state(&get_state.id),
1091 None => resource::State {
1092 id: get_state.id.clone(),
1093 status: resource::Status::NotExist,
1094 state: None,
1095 generation: 0,
1096 timestamp: std::time::SystemTime::now(),
1097 },
1098 };
1099
1100 get_state.reply.post(cx, state);
1101 Ok(())
1102 }
1103}
1104
1105#[async_trait]
1106impl Handler<resource::StreamState<ActorState>> for ProcAgent {
1107 async fn handle(
1108 &mut self,
1109 cx: &Context<Self>,
1110 stream_state: resource::StreamState<ActorState>,
1111 ) -> anyhow::Result<()> {
1112 let state = match self.actor_states.get_mut(&stream_state.id) {
1113 Some(instance) => {
1114 let state = instance.to_state(&stream_state.id);
1115 instance.subscribers.push(stream_state.subscriber.clone());
1116 state
1117 }
1118 None => resource::State {
1119 id: stream_state.id.clone(),
1120 status: resource::Status::NotExist,
1121 state: None,
1122 generation: 0,
1123 timestamp: std::time::SystemTime::now(),
1124 },
1125 };
1126
1127 let mut headers = Flattrs::new();
1129 headers.set(STREAM_STATE_SUBSCRIBER, true);
1130 stream_state
1131 .subscriber
1132 .post_with_headers(cx, headers, state);
1133 Ok(())
1134 }
1135}
1136
1137#[async_trait]
1138impl Handler<resource::KeepaliveGetState<ActorState>> for ProcAgent {
1139 async fn handle(
1140 &mut self,
1141 cx: &Context<Self>,
1142 message: resource::KeepaliveGetState<ActorState>,
1143 ) -> anyhow::Result<()> {
1144 if let Ok(instance_state) =
1146 self.actor_states
1147 .get_mut(&message.get_state.id)
1148 .ok_or_else(|| {
1149 anyhow::anyhow!(
1150 "attempting to register a keepalive for an actor that doesn't exist: {}",
1151 message.get_state.id
1152 )
1153 })
1154 {
1155 instance_state.expiry_time = Some(message.expires_after);
1156 }
1157
1158 <Self as Handler<resource::GetState<ActorState>>>::handle(self, cx, message.get_state).await
1160 }
1161}
1162
1163#[derive(Debug, hyperactor::Handler, hyperactor::HandleClient)]
1166pub struct NewClientInstance {
1167 #[reply]
1168 pub client_instance: PortHandle<Client>,
1169}
1170
1171#[async_trait]
1172impl Handler<NewClientInstance> for ProcAgent {
1173 async fn handle(
1174 &mut self,
1175 cx: &Context<Self>,
1176 NewClientInstance { client_instance }: NewClientInstance,
1177 ) -> anyhow::Result<()> {
1178 let client = self.proc.client("client");
1179 client_instance.post(cx, client);
1180 Ok(())
1181 }
1182}
1183
1184#[derive(Debug, hyperactor::Handler, hyperactor::HandleClient)]
1187pub struct GetProc {
1188 #[reply]
1189 pub proc: PortHandle<Proc>,
1190}
1191
1192#[async_trait]
1193impl Handler<GetProc> for ProcAgent {
1194 async fn handle(
1195 &mut self,
1196 cx: &Context<Self>,
1197 GetProc { proc }: GetProc,
1198 ) -> anyhow::Result<()> {
1199 proc.post(cx, self.proc.clone());
1200 Ok(())
1201 }
1202}
1203
1204#[async_trait]
1205impl Handler<SelfCheck> for ProcAgent {
1206 async fn handle(&mut self, cx: &Context<Self>, _: SelfCheck) -> anyhow::Result<()> {
1207 let Some(duration) = &self.mesh_orphan_timeout else {
1213 return Ok(());
1214 };
1215 let duration = *duration;
1216 let now = std::time::SystemTime::now();
1217
1218 let expired: Vec<(ResourceId, ActorAddr)> = self
1220 .actor_states
1221 .iter()
1222 .filter_map(|(id, state)| {
1223 let expiry = state.expiry_time?;
1224 if now > expiry
1226 && !state.stop_initiated
1227 && let Ok(actor_id) = &state.spawn
1228 {
1229 return Some((id.clone(), actor_id.clone()));
1230 }
1231 None
1232 })
1233 .collect();
1234
1235 if !expired.is_empty() {
1236 tracing::info!(
1237 "stopping {} orphaned actors past their keepalive expiry",
1238 expired.len(),
1239 );
1240 }
1241
1242 for (id, actor_id) in expired {
1243 if let Some(state) = self.actor_states.get_mut(&id) {
1244 state.stop_initiated = true;
1245 }
1246 self.stop_actor_by_id(&actor_id, "orphaned");
1247 }
1248
1249 cx.post_after(cx, SelfCheck::default(), duration);
1251 Ok(())
1252 }
1253}
1254
1255#[cfg(test)]
1256mod tests {
1257 use std::sync::Arc;
1258
1259 use hyperactor::ActorRef;
1260
1261 use super::*;
1262
1263 #[derive(Debug, Default, Serialize, Deserialize)]
1265 #[hyperactor::export(handlers = [])]
1266 struct ExtraActor;
1267 impl hyperactor::Actor for ExtraActor {}
1268 hyperactor::register_spawnable!(ExtraActor);
1269 #[tokio::test]
1283 async fn test_query_child_proc_returns_live_children() {
1284 use hyperactor::Proc;
1285 use hyperactor::actor::ActorStatus;
1286 use hyperactor::channel::ChannelTransport;
1287 use hyperactor::introspect::IntrospectMessage;
1288 use hyperactor::introspect::IntrospectResult;
1289
1290 let proc = Proc::direct(ChannelTransport::Unix.any(), "test_proc".to_string()).unwrap();
1291 let agent_handle = ProcAgent::boot_v1(proc.clone(), None).unwrap();
1292
1293 agent_handle
1295 .status()
1296 .wait_for(|s| matches!(s, ActorStatus::Idle))
1297 .await
1298 .unwrap();
1299
1300 let client_proc = Proc::direct(ChannelTransport::Unix.any(), "client".to_string()).unwrap();
1302 let client = client_proc.client("client");
1303
1304 let agent_id: ActorAddr = proc.proc_addr().actor_addr(PROC_AGENT_ACTOR_NAME);
1305 let port = agent_id.introspect_port();
1306
1307 let query = |client: &hyperactor::Client| {
1310 let (reply_port, reply_rx) = client.open_once_port::<IntrospectResult>();
1311 port.post(
1312 client,
1313 IntrospectMessage::QueryChild {
1314 child_ref: Addr::Proc(proc.proc_addr().clone()),
1315 reply: reply_port.bind(),
1316 },
1317 );
1318 reply_rx
1319 };
1320 let recv = |rx: hyperactor::mailbox::OncePortReceiver<IntrospectResult>| async move {
1321 tokio::time::timeout(std::time::Duration::from_secs(5), rx.recv())
1322 .await
1323 .expect("QueryChild(Proc) timed out — reply never delivered")
1324 .expect("reply channel closed")
1325 };
1326
1327 let payload = recv(query(&client)).await;
1329 let attrs: hyperactor_config::Attrs =
1331 serde_json::from_str(&payload.attrs).expect("valid attrs JSON");
1332 assert_eq!(
1333 attrs.get(crate::introspect::NODE_TYPE).map(String::as_str),
1334 Some("proc"),
1335 "expected node_type=proc in attrs, got {:?}",
1336 payload.attrs
1337 );
1338 assert!(
1339 payload
1340 .children
1341 .iter()
1342 .any(|c| c.to_string().contains(PROC_AGENT_ACTOR_NAME)),
1343 "initial children {:?} should contain proc_agent",
1344 payload.children
1345 );
1346 let initial_count = payload.children.len();
1347
1348 proc.spawn_with_label("extra_actor", ExtraActor);
1352
1353 let payload2 = recv(query(&client)).await;
1355 let attrs2: hyperactor_config::Attrs =
1356 serde_json::from_str(&payload2.attrs).expect("valid attrs JSON");
1357 assert_eq!(
1358 attrs2.get(crate::introspect::NODE_TYPE).map(String::as_str),
1359 Some("proc"),
1360 "expected node_type=proc in attrs, got {:?}",
1361 payload2.attrs
1362 );
1363 assert!(
1364 payload2
1365 .children
1366 .iter()
1367 .any(|c| c.to_string().contains("extra_actor")),
1368 "after direct spawn, children {:?} should contain extra_actor",
1369 payload2.children
1370 );
1371 assert!(
1372 payload2.children.len() > initial_count,
1373 "expected at least {} children after direct spawn, got {:?}",
1374 initial_count + 1,
1375 payload2.children
1376 );
1377 }
1378
1379 #[tokio::test]
1386 async fn test_rapid_spawn_stop_does_not_stall_proc_agent() {
1387 use std::sync::Arc;
1388 use std::sync::atomic::AtomicUsize;
1389 use std::sync::atomic::Ordering;
1390
1391 use hyperactor::Proc;
1392 use hyperactor::actor::ActorStatus;
1393 use hyperactor::channel::ChannelTransport;
1394 use hyperactor::introspect::IntrospectMessage;
1395 use hyperactor::introspect::IntrospectResult;
1396
1397 let proc = Proc::direct(ChannelTransport::Unix.any(), "test_proc".to_string()).unwrap();
1398 let agent_handle = ProcAgent::boot_v1(proc.clone(), None).unwrap();
1399
1400 agent_handle
1401 .status()
1402 .wait_for(|s| matches!(s, ActorStatus::Idle))
1403 .await
1404 .unwrap();
1405
1406 let client_proc = Proc::direct(ChannelTransport::Unix.any(), "client".to_string()).unwrap();
1407 let client = client_proc.client("client");
1408
1409 let agent_id: ActorAddr = proc.proc_addr().actor_addr(PROC_AGENT_ACTOR_NAME);
1410 let port = agent_id.introspect_port();
1411
1412 let query_client_proc =
1414 Proc::direct(ChannelTransport::Unix.any(), "query_client".to_string()).unwrap();
1415 let query_client = query_client_proc.client("qc");
1416 let query_port = port.clone();
1417 let query_proc_id = proc.proc_addr().clone();
1418 let query_count = Arc::new(AtomicUsize::new(0));
1419 let query_count_clone = query_count.clone();
1420 let query_task = tokio::spawn(async move {
1421 loop {
1422 let (reply_port, reply_rx) = query_client.open_once_port::<IntrospectResult>();
1423 query_port.post(
1424 &query_client,
1425 IntrospectMessage::QueryChild {
1426 child_ref: Addr::Proc(query_proc_id.clone()),
1427 reply: reply_port.bind(),
1428 },
1429 );
1430 match tokio::time::timeout(std::time::Duration::from_secs(2), reply_rx.recv()).await
1431 {
1432 Ok(Ok(_)) => {
1433 query_count_clone.fetch_add(1, Ordering::Relaxed);
1434 }
1435 _ => {} }
1437 tokio::time::sleep(std::time::Duration::from_millis(10)).await;
1438 }
1439 });
1440
1441 const ITERATIONS: usize = 200;
1443 let mut completed = 0usize;
1444 let result = tokio::time::timeout(std::time::Duration::from_secs(30), async {
1445 for i in 0..ITERATIONS {
1446 let name = format!("churn_{}", i);
1447 let handle = proc.spawn_with_label(&name, ExtraActor);
1448 let actor_id = handle.actor_addr().clone();
1449 if let Some(mut status) = proc.stop_actor(actor_id.id(), "churn".to_string()) {
1450 let _ = tokio::time::timeout(
1451 std::time::Duration::from_secs(5),
1452 status.wait_for(ActorStatus::is_terminal),
1453 )
1454 .await;
1455 }
1456 completed += 1;
1457 }
1458 })
1459 .await;
1460
1461 query_task.abort();
1462 let _ = query_task.await; assert!(
1465 result.is_ok(),
1466 "spawn/stop loop stalled after {completed}/{ITERATIONS} iterations — \
1467 DashMap convoy starvation likely"
1468 );
1469 assert_eq!(
1470 completed, ITERATIONS,
1471 "expected {ITERATIONS} completed iterations, got {completed}"
1472 );
1473 assert!(
1474 query_count.load(Ordering::Relaxed) > 0,
1475 "concurrent QueryChild queries never succeeded — query task may not have run"
1476 );
1477
1478 let (reply_port, reply_rx) = client.open_once_port::<IntrospectResult>();
1480 port.post(
1481 &client,
1482 IntrospectMessage::QueryChild {
1483 child_ref: Addr::Proc(proc.proc_addr().clone()),
1484 reply: reply_port.bind(),
1485 },
1486 );
1487 let final_payload =
1488 tokio::time::timeout(std::time::Duration::from_secs(5), reply_rx.recv())
1489 .await
1490 .expect("final QueryChild timed out")
1491 .expect("final QueryChild channel closed");
1492 let attrs: hyperactor_config::Attrs =
1493 serde_json::from_str(&final_payload.attrs).expect("valid attrs JSON");
1494 assert_eq!(
1495 attrs.get(crate::introspect::NODE_TYPE).map(String::as_str),
1496 Some("proc"),
1497 );
1498 }
1499
1500 #[tokio::test]
1501 async fn test_stream_state_and_unsubscribe() {
1502 use hyperactor::Proc;
1503 use hyperactor::actor::ActorStatus;
1504 use hyperactor::channel::ChannelTransport;
1505
1506 use crate::resource::CreateOrUpdateClient;
1507 use crate::resource::GetStateClient;
1508 use crate::resource::StopClient;
1509 use crate::resource::StreamStateClient;
1510
1511 let proc = Proc::direct(ChannelTransport::Unix.any(), "test_proc".to_string()).unwrap();
1512 let agent_handle = ProcAgent::boot_v1(proc.clone(), None).unwrap();
1513 agent_handle
1514 .status()
1515 .wait_for(|s| matches!(s, ActorStatus::Idle))
1516 .await
1517 .unwrap();
1518
1519 let client = proc.client("client");
1520 let agent_ref: ActorRef<ProcAgent> = agent_handle.bind();
1521
1522 let actor_type = hyperactor::actor::remote::Remote::collect()
1523 .name_of::<ExtraActor>()
1524 .unwrap()
1525 .to_string();
1526 let actor_params =
1527 bincode::serde::encode_to_vec(&ExtraActor, bincode::config::legacy()).unwrap();
1528 let actor_name = ResourceId::singleton(hyperactor::id::Label::new("test-actor").unwrap());
1529
1530 agent_ref
1532 .create_or_update(
1533 &client,
1534 actor_name.clone(),
1535 resource::Rank::new(0),
1536 ActorSpec {
1537 actor_type: actor_type.clone(),
1538 params_data: actor_params.clone(),
1539 },
1540 )
1541 .await
1542 .unwrap();
1543
1544 let (sub_port, mut sub_rx) = client.open_port::<resource::State<ActorState>>();
1546 agent_ref
1547 .stream_state(&client, actor_name.clone(), sub_port.bind())
1548 .await
1549 .unwrap();
1550
1551 let initial = sub_rx.recv().await.expect("subscriber channel error");
1553 assert_eq!(initial.status, resource::Status::Running);
1554 assert!(initial.state.is_some());
1555
1556 agent_ref
1558 .stop(&client, actor_name.clone(), "test".to_string())
1559 .await
1560 .unwrap();
1561
1562 let stopping = sub_rx.recv().await.expect("subscriber channel error");
1563 assert_eq!(stopping.status, resource::Status::Stopping);
1564
1565 let stopped = sub_rx.recv().await.expect("subscriber channel error");
1567 assert_eq!(stopped.status, resource::Status::Stopped);
1568
1569 let actor_name_2 =
1571 ResourceId::singleton(hyperactor::id::Label::new("test-actor-2").unwrap());
1572 agent_ref
1573 .create_or_update(
1574 &client,
1575 actor_name_2.clone(),
1576 resource::Rank::new(1),
1577 ActorSpec {
1578 actor_type: actor_type.clone(),
1579 params_data: actor_params.clone(),
1580 },
1581 )
1582 .await
1583 .unwrap();
1584
1585 let (sub_port_2, mut sub_rx_2) = client.open_port::<resource::State<ActorState>>();
1586 agent_ref
1587 .stream_state(&client, actor_name_2.clone(), sub_port_2.bind())
1588 .await
1589 .unwrap();
1590
1591 let initial_2 = sub_rx_2.recv().await.expect("subscriber 2 channel error");
1592 assert_eq!(initial_2.status, resource::Status::Running);
1593
1594 drop(sub_rx_2);
1596
1597 agent_ref
1601 .stop(
1602 &client,
1603 actor_name_2.clone(),
1604 "test unsubscribe".to_string(),
1605 )
1606 .await
1607 .unwrap();
1608
1609 let (sub_port_3, mut sub_rx_3) = client.open_port::<resource::State<ActorState>>();
1611 agent_ref
1612 .stream_state(&client, actor_name_2.clone(), sub_port_3.bind())
1613 .await
1614 .unwrap();
1615 loop {
1616 let state = sub_rx_3.recv().await.expect("subscriber 3 channel error");
1617 if state.status.is_terminating() {
1618 break;
1619 }
1620 }
1621
1622 let state = agent_ref
1624 .get_state(&client, actor_name_2.clone())
1625 .await
1626 .unwrap();
1627 assert!(
1628 state.status.is_terminating(),
1629 "expected terminating status, got {:?}",
1630 state.status,
1631 );
1632 }
1633
1634 #[derive(Debug, Default, Serialize, Deserialize)]
1640 #[hyperactor::export(handlers = [BlockMsg])]
1641 struct BlockActor {
1642 #[serde(skip)]
1643 gate: Option<Arc<tokio::sync::Notify>>,
1644 }
1645 impl hyperactor::Actor for BlockActor {}
1646
1647 #[derive(
1648 Debug,
1649 Clone,
1650 Serialize,
1651 Deserialize,
1652 Named,
1653 hyperactor::Handler,
1654 hyperactor::HandleClient
1655 )]
1656 enum BlockMsg {
1657 Block(),
1659 Noop(),
1661 }
1662 wirevalue::register_type!(BlockMsg);
1663
1664 #[async_trait::async_trait]
1665 #[hyperactor::handle(BlockMsg)]
1666 impl BlockMsgHandler for BlockActor {
1667 async fn block(&mut self, _cx: &hyperactor::Context<Self>) -> Result<(), anyhow::Error> {
1668 if let Some(gate) = &self.gate {
1669 gate.notified().await;
1670 }
1671 Ok(())
1672 }
1673 async fn noop(&mut self, _cx: &hyperactor::Context<Self>) -> Result<(), anyhow::Error> {
1674 Ok(())
1675 }
1676 }
1677
1678 #[tokio::test]
1686 async fn test_query_child_proc_queue_depth_under_pressure() {
1687 use hyperactor::Proc;
1688 use hyperactor::actor::ActorStatus;
1689 use hyperactor::channel::ChannelTransport;
1690 use hyperactor::introspect::IntrospectMessage;
1691 use hyperactor::introspect::IntrospectResult;
1692
1693 let proc = Proc::direct(ChannelTransport::Unix.any(), "qd_proc".to_string()).unwrap();
1694 let agent_handle = ProcAgent::boot_v1(proc.clone(), None).unwrap();
1695
1696 agent_handle
1697 .status()
1698 .wait_for(|s| matches!(s, ActorStatus::Idle))
1699 .await
1700 .unwrap();
1701
1702 let client_proc =
1703 Proc::direct(ChannelTransport::Unix.any(), "qd_client".to_string()).unwrap();
1704 let client = client_proc.client("client");
1705
1706 let gate = Arc::new(tokio::sync::Notify::new());
1708 let blocker = proc.spawn(BlockActor {
1709 gate: Some(Arc::clone(&gate)),
1710 });
1711
1712 blocker.block(&client).await.unwrap();
1714 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1716 blocker.noop(&client).await.unwrap();
1717 blocker.noop(&client).await.unwrap();
1718
1719 let agent_id: ActorAddr = proc.proc_addr().actor_addr(PROC_AGENT_ACTOR_NAME);
1722 let port = agent_id.introspect_port();
1723
1724 let deadline = tokio::time::Instant::now() + std::time::Duration::from_secs(5);
1726 loop {
1727 let (reply_port, reply_rx) = client.open_once_port::<IntrospectResult>();
1728 port.post(
1729 &client,
1730 IntrospectMessage::QueryChild {
1731 child_ref: Addr::Proc(proc.proc_addr().clone()),
1732 reply: reply_port.bind(),
1733 },
1734 );
1735 let payload = tokio::time::timeout(std::time::Duration::from_secs(3), reply_rx.recv())
1736 .await
1737 .expect("QueryChild timed out")
1738 .expect("reply channel closed");
1739
1740 let attrs: hyperactor_config::Attrs =
1741 serde_json::from_str(&payload.attrs).expect("valid attrs JSON");
1742
1743 let total = attrs
1744 .get(crate::introspect::ACTOR_WORK_QUEUE_DEPTH_TOTAL)
1745 .copied()
1746 .unwrap_or(0);
1747 let max = attrs
1748 .get(crate::introspect::ACTOR_WORK_QUEUE_DEPTH_MAX)
1749 .copied()
1750 .unwrap_or(0);
1751
1752 if total > 0 {
1753 assert!(max > 0, "max should be > 0 when total is {total}");
1754 assert!(max <= total, "PD-1: max ({max}) <= total ({total})");
1755 break;
1756 }
1757
1758 assert!(
1759 tokio::time::Instant::now() < deadline,
1760 "timed out waiting for non-zero queue depth in QueryChild(Proc)",
1761 );
1762 tokio::time::sleep(std::time::Duration::from_millis(50)).await;
1763 }
1764
1765 gate.notify_one();
1767 }
1768}