Skip to main content

spatialrust_distribute/
transfer.rs

1//! Named host/device/network transfers and measurable plans.
2
3use crate::{DistributeError, DistributeResult, PartitionGraph};
4
5/// Transfer direction.
6#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
7pub enum TransferDirection {
8    /// Host to device.
9    HostToDevice,
10    /// Device to host.
11    DeviceToHost,
12    /// Host to remote host.
13    HostToNetwork,
14    /// Remote host to host.
15    NetworkToHost,
16}
17
18/// Kind of payload transfer.
19#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
20pub enum TransferKind {
21    /// Explicit named copy (never implied).
22    ExplicitCopy,
23    /// Zero-copy handoff when both ends agree.
24    ZeroCopyHandoff,
25}
26
27/// One named transfer in a distributed plan.
28#[derive(Clone, Debug, PartialEq, Eq)]
29pub struct NamedTransfer {
30    /// Transfer name used in telemetry.
31    pub name: String,
32    /// Direction.
33    pub direction: TransferDirection,
34    /// Kind.
35    pub kind: TransferKind,
36    /// Source node id.
37    pub from: String,
38    /// Destination node id.
39    pub to: String,
40    /// Approximate payload bytes.
41    pub bytes: u64,
42}
43
44impl NamedTransfer {
45    /// Creates a validated named transfer.
46    pub fn try_new(
47        name: impl Into<String>,
48        direction: TransferDirection,
49        kind: TransferKind,
50        from: impl Into<String>,
51        to: impl Into<String>,
52        bytes: u64,
53    ) -> DistributeResult<Self> {
54        let name = name.into();
55        let from = from.into();
56        let to = to.into();
57        if name.is_empty() || from.is_empty() || to.is_empty() {
58            return Err(DistributeError::InvalidConfiguration(
59                "transfer name/from/to must be non-empty".into(),
60            ));
61        }
62        if from == to {
63            return Err(DistributeError::InvalidConfiguration(
64                "transfer endpoints must differ".into(),
65            ));
66        }
67        Ok(Self { name, direction, kind, from, to, bytes })
68    }
69
70    /// Bytes counted as measurable copies (zero-copy handoffs are 0).
71    #[must_use]
72    pub fn counted_copy_bytes(&self) -> u64 {
73        match self.kind {
74            TransferKind::ExplicitCopy => self.bytes,
75            TransferKind::ZeroCopyHandoff => 0,
76        }
77    }
78}
79
80/// Ordered plan of named transfers.
81#[derive(Clone, Debug, Default, PartialEq, Eq)]
82pub struct TransferPlan {
83    transfers: Vec<NamedTransfer>,
84}
85
86impl TransferPlan {
87    /// Creates an empty plan.
88    #[must_use]
89    pub fn new() -> Self {
90        Self::default()
91    }
92
93    /// Appends a transfer.
94    pub fn push(&mut self, transfer: NamedTransfer) {
95        self.transfers.push(transfer);
96    }
97
98    /// Returns transfers.
99    #[must_use]
100    pub fn transfers(&self) -> &[NamedTransfer] {
101        &self.transfers
102    }
103
104    /// Sum of payload bytes (including zero-copy declared sizes).
105    #[must_use]
106    pub fn total_bytes(&self) -> u64 {
107        self.transfers.iter().map(|t| t.bytes).sum()
108    }
109
110    /// Sum of measurable explicit-copy bytes.
111    #[must_use]
112    pub fn counted_copy_bytes(&self) -> u64 {
113        self.transfers.iter().map(NamedTransfer::counted_copy_bytes).sum()
114    }
115
116    /// Ensures every endpoint exists in the partition graph.
117    pub fn validate_against(&self, graph: &PartitionGraph) -> DistributeResult<()> {
118        for transfer in &self.transfers {
119            if graph.partition_of_node(&transfer.from).is_none() {
120                return Err(DistributeError::Missing(format!(
121                    "transfer source node `{}`",
122                    transfer.from
123                )));
124            }
125            if graph.partition_of_node(&transfer.to).is_none() {
126                return Err(DistributeError::Missing(format!(
127                    "transfer destination node `{}`",
128                    transfer.to
129                )));
130            }
131        }
132        Ok(())
133    }
134}
135
136/// Append-only ledger of completed named transfers for telemetry.
137#[derive(Clone, Debug, Default, PartialEq, Eq)]
138pub struct TransferLedger {
139    completed: Vec<NamedTransfer>,
140}
141
142impl TransferLedger {
143    /// Creates an empty ledger.
144    #[must_use]
145    pub fn new() -> Self {
146        Self::default()
147    }
148
149    /// Records a completed transfer.
150    pub fn record(&mut self, transfer: NamedTransfer) {
151        self.completed.push(transfer);
152    }
153
154    /// Returns completed transfers.
155    #[must_use]
156    pub fn completed(&self) -> &[NamedTransfer] {
157        &self.completed
158    }
159
160    /// Total measurable copy bytes recorded.
161    #[must_use]
162    pub fn counted_copy_bytes(&self) -> u64 {
163        self.completed.iter().map(NamedTransfer::counted_copy_bytes).sum()
164    }
165}
166
167#[cfg(test)]
168mod tests {
169    use super::{NamedTransfer, TransferDirection, TransferKind, TransferLedger, TransferPlan};
170    use crate::{ExecutionPartition, PartitionGraph};
171
172    #[test]
173    fn plan_validates_and_counts_copies() {
174        let mut graph = PartitionGraph::new();
175        graph
176            .insert_partition(ExecutionPartition::try_new("edge", vec!["cam".into()]).unwrap())
177            .unwrap();
178        graph
179            .insert_partition(ExecutionPartition::try_new("host", vec!["scene".into()]).unwrap())
180            .unwrap();
181        let mut plan = TransferPlan::new();
182        plan.push(
183            NamedTransfer::try_new(
184                "cam-to-scene",
185                TransferDirection::HostToNetwork,
186                TransferKind::ExplicitCopy,
187                "cam",
188                "scene",
189                1024,
190            )
191            .unwrap(),
192        );
193        plan.push(
194            NamedTransfer::try_new(
195                "handoff",
196                TransferDirection::HostToDevice,
197                TransferKind::ZeroCopyHandoff,
198                "scene",
199                "cam",
200                4096,
201            )
202            .unwrap(),
203        );
204        plan.validate_against(&graph).unwrap();
205        assert_eq!(plan.total_bytes(), 5120);
206        assert_eq!(plan.counted_copy_bytes(), 1024);
207
208        let mut ledger = TransferLedger::new();
209        ledger.record(plan.transfers()[0].clone());
210        assert_eq!(ledger.counted_copy_bytes(), 1024);
211    }
212}