Skip to main content

robot_bus/runtime/node/
action_clients.rs

1use std::marker::PhantomData;
2use std::sync::atomic::{AtomicBool, Ordering};
3use std::sync::mpsc::{self, Receiver, Sender};
4use std::sync::{Arc, Mutex};
5use std::thread;
6use std::time::{Duration, Instant};
7
8use prost::Message;
9
10use crate::action_bus::{ActionClient as BusActionClient, ActionKind, ActionMessage};
11use crate::errors::{BusError, Result, parse_error_body};
12use crate::runtime::console_ready::{self, ReadyKind};
13use crate::runtime::topology_register::TopologyEndpointGuard;
14#[cfg(feature = "ws")]
15use crate::runtime::ws_runtime::WsClientContext;
16use crate::typed::{Action, ActionOutcome};
17use crate::zmq_helpers::HighWaterMark;
18
19/// Action server handle returned by [`Node::create_action_server`] /
20/// [`Node::create_action_server_raw`].
21#[derive(Clone, Debug)]
22pub struct NodeActionServer {
23    pub(super) id: u64,
24    pub(super) action_name: String,
25}
26
27impl NodeActionServer {
28    pub fn id(&self) -> u64 {
29        self.id
30    }
31
32    pub fn action_name(&self) -> &str {
33        &self.action_name
34    }
35}
36
37/// Callback invoked for each action feedback as it arrives.
38pub type RawActionFeedbackCallback = Arc<dyn Fn(&ActionMessage) + Send + Sync + 'static>;
39
40pub(super) fn spawn_zmq_goal(
41    context: zmq::Context,
42    endpoint: String,
43    action_name: String,
44    body: Vec<u8>,
45    requested_goal_id: Option<String>,
46    timeout: Option<Duration>,
47    hwm: HighWaterMark,
48    feedback_callback: Option<RawActionFeedbackCallback>,
49) -> Result<RawGoalHandle> {
50    let (ready_tx, ready_rx) = mpsc::sync_channel(1);
51    let (event_tx, event_rx) = mpsc::channel();
52    let (command_tx, command_rx) = mpsc::channel();
53    let thread_action_name = action_name.clone();
54
55    thread::Builder::new()
56        .name(format!("action-{}", action_name))
57        .spawn(move || {
58            let client = match BusActionClient::with_context_hwm(&context, Some(&endpoint), hwm) {
59                Ok(client) => client,
60                Err(err) => {
61                    let _ = ready_tx.send(Err(err));
62                    return;
63                }
64            };
65            let goal_id = match client.submit_goal(
66                &thread_action_name,
67                &body,
68                requested_goal_id.as_deref(),
69            ) {
70                Ok(goal_id) => goal_id,
71                Err(err) => {
72                    let _ = ready_tx.send(Err(err));
73                    return;
74                }
75            };
76            if ready_tx.send(Ok(goal_id.clone())).is_err() {
77                return;
78            }
79
80            let deadline = timeout.map(|duration| Instant::now() + duration);
81            loop {
82                while let Ok(command) = command_rx.try_recv() {
83                    match command {
84                        GoalCommand::Cancel(body) => {
85                            if let Err(err) =
86                                client.submit_cancel(&thread_action_name, &goal_id, &body)
87                            {
88                                let _ = event_tx.send(Err(err));
89                                return;
90                            }
91                        }
92                    }
93                }
94
95                if deadline.is_some_and(|deadline| Instant::now() >= deadline) {
96                    let _ = client.submit_cancel(&thread_action_name, &goal_id, b"");
97                    let _ = event_tx.send(Err(BusError::Timeout(format!(
98                        "action client timed out after {}s",
99                        timeout.unwrap_or_default().as_secs_f64()
100                    ))));
101                    return;
102                }
103
104                let poll_timeout = deadline
105                    .map(|deadline| {
106                        deadline
107                            .saturating_duration_since(Instant::now())
108                            .min(Duration::from_millis(20))
109                    })
110                    .unwrap_or(Duration::from_millis(20));
111                let message = match client.recv_message(Some(poll_timeout)) {
112                    Ok(message) => message,
113                    Err(BusError::Timeout(_)) => continue,
114                    Err(err) => {
115                        let _ = event_tx.send(Err(err));
116                        return;
117                    }
118                };
119                if message.action_name != thread_action_name || message.goal_id != goal_id {
120                    let _ = event_tx.send(Err(BusError::Protocol(format!(
121                        "unexpected message for {:?}/{:?}",
122                        message.action_name, message.goal_id
123                    ))));
124                    return;
125                }
126                if message.kind == ActionKind::Feedback {
127                    if let Some(callback) = &feedback_callback {
128                        callback(&message);
129                    }
130                }
131                let done = message.kind == ActionKind::Result;
132                if done {
133                    if let Some(err) = parse_error_body(&message.body) {
134                        let _ = event_tx.send(Err(err));
135                        return;
136                    }
137                }
138                if event_tx.send(Ok(message)).is_err() || done {
139                    return;
140                }
141            }
142        })
143        .map_err(|err| BusError::Protocol(format!("spawn action thread: {err}")))?;
144
145    let goal_id = ready_rx
146        .recv()
147        .map_err(|_| BusError::Protocol("action thread ended before submitting goal".into()))??;
148    Ok(RawGoalHandle {
149        inner: Arc::new(GoalHandleCore {
150            action_name,
151            goal_id,
152            events: Mutex::new(event_rx),
153            messages: Mutex::new(Vec::new()),
154            control: GoalControl::Zmq(command_tx),
155            completed: AtomicBool::new(false),
156        }),
157    })
158}
159
160pub(super) enum GoalControl {
161    Zmq(Sender<GoalCommand>),
162    #[cfg(feature = "ws")]
163    Ws(crate::runtime::ws_runtime::WsCancelHandle),
164}
165
166pub(super) enum GoalCommand {
167    Cancel(Vec<u8>),
168}
169
170pub(super) struct GoalHandleCore {
171    action_name: String,
172    goal_id: String,
173    events: Mutex<Receiver<Result<ActionMessage>>>,
174    messages: Mutex<Vec<ActionMessage>>,
175    control: GoalControl,
176    completed: AtomicBool,
177}
178
179impl GoalHandleCore {
180    fn wait_result(&self) -> Result<ActionMessage> {
181        if let Some(result) = self
182            .messages
183            .lock()
184            .map_err(|_| BusError::Protocol("action messages mutex poisoned".into()))?
185            .iter()
186            .find(|message| message.kind == ActionKind::Result)
187            .cloned()
188        {
189            return Ok(result);
190        }
191
192        loop {
193            let event = self
194                .events
195                .lock()
196                .map_err(|_| BusError::Protocol("action event mutex poisoned".into()))?
197                .recv()
198                .map_err(|_| {
199                    BusError::Protocol(format!(
200                        "action '{}' goal '{}' ended without RESULT",
201                        self.action_name, self.goal_id
202                    ))
203                })??;
204            let done = event.kind == ActionKind::Result;
205            self.messages
206                .lock()
207                .map_err(|_| BusError::Protocol("action messages mutex poisoned".into()))?
208                .push(event.clone());
209            if done {
210                self.completed.store(true, Ordering::Release);
211                return Ok(event);
212            }
213        }
214    }
215
216    fn collect(&self) -> Result<Vec<ActionMessage>> {
217        self.wait_result()?;
218        self.messages
219            .lock()
220            .map(|messages| messages.clone())
221            .map_err(|_| BusError::Protocol("action messages mutex poisoned".into()))
222    }
223
224    fn cancel(&self, body: &[u8]) -> Result<()> {
225        match &self.control {
226            GoalControl::Zmq(commands) => commands
227                .send(GoalCommand::Cancel(body.to_vec()))
228                .map_err(|_| BusError::Closed),
229            #[cfg(feature = "ws")]
230            GoalControl::Ws(abort) => {
231                abort.abort();
232                Ok(())
233            }
234        }
235    }
236}
237
238impl Drop for GoalHandleCore {
239    fn drop(&mut self) {
240        if self.completed.load(Ordering::Acquire) {
241            return;
242        }
243        match &self.control {
244            GoalControl::Zmq(commands) => {
245                let _ = commands.send(GoalCommand::Cancel(Vec::new()));
246            }
247            #[cfg(feature = "ws")]
248            GoalControl::Ws(abort) => abort.abort(),
249        }
250    }
251}
252
253/// Live handle for one raw (opaque bytes) action goal.
254#[derive(Clone)]
255pub struct RawGoalHandle {
256    pub(super) inner: Arc<GoalHandleCore>,
257}
258
259impl RawGoalHandle {
260    pub fn goal_id(&self) -> &str {
261        &self.inner.goal_id
262    }
263
264    pub fn action_name(&self) -> &str {
265        &self.inner.action_name
266    }
267
268    pub fn wait_result(&self) -> Result<ActionMessage> {
269        self.inner.wait_result()
270    }
271
272    pub fn collect(&self) -> Result<Vec<ActionMessage>> {
273        self.inner.collect()
274    }
275
276    /// Best-effort cancellation. This does not wait for server acknowledgement.
277    pub fn cancel(&self) -> Result<()> {
278        self.inner.cancel(&[])
279    }
280
281    /// Best-effort cancellation with an opaque ZMQ cancel payload.
282    ///
283    /// On native gRPC this aborts the response stream (body ignored). Browser
284    /// WebSocket clients send an explicit CANCEL frame instead.
285    pub fn cancel_with_body(&self, body: &[u8]) -> Result<()> {
286        self.inner.cancel(body)
287    }
288}
289
290/// Live handle for one typed action goal.
291pub struct GoalHandle<A: Action> {
292    pub(super) inner: RawGoalHandle,
293    pub(super) _marker: PhantomData<A>,
294}
295
296impl<A: Action> Clone for GoalHandle<A> {
297    fn clone(&self) -> Self {
298        Self {
299            inner: self.inner.clone(),
300            _marker: PhantomData,
301        }
302    }
303}
304
305impl<A: Action> GoalHandle<A> {
306    pub fn goal_id(&self) -> &str {
307        self.inner.goal_id()
308    }
309
310    pub fn action_name(&self) -> &str {
311        self.inner.action_name()
312    }
313
314    pub fn wait_result(&self) -> Result<A::Result> {
315        let message = self.inner.wait_result()?;
316        A::Result::decode(message.body.as_slice()).map_err(|err| {
317            BusError::Protocol(format!(
318                "action '{}' result decode failed: {err}",
319                self.action_name()
320            ))
321        })
322    }
323
324    pub fn cancel(&self) -> Result<()> {
325        self.inner.cancel()
326    }
327}
328
329/// Raw (opaque bytes) action client from [`Node::create_action_client_raw`].
330pub struct NodeActionClientRaw {
331    pub(super) inner: ActionClientInner,
332    pub(super) action_name: String,
333    pub(super) console_url: Option<String>,
334    pub(super) _topology: Option<Arc<TopologyEndpointGuard>>,
335}
336
337pub(super) enum ActionClientInner {
338    Zmq {
339        context: zmq::Context,
340        endpoint: String,
341        hwm: Mutex<HighWaterMark>,
342    },
343    #[cfg(feature = "ws")]
344    Ws(WsClientContext),
345}
346
347impl NodeActionClientRaw {
348    pub fn action_name(&self) -> &str {
349        &self.action_name
350    }
351
352    /// Best-effort: console reports `workers > 0` for this action.
353    pub fn action_server_is_ready(&self) -> bool {
354        console_ready::is_ready(
355            self.console_url.as_deref(),
356            ReadyKind::Action,
357            &self.action_name,
358        )
359    }
360
361    /// Poll until [`action_server_is_ready`](Self::action_server_is_ready) or `timeout`.
362    pub fn wait_for_action_server(&self, timeout: Option<Duration>) -> bool {
363        console_ready::wait_until_ready(
364            self.console_url.as_deref(),
365            ReadyKind::Action,
366            &self.action_name,
367            timeout,
368        )
369    }
370
371    pub fn send_goal(
372        &self,
373        body: &[u8],
374        goal_id: Option<&str>,
375        timeout: Option<Duration>,
376        feedback_callback: Option<RawActionFeedbackCallback>,
377    ) -> Result<RawGoalHandle> {
378        match &self.inner {
379            ActionClientInner::Zmq {
380                context,
381                endpoint,
382                hwm,
383            } => {
384                let hwm = *hwm
385                    .lock()
386                    .map_err(|_| BusError::Protocol("action HWM mutex poisoned".into()))?;
387                spawn_zmq_goal(
388                    context.clone(),
389                    endpoint.clone(),
390                    self.action_name.clone(),
391                    body.to_vec(),
392                    goal_id.map(str::to_string),
393                    timeout,
394                    hwm,
395                    feedback_callback,
396                )
397            }
398            #[cfg(feature = "ws")]
399            ActionClientInner::Ws(ctx) => ctx
400                .send_goal(&self.action_name, body, goal_id, timeout, feedback_callback)
401                .map(|session| RawGoalHandle {
402                    inner: Arc::new(GoalHandleCore {
403                        action_name: self.action_name.clone(),
404                        goal_id: session.goal_id,
405                        events: Mutex::new(session.events),
406                        messages: Mutex::new(Vec::new()),
407                        control: GoalControl::Ws(session.abort),
408                        completed: AtomicBool::new(false),
409                    }),
410                }),
411        }
412    }
413
414    /// Compatibility helper that waits for and collects FEEDBACK/RESULT.
415    pub fn send_goal_and_wait(
416        &self,
417        body: &[u8],
418        goal_id: Option<&str>,
419        timeout: Option<Duration>,
420    ) -> Result<Vec<ActionMessage>> {
421        self.send_goal(body, goal_id, timeout, None)?.collect()
422    }
423
424    /// Alias for [`send_goal_and_wait`](Self::send_goal_and_wait).
425    pub fn collect(
426        &self,
427        body: &[u8],
428        goal_id: Option<&str>,
429        timeout: Option<Duration>,
430    ) -> Result<Vec<ActionMessage>> {
431        self.send_goal_and_wait(body, goal_id, timeout)
432    }
433
434    pub fn high_water_mark(&self) -> Result<HighWaterMark> {
435        match &self.inner {
436            ActionClientInner::Zmq { hwm, .. } => hwm
437                .lock()
438                .map(|hwm| *hwm)
439                .map_err(|_| BusError::Protocol("action HWM mutex poisoned".into())),
440            #[cfg(feature = "ws")]
441            ActionClientInner::Ws(_) => Err(BusError::Protocol(
442                "high_water_mark is not available in WebSocket RPC node mode".into(),
443            )),
444        }
445    }
446
447    pub fn set_high_water_mark(&self, hwm: HighWaterMark) -> Result<()> {
448        match &self.inner {
449            ActionClientInner::Zmq { hwm: current, .. } => {
450                *current
451                    .lock()
452                    .map_err(|_| BusError::Protocol("action HWM mutex poisoned".into()))? = hwm;
453                Ok(())
454            }
455            #[cfg(feature = "ws")]
456            ActionClientInner::Ws(_) => Err(BusError::Protocol(
457                "set_high_water_mark is not available in WebSocket RPC node mode".into(),
458            )),
459        }
460    }
461}
462
463/// Typed action client returned by [`Node::create_action_client`] (ROS 2 style).
464pub struct NodeActionClient<A: Action> {
465    pub(super) inner: NodeActionClientRaw,
466    pub(super) _marker: PhantomData<A>,
467}
468
469impl<A: Action> NodeActionClient<A> {
470    pub fn action_name(&self) -> &str {
471        self.inner.action_name()
472    }
473
474    pub fn action_server_is_ready(&self) -> bool {
475        self.inner.action_server_is_ready()
476    }
477
478    pub fn wait_for_action_server(&self, timeout: Option<Duration>) -> bool {
479        self.inner.wait_for_action_server(timeout)
480    }
481
482    pub fn send_goal(
483        &self,
484        goal: &A::Goal,
485        goal_id: Option<&str>,
486        timeout: Option<Duration>,
487        feedback_callback: Option<Arc<dyn Fn(A::Feedback) + Send + Sync + 'static>>,
488    ) -> Result<GoalHandle<A>> {
489        let action_name = self.action_name().to_string();
490        let raw_callback = feedback_callback.map(|callback| {
491            Arc::new(move |message: &ActionMessage| {
492                match A::Feedback::decode(message.body.as_slice()) {
493                    Ok(feedback) => callback(feedback),
494                    Err(err) => {
495                        log::warn!("action '{}' feedback decode failed: {err}", action_name)
496                    }
497                }
498            }) as RawActionFeedbackCallback
499        });
500        let inner = self
501            .inner
502            .send_goal(&goal.encode_to_vec(), goal_id, timeout, raw_callback)?;
503        Ok(GoalHandle {
504            inner,
505            _marker: PhantomData,
506        })
507    }
508
509    /// Compatibility helper that waits for a result and collects feedback.
510    pub fn send_goal_and_wait(
511        &self,
512        goal: &A::Goal,
513        goal_id: Option<&str>,
514        timeout: Option<Duration>,
515    ) -> Result<ActionOutcome<A>> {
516        let messages = self
517            .inner
518            .send_goal_and_wait(&goal.encode_to_vec(), goal_id, timeout)?;
519        let mut feedbacks = Vec::new();
520        let mut result = None;
521        for msg in messages {
522            match msg.kind {
523                ActionKind::Feedback => {
524                    let fb = A::Feedback::decode(msg.body.as_slice()).map_err(|err| {
525                        BusError::Protocol(format!(
526                            "action '{}' feedback decode failed: {err}",
527                            self.action_name()
528                        ))
529                    })?;
530                    feedbacks.push(fb);
531                }
532                ActionKind::Result => {
533                    let res = A::Result::decode(msg.body.as_slice()).map_err(|err| {
534                        BusError::Protocol(format!(
535                            "action '{}' result decode failed: {err}",
536                            self.action_name()
537                        ))
538                    })?;
539                    result = Some(res);
540                }
541                ActionKind::Goal | ActionKind::Cancel => {}
542            }
543        }
544        let result = result.ok_or_else(|| {
545            BusError::Protocol(format!(
546                "action '{}' completed without RESULT",
547                self.action_name()
548            ))
549        })?;
550        Ok(ActionOutcome { feedbacks, result })
551    }
552
553    pub fn high_water_mark(&self) -> Result<HighWaterMark> {
554        self.inner.high_water_mark()
555    }
556
557    pub fn set_high_water_mark(&self, hwm: HighWaterMark) -> Result<()> {
558        self.inner.set_high_water_mark(hwm)
559    }
560}