spatialrust_distribute/
transfer.rs1use crate::{DistributeError, DistributeResult, PartitionGraph};
4
5#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
7pub enum TransferDirection {
8 HostToDevice,
10 DeviceToHost,
12 HostToNetwork,
14 NetworkToHost,
16}
17
18#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
20pub enum TransferKind {
21 ExplicitCopy,
23 ZeroCopyHandoff,
25}
26
27#[derive(Clone, Debug, PartialEq, Eq)]
29pub struct NamedTransfer {
30 pub name: String,
32 pub direction: TransferDirection,
34 pub kind: TransferKind,
36 pub from: String,
38 pub to: String,
40 pub bytes: u64,
42}
43
44impl NamedTransfer {
45 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 #[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#[derive(Clone, Debug, Default, PartialEq, Eq)]
82pub struct TransferPlan {
83 transfers: Vec<NamedTransfer>,
84}
85
86impl TransferPlan {
87 #[must_use]
89 pub fn new() -> Self {
90 Self::default()
91 }
92
93 pub fn push(&mut self, transfer: NamedTransfer) {
95 self.transfers.push(transfer);
96 }
97
98 #[must_use]
100 pub fn transfers(&self) -> &[NamedTransfer] {
101 &self.transfers
102 }
103
104 #[must_use]
106 pub fn total_bytes(&self) -> u64 {
107 self.transfers.iter().map(|t| t.bytes).sum()
108 }
109
110 #[must_use]
112 pub fn counted_copy_bytes(&self) -> u64 {
113 self.transfers.iter().map(NamedTransfer::counted_copy_bytes).sum()
114 }
115
116 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#[derive(Clone, Debug, Default, PartialEq, Eq)]
138pub struct TransferLedger {
139 completed: Vec<NamedTransfer>,
140}
141
142impl TransferLedger {
143 #[must_use]
145 pub fn new() -> Self {
146 Self::default()
147 }
148
149 pub fn record(&mut self, transfer: NamedTransfer) {
151 self.completed.push(transfer);
152 }
153
154 #[must_use]
156 pub fn completed(&self) -> &[NamedTransfer] {
157 &self.completed
158 }
159
160 #[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}