moirai_transport/
remote_task.rs

1//! Remote task envelopes over transport-owned bytes.
2
3use crate::{
4    Address, NetworkTransport, RemoteAddress, Transport, TransportError, TransportResult,
5    payload::{ServerPayloadRegion, TransportPayload, archive_transport_payload},
6    safe_channel::{ArchiveSerialize, ArchiveView, ArchivedMessage},
7};
8
9mod capability;
10mod server;
11
12pub use capability::{
13    EchoBytesCapability, IntoRemoteOperation, RemoteCapability, RemoteCapabilityToken,
14    RemoteTaskOperationKind, SumU64Capability, build_remote_operation,
15};
16pub use server::{
17    BoundedRemoteTaskServer, RemoteTaskQueueCapacity, RemoteTaskRequestLimit, RemoteTaskServer,
18    RemoteTaskServerStats, RemoteTaskWorkerCount,
19};
20
21const OP_ECHO_BYTES: u8 = 1;
22const OP_SUM_U64: u8 = 2;
23const RESULT_BYTES: u8 = 1;
24const RESULT_U64: u8 = 2;
25
26/// Remote task identifier.
27#[repr(transparent)]
28#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
29pub struct RemoteTaskId(u64);
30
31impl RemoteTaskId {
32    /// Construct a remote task id.
33    pub const fn new(id: u64) -> Self {
34        Self(id)
35    }
36
37    /// Return the raw id.
38    pub const fn get(self) -> u64 {
39        self.0
40    }
41}
42
43/// Built-in remote task operation.
44#[derive(Debug, Clone, PartialEq, Eq)]
45pub enum RemoteTaskOperation {
46    /// Return the same byte payload.
47    EchoBytes(Vec<u8>),
48    /// Sum `u64` values with wrapping arithmetic.
49    SumU64(Vec<u64>),
50}
51
52/// Owned remote task envelope.
53#[derive(Debug, Clone, PartialEq, Eq)]
54pub struct RemoteTaskEnvelope {
55    /// Task id copied into the result envelope.
56    pub task_id: RemoteTaskId,
57    /// Address where the server sends the result.
58    pub reply_to: RemoteAddress,
59    /// Operation to execute.
60    pub operation: RemoteTaskOperation,
61}
62
63/// Borrowed remote task operation view.
64#[derive(Debug, Clone, Copy, PartialEq, Eq)]
65pub enum RemoteTaskOperationView<'a> {
66    /// Borrowed byte payload.
67    EchoBytes(&'a [u8]),
68    /// Borrowed little-endian u64 list.
69    SumU64(RemoteU64List<'a>),
70}
71
72/// Borrowed `u64` list backed by archive bytes.
73#[derive(Debug, Clone, Copy, PartialEq, Eq)]
74pub struct RemoteU64List<'a> {
75    bytes: &'a [u8],
76    len: usize,
77}
78
79impl RemoteU64List<'_> {
80    /// Number of `u64` values in the list.
81    pub const fn len(self) -> usize {
82        self.len
83    }
84
85    /// Whether the list is empty.
86    pub const fn is_empty(self) -> bool {
87        self.len == 0
88    }
89
90    /// Compute the wrapping sum without materializing an owned vector.
91    pub fn wrapping_sum(self) -> u64 {
92        self.bytes
93            .chunks_exact(core::mem::size_of::<u64>())
94            .fold(0u64, |sum, chunk| {
95                let bytes: [u8; 8] = chunk.try_into().expect("chunk size is fixed");
96                sum.wrapping_add(u64::from_le_bytes(bytes))
97            })
98    }
99}
100
101/// Borrowed remote task envelope view.
102#[derive(Debug, Clone, PartialEq, Eq)]
103pub struct RemoteTaskEnvelopeView<'a> {
104    /// Task id copied into the result envelope.
105    pub task_id: RemoteTaskId,
106    /// Address where the server sends the result.
107    pub reply_to: RemoteAddress,
108    /// Borrowed operation view.
109    pub operation: RemoteTaskOperationView<'a>,
110}
111
112/// Built-in remote task result.
113#[derive(Debug, Clone, PartialEq, Eq)]
114pub enum RemoteTaskOutput {
115    /// Returned byte payload.
116    Bytes(Vec<u8>),
117    /// Returned `u64` value.
118    U64(u64),
119}
120
121/// Owned remote task result envelope.
122#[derive(Debug, Clone, PartialEq, Eq)]
123pub struct RemoteTaskResult {
124    /// Completed task id.
125    pub task_id: RemoteTaskId,
126    /// Operation output.
127    pub output: RemoteTaskOutput,
128}
129
130/// Borrowed remote task result output.
131#[derive(Debug, Clone, Copy, PartialEq, Eq)]
132pub enum RemoteTaskOutputView<'a> {
133    /// Borrowed result bytes.
134    Bytes(&'a [u8]),
135    /// Result integer.
136    U64(u64),
137}
138
139/// Borrowed remote task result envelope view.
140#[derive(Debug, Clone, Copy, PartialEq, Eq)]
141pub struct RemoteTaskResultView<'a> {
142    /// Completed task id.
143    pub task_id: RemoteTaskId,
144    /// Borrowed operation output.
145    pub output: RemoteTaskOutputView<'a>,
146}
147
148/// Remote task client for a fixed server and reply endpoint.
149#[derive(Debug, Clone)]
150pub struct RemoteTaskClient {
151    server: RemoteAddress,
152    reply_to: RemoteAddress,
153}
154
155impl RemoteTaskClient {
156    /// Construct a remote task client.
157    pub fn new(server: RemoteAddress, reply_to: RemoteAddress) -> Self {
158        Self { server, reply_to }
159    }
160
161    /// Send a task, wait for its result, and validate the returned id.
162    pub fn execute(
163        &self,
164        task_id: RemoteTaskId,
165        operation: RemoteTaskOperation,
166    ) -> TransportResult<RemoteTaskResult> {
167        // The receiver must exist before publishing its address. Binding after
168        // send races a fast server reply and can exhaust connection attempts.
169        let reply = NetworkTransport {}.listen(&self.reply_to)?;
170        let envelope = RemoteTaskEnvelope {
171            task_id,
172            reply_to: self.reply_to.clone(),
173            operation,
174        };
175
176        let payload = archive_transport_payload::<ServerPayloadRegion, _>(&envelope)?;
177        NetworkTransport {}.send(&Address::Remote(self.server.clone()), payload.into_bytes())?;
178        let bytes = reply.recv()?;
179        let payload = TransportPayload::<ServerPayloadRegion>::from_bytes(bytes);
180        let message = ArchivedMessage::<RemoteTaskResult>::from_bytes(payload.into_bytes());
181        let view = message.get()?;
182        if view.task_id != task_id {
183            return Err(TransportError::Closed);
184        }
185
186        Ok(view.into_owned())
187    }
188}
189
190impl ArchiveSerialize for RemoteTaskEnvelope {
191    fn archive_size_hint(&self) -> usize {
192        core::mem::size_of::<u64>()
193            + remote_address_size(&self.reply_to)
194            + operation_size(&self.operation)
195    }
196
197    fn encode_archive(&self, output: &mut Vec<u8>) -> TransportResult<()> {
198        output.extend_from_slice(&self.task_id.get().to_le_bytes());
199        encode_remote_address(&self.reply_to, output)?;
200        match &self.operation {
201            RemoteTaskOperation::EchoBytes(bytes) => {
202                output.push(OP_ECHO_BYTES);
203                encode_len_prefixed_bytes(bytes, output)?;
204            }
205            RemoteTaskOperation::SumU64(values) => {
206                output.push(OP_SUM_U64);
207                let len = u32::try_from(values.len()).map_err(|_| TransportError::Closed)?;
208                output.extend_from_slice(&len.to_le_bytes());
209                for value in values {
210                    output.extend_from_slice(&value.to_le_bytes());
211                }
212            }
213        }
214        Ok(())
215    }
216}
217
218impl ArchiveView for RemoteTaskEnvelope {
219    type Archived<'a> = RemoteTaskEnvelopeView<'a>;
220
221    fn view_archive(bytes: &[u8]) -> TransportResult<Self::Archived<'_>> {
222        let mut cursor = ByteCursor::new(bytes);
223        let task_id = RemoteTaskId::new(u64::from_le_bytes(cursor.read_array()?));
224        let reply_to = cursor.read_remote_address()?;
225        let op = cursor.read_array::<1>()?[0];
226        let operation = match op {
227            OP_ECHO_BYTES => RemoteTaskOperationView::EchoBytes(cursor.read_len_prefixed_bytes()?),
228            OP_SUM_U64 => {
229                let len = usize::try_from(u32::from_le_bytes(cursor.read_array()?))
230                    .map_err(|_| TransportError::Closed)?;
231                let byte_len = len
232                    .checked_mul(core::mem::size_of::<u64>())
233                    .ok_or(TransportError::Closed)?;
234                let payload = cursor.read_exact(byte_len)?;
235                RemoteTaskOperationView::SumU64(RemoteU64List {
236                    bytes: payload,
237                    len,
238                })
239            }
240            _ => return Err(TransportError::Closed),
241        };
242        cursor.finish()?;
243
244        Ok(RemoteTaskEnvelopeView {
245            task_id,
246            reply_to,
247            operation,
248        })
249    }
250}
251
252impl ArchiveSerialize for RemoteTaskResult {
253    fn archive_size_hint(&self) -> usize {
254        core::mem::size_of::<u64>()
255            + 1
256            + match &self.output {
257                RemoteTaskOutput::Bytes(bytes) => core::mem::size_of::<u32>() + bytes.len(),
258                RemoteTaskOutput::U64(_) => core::mem::size_of::<u64>(),
259            }
260    }
261
262    fn encode_archive(&self, output: &mut Vec<u8>) -> TransportResult<()> {
263        output.extend_from_slice(&self.task_id.get().to_le_bytes());
264        match &self.output {
265            RemoteTaskOutput::Bytes(bytes) => {
266                output.push(RESULT_BYTES);
267                encode_len_prefixed_bytes(bytes, output)?;
268            }
269            RemoteTaskOutput::U64(value) => {
270                output.push(RESULT_U64);
271                output.extend_from_slice(&value.to_le_bytes());
272            }
273        }
274        Ok(())
275    }
276}
277
278impl ArchiveView for RemoteTaskResult {
279    type Archived<'a> = RemoteTaskResultView<'a>;
280
281    fn view_archive(bytes: &[u8]) -> TransportResult<Self::Archived<'_>> {
282        let mut cursor = ByteCursor::new(bytes);
283        let task_id = RemoteTaskId::new(u64::from_le_bytes(cursor.read_array()?));
284        let tag = cursor.read_array::<1>()?[0];
285        let output = match tag {
286            RESULT_BYTES => RemoteTaskOutputView::Bytes(cursor.read_len_prefixed_bytes()?),
287            RESULT_U64 => RemoteTaskOutputView::U64(u64::from_le_bytes(cursor.read_array()?)),
288            _ => return Err(TransportError::Closed),
289        };
290        cursor.finish()?;
291
292        Ok(RemoteTaskResultView { task_id, output })
293    }
294}
295
296impl RemoteTaskResultView<'_> {
297    /// Materialize the borrowed result view into an owned result.
298    pub fn into_owned(self) -> RemoteTaskResult {
299        RemoteTaskResult {
300            task_id: self.task_id,
301            output: match self.output {
302                RemoteTaskOutputView::Bytes(bytes) => RemoteTaskOutput::Bytes(bytes.to_vec()),
303                RemoteTaskOutputView::U64(value) => RemoteTaskOutput::U64(value),
304            },
305        }
306    }
307}
308
309pub(super) fn execute_remote_task(request: &RemoteTaskEnvelopeView<'_>) -> RemoteTaskResult {
310    let output = match request.operation {
311        RemoteTaskOperationView::EchoBytes(bytes) => RemoteTaskOutput::Bytes(bytes.to_vec()),
312        RemoteTaskOperationView::SumU64(values) => RemoteTaskOutput::U64(values.wrapping_sum()),
313    };
314
315    RemoteTaskResult {
316        task_id: request.task_id,
317        output,
318    }
319}
320
321fn encode_remote_address(address: &RemoteAddress, output: &mut Vec<u8>) -> TransportResult<()> {
322    encode_len_prefixed_str(&address.host, output)?;
323    output.extend_from_slice(&address.port.to_le_bytes());
324    encode_len_prefixed_str(&address.service, output)
325}
326
327fn remote_address_size(address: &RemoteAddress) -> usize {
328    core::mem::size_of::<u32>()
329        + address.host.len()
330        + core::mem::size_of::<u16>()
331        + core::mem::size_of::<u32>()
332        + address.service.len()
333}
334
335fn operation_size(operation: &RemoteTaskOperation) -> usize {
336    1 + match operation {
337        RemoteTaskOperation::EchoBytes(bytes) => core::mem::size_of::<u32>() + bytes.len(),
338        RemoteTaskOperation::SumU64(values) => {
339            core::mem::size_of::<u32>() + values.len() * core::mem::size_of::<u64>()
340        }
341    }
342}
343
344fn encode_len_prefixed_str(value: &str, output: &mut Vec<u8>) -> TransportResult<()> {
345    encode_len_prefixed_bytes(value.as_bytes(), output)
346}
347
348fn encode_len_prefixed_bytes(value: &[u8], output: &mut Vec<u8>) -> TransportResult<()> {
349    let len = u32::try_from(value.len()).map_err(|_| TransportError::Closed)?;
350    output.extend_from_slice(&len.to_le_bytes());
351    output.extend_from_slice(value);
352    Ok(())
353}
354
355struct ByteCursor<'a> {
356    bytes: &'a [u8],
357    offset: usize,
358}
359
360impl<'a> ByteCursor<'a> {
361    fn new(bytes: &'a [u8]) -> Self {
362        Self { bytes, offset: 0 }
363    }
364
365    /// Read exactly `N` bytes and return them as a fixed-size array.
366    ///
367    /// One generic reader replaces the per-width `read_u16`/`read_u32`/
368    /// `read_u64` clones that previously differed only in the array length and
369    /// the `from_le_bytes` call. `N` is inferred at the call site from the
370    /// integer type it is converted into, so callers stay width-explicit
371    /// (`u32::from_le_bytes(cursor.read_array()?)`) without a per-width
372    /// method.
373    fn read_array<const N: usize>(&mut self) -> TransportResult<[u8; N]> {
374        self.read_exact(N)?
375            .try_into()
376            .map_err(|_| TransportError::Closed)
377    }
378
379    fn read_len_prefixed_bytes(&mut self) -> TransportResult<&'a [u8]> {
380        let len = usize::try_from(u32::from_le_bytes(self.read_array()?))
381            .map_err(|_| TransportError::Closed)?;
382        self.read_exact(len)
383    }
384
385    fn read_len_prefixed_string(&mut self) -> TransportResult<String> {
386        let bytes = self.read_len_prefixed_bytes()?;
387        let value = core::str::from_utf8(bytes).map_err(|_| TransportError::Closed)?;
388        Ok(value.to_string())
389    }
390
391    fn read_remote_address(&mut self) -> TransportResult<RemoteAddress> {
392        let host = self.read_len_prefixed_string()?;
393        let port = u16::from_le_bytes(self.read_array()?);
394        let service = self.read_len_prefixed_string()?;
395        Ok(RemoteAddress {
396            host,
397            port,
398            service,
399        })
400    }
401
402    fn read_exact(&mut self, len: usize) -> TransportResult<&'a [u8]> {
403        let end = self.offset.checked_add(len).ok_or(TransportError::Closed)?;
404        let bytes = self
405            .bytes
406            .get(self.offset..end)
407            .ok_or(TransportError::Closed)?;
408        self.offset = end;
409        Ok(bytes)
410    }
411
412    fn finish(self) -> TransportResult<()> {
413        if self.offset == self.bytes.len() {
414            Ok(())
415        } else {
416            Err(TransportError::Closed)
417        }
418    }
419}
420
421#[cfg(test)]
422mod tests;
moirai_transport/remote_task.rs

moirai_transport/
remote_task.rs
