robot_bus/runtime/
registrations.rs

1//! Registration types for sockets managed by [`super::Executor`].
2
3use std::collections::HashMap;
4use std::sync::atomic::{AtomicBool, Ordering};
5use std::sync::{Arc, Mutex};
6use std::time::{Duration, Instant};
7
8use uuid::Uuid;
9use zmq::{Context, Socket, SocketType};
10
11use crate::runtime::callback_group::CallbackGroup;
12use crate::runtime::queues::ActionMessageCallback;
13use crate::zmq_helpers::{HighWaterMark, apply_action_options_with, apply_rpc_options_with};
14
15pub type MessageCallback = Arc<dyn Fn(&[u8]) + Send + Sync>;
16pub type ServiceHandler = Arc<dyn Fn(&[u8]) -> Vec<u8> + Send + Sync>;
17pub type ActionGoalHandler = Arc<dyn Fn(&[u8]) -> Vec<(String, Vec<u8>)> + Send + Sync>;
18/// Live action handler: may call [`ActionGoalContext::publish_feedback`] before returning
19/// the RESULT body. Poll [`ActionGoalContext::cancel_requested`] to honor CANCEL.
20pub type ActionGoalLiveHandler = Arc<dyn Fn(&[u8], &ActionGoalContext) -> Vec<u8> + Send + Sync>;
21
22/// Per-goal context for a live action server handler.
23#[derive(Clone)]
24pub struct ActionGoalContext {
25    goal_id: String,
26    cancel: Arc<AtomicBool>,
27    on_feedback: Arc<dyn Fn(&[u8]) + Send + Sync>,
28}
29
30impl ActionGoalContext {
31    pub fn new(
32        goal_id: impl Into<String>,
33        cancel: Arc<AtomicBool>,
34        on_feedback: Arc<dyn Fn(&[u8]) + Send + Sync>,
35    ) -> Self {
36        Self {
37            goal_id: goal_id.into(),
38            cancel,
39            on_feedback,
40        }
41    }
42
43    pub fn goal_id(&self) -> &str {
44        &self.goal_id
45    }
46
47    pub fn cancel_requested(&self) -> bool {
48        self.cancel.load(Ordering::SeqCst)
49    }
50
51    pub fn publish_feedback(&self, body: &[u8]) {
52        (self.on_feedback)(body);
53    }
54}
55
56/// Wrap a batch handler so FEEDBACK/RESULT are emitted after it returns.
57pub fn wrap_batch_action_handler(handler: ActionGoalHandler) -> ActionGoalLiveHandler {
58    Arc::new(move |body, ctx| {
59        let mut result = Vec::new();
60        for (phase, chunk) in handler(body) {
61            if phase.eq_ignore_ascii_case("FEEDBACK") {
62                ctx.publish_feedback(&chunk);
63            } else if phase.eq_ignore_ascii_case("RESULT") {
64                result = chunk;
65            }
66        }
67        result
68    })
69}
70
71pub enum RegistrationKind {
72    Sub,
73    Service,
74    Action,
75    ActionClient,
76}
77
78pub struct SubRegistration {
79    pub socket: Socket,
80    pub endpoint: String,
81}
82
83impl SubRegistration {
84    pub fn kind(&self) -> RegistrationKind {
85        RegistrationKind::Sub
86    }
87}
88
89pub struct ServiceRegistration {
90    pub id: u64,
91    pub socket: Socket,
92    pub service_name: String,
93    pub handler: ServiceHandler,
94    pub callback_group: CallbackGroup,
95    pub identity: Vec<u8>,
96    pub heartbeat_interval: Duration,
97    pub last_heartbeat: Instant,
98}
99
100impl ServiceRegistration {
101    pub fn create(
102        id: u64,
103        context: &Context,
104        service_name: &str,
105        handler: ServiceHandler,
106        callback_group: CallbackGroup,
107        endpoint: &str,
108        identity: Option<&str>,
109        heartbeat_interval_ms: u64,
110        hwm: HighWaterMark,
111    ) -> crate::errors::Result<Self> {
112        let socket = context.socket(SocketType::DEALER)?;
113        apply_rpc_options_with(&socket, hwm)?;
114        let worker_id = identity
115            .map(str::to_string)
116            .unwrap_or_else(|| format!("worker-{}", &Uuid::new_v4().simple().to_string()[..8]));
117        let identity = worker_id.into_bytes();
118        socket.set_identity(&identity)?;
119        socket.connect(endpoint)?;
120        let mut reg = Self {
121            id,
122            socket,
123            service_name: service_name.to_string(),
124            handler,
125            callback_group,
126            identity,
127            heartbeat_interval: Duration::from_millis(heartbeat_interval_ms),
128            last_heartbeat: Instant::now(),
129        };
130        reg.send_control(b"READY")?;
131        reg.last_heartbeat = Instant::now();
132        Ok(reg)
133    }
134
135    pub fn kind(&self) -> RegistrationKind {
136        RegistrationKind::Service
137    }
138
139    pub fn send_control(&self, command: &[u8]) -> crate::errors::Result<()> {
140        self.socket
141            .send_multipart([command, self.service_name.as_bytes()], 0)?;
142        Ok(())
143    }
144
145    pub fn send_heartbeat(&self) -> crate::errors::Result<()> {
146        self.send_control(b"HEARTBEAT")
147    }
148
149    pub fn disconnect(&self) {
150        let _ = self.send_control(b"DISCONNECT");
151    }
152}
153
154pub struct ActionRegistration {
155    pub id: u64,
156    pub socket: Socket,
157    pub action_name: String,
158    pub handler: ActionGoalLiveHandler,
159    pub inflight: Arc<Mutex<HashMap<String, Arc<AtomicBool>>>>,
160    pub callback_group: CallbackGroup,
161    pub identity: Vec<u8>,
162    pub heartbeat_interval: Duration,
163    pub last_heartbeat: Instant,
164}
165
166impl ActionRegistration {
167    pub fn create(
168        id: u64,
169        context: &Context,
170        action_name: &str,
171        handler: ActionGoalHandler,
172        callback_group: CallbackGroup,
173        endpoint: &str,
174        identity: Option<&str>,
175        heartbeat_interval_ms: u64,
176        hwm: HighWaterMark,
177    ) -> crate::errors::Result<Self> {
178        Self::create_live(
179            id,
180            context,
181            action_name,
182            wrap_batch_action_handler(handler),
183            callback_group,
184            endpoint,
185            identity,
186            heartbeat_interval_ms,
187            hwm,
188        )
189    }
190
191    pub fn create_live(
192        id: u64,
193        context: &Context,
194        action_name: &str,
195        handler: ActionGoalLiveHandler,
196        callback_group: CallbackGroup,
197        endpoint: &str,
198        identity: Option<&str>,
199        heartbeat_interval_ms: u64,
200        hwm: HighWaterMark,
201    ) -> crate::errors::Result<Self> {
202        let socket = context.socket(SocketType::DEALER)?;
203        apply_action_options_with(&socket, hwm)?;
204        let worker_id = identity
205            .map(str::to_string)
206            .unwrap_or_else(|| format!("worker-{}", &Uuid::new_v4().simple().to_string()[..8]));
207        let identity = worker_id.into_bytes();
208        socket.set_identity(&identity)?;
209        socket.connect(endpoint)?;
210        let mut reg = Self {
211            id,
212            socket,
213            action_name: action_name.to_string(),
214            handler,
215            inflight: Arc::new(Mutex::new(HashMap::new())),
216            callback_group,
217            identity,
218            heartbeat_interval: Duration::from_millis(heartbeat_interval_ms),
219            last_heartbeat: Instant::now(),
220        };
221        reg.send_control(b"READY")?;
222        reg.last_heartbeat = Instant::now();
223        Ok(reg)
224    }
225
226    pub fn kind(&self) -> RegistrationKind {
227        RegistrationKind::Action
228    }
229
230    pub fn send_control(&self, command: &[u8]) -> crate::errors::Result<()> {
231        self.socket
232            .send_multipart([command, self.action_name.as_bytes()], 0)?;
233        Ok(())
234    }
235
236    pub fn send_heartbeat(&self) -> crate::errors::Result<()> {
237        self.send_control(b"HEARTBEAT")
238    }
239
240    pub fn disconnect(&self) {
241        let _ = self.send_control(b"DISCONNECT");
242    }
243}
244
245pub struct ActionClientRegistration {
246    pub socket: Socket,
247    pub endpoint: String,
248    pub goal_callbacks: std::collections::HashMap<String, ActionMessageCallback>,
249}
250
251impl ActionClientRegistration {
252    pub fn create(
253        context: &Context,
254        endpoint: &str,
255        hwm: HighWaterMark,
256    ) -> crate::errors::Result<Self> {
257        let socket = context.socket(SocketType::DEALER)?;
258        apply_action_options_with(&socket, hwm)?;
259        socket.connect(endpoint)?;
260        Ok(Self {
261            socket,
262            endpoint: endpoint.to_string(),
263            goal_callbacks: std::collections::HashMap::new(),
264        })
265    }
266
267    pub fn kind(&self) -> RegistrationKind {
268        RegistrationKind::ActionClient
269    }
270}
271
272pub enum Registration {
273    Sub(SubRegistration),
274    Service(ServiceRegistration),
275    Action(ActionRegistration),
276    ActionClient(ActionClientRegistration),
277}
278
279impl Registration {
280    pub fn socket(&self) -> &Socket {
281        match self {
282            Self::Sub(reg) => &reg.socket,
283            Self::Service(reg) => &reg.socket,
284            Self::Action(reg) => &reg.socket,
285            Self::ActionClient(reg) => &reg.socket,
286        }
287    }
288
289    pub fn kind(&self) -> RegistrationKind {
290        match self {
291            Self::Sub(reg) => reg.kind(),
292            Self::Service(reg) => reg.kind(),
293            Self::Action(reg) => reg.kind(),
294            Self::ActionClient(reg) => reg.kind(),
295        }
296    }
297}
robot_bus/runtime/registrations.rs

robot_bus/runtime/
registrations.rs
