soma_fleet/
transfer_guard.rs1use std::sync::{Arc, Mutex, MutexGuard};
2
3use crate::{FleetError, FleetResult, HostId, TransferReceipt, TransferRequest};
4
5#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum TransferGuardState {
8 InProgress {
10 bytes: u64,
12 max_bytes: u64,
14 },
15 Completed {
17 bytes: u64,
19 verified: bool,
21 },
22 Cancelled {
24 bytes: u64,
26 },
27 Failed {
29 bytes: u64,
31 message: String,
33 },
34 Abandoned {
36 bytes: u64,
38 },
39}
40
41#[derive(Debug, Clone)]
43pub struct TransferLifecycle {
44 source: HostId,
45 destination: HostId,
46 state: Arc<Mutex<TransferGuardState>>,
47}
48
49impl TransferLifecycle {
50 #[must_use]
52 pub fn start(request: &TransferRequest) -> (Self, TransferGuard) {
53 let state = Arc::new(Mutex::new(TransferGuardState::InProgress {
54 bytes: 0,
55 max_bytes: request.max_bytes(),
56 }));
57 let lifecycle = Self {
58 source: request.source_host().clone(),
59 destination: request.destination_host().clone(),
60 state: Arc::clone(&state),
61 };
62 let guard = TransferGuard {
63 state,
64 terminal: false,
65 };
66 (lifecycle, guard)
67 }
68
69 #[must_use]
71 pub fn source(&self) -> &HostId {
72 &self.source
73 }
74 #[must_use]
76 pub fn destination(&self) -> &HostId {
77 &self.destination
78 }
79 #[must_use]
81 pub fn snapshot(&self) -> TransferGuardState {
82 lock(&self.state).clone()
83 }
84}
85
86pub struct TransferGuard {
88 state: Arc<Mutex<TransferGuardState>>,
89 terminal: bool,
90}
91
92impl TransferGuard {
93 pub fn record_chunk(&mut self, bytes: u64) -> FleetResult<u64> {
95 let mut state = lock(&self.state);
96 let TransferGuardState::InProgress {
97 bytes: observed,
98 max_bytes,
99 } = &mut *state
100 else {
101 return Err(FleetError::TransferLifecycle(
102 "cannot record bytes after a terminal state".into(),
103 ));
104 };
105 let next = observed.checked_add(bytes).ok_or_else(|| {
106 FleetError::TransferLifecycle("transfer byte counter overflow".into())
107 })?;
108 if next > *max_bytes {
109 return Err(FleetError::TransferLifecycle(format!(
110 "transfer exceeded maximum of {max_bytes} bytes"
111 )));
112 }
113 *observed = next;
114 Ok(next)
115 }
116
117 pub fn complete(mut self, receipt: TransferReceipt) -> FleetResult<TransferReceipt> {
119 let observed = match &*lock(&self.state) {
120 TransferGuardState::InProgress { bytes, .. } => *bytes,
121 _ => {
122 return Err(FleetError::TransferLifecycle(
123 "cannot complete a terminal transfer".into(),
124 ));
125 }
126 };
127 if observed != receipt.bytes() {
128 return Err(FleetError::TransferLifecycle(format!(
129 "receipt reports {} bytes but lifecycle observed {observed}",
130 receipt.bytes()
131 )));
132 }
133 *lock(&self.state) = TransferGuardState::Completed {
134 bytes: observed,
135 verified: receipt.verified(),
136 };
137 self.terminal = true;
138 Ok(receipt)
139 }
140
141 pub fn cancel(mut self) -> FleetResult<()> {
143 let bytes = in_progress_bytes(&self.state)?;
144 *lock(&self.state) = TransferGuardState::Cancelled { bytes };
145 self.terminal = true;
146 Ok(())
147 }
148
149 pub fn fail(mut self, message: impl Into<String>) -> FleetResult<()> {
151 let message = message.into();
152 if message.is_empty()
153 || message.chars().count() > 1024
154 || message.chars().any(char::is_control)
155 {
156 return Err(FleetError::TransferLifecycle(
157 "transfer failure detail is invalid".into(),
158 ));
159 }
160 let bytes = in_progress_bytes(&self.state)?;
161 *lock(&self.state) = TransferGuardState::Failed { bytes, message };
162 self.terminal = true;
163 Ok(())
164 }
165}
166
167impl Drop for TransferGuard {
168 fn drop(&mut self) {
169 if self.terminal {
170 return;
171 }
172 let mut state = lock(&self.state);
173 if let TransferGuardState::InProgress { bytes, .. } = *state {
174 *state = TransferGuardState::Abandoned { bytes };
175 }
176 }
177}
178
179fn in_progress_bytes(state: &Mutex<TransferGuardState>) -> FleetResult<u64> {
180 match &*lock(state) {
181 TransferGuardState::InProgress { bytes, .. } => Ok(*bytes),
182 _ => Err(FleetError::TransferLifecycle(
183 "transfer is already terminal".into(),
184 )),
185 }
186}
187
188fn lock(state: &Mutex<TransferGuardState>) -> MutexGuard<'_, TransferGuardState> {
189 state
190 .lock()
191 .unwrap_or_else(std::sync::PoisonError::into_inner)
192}
193
194#[cfg(test)]
195#[path = "transfer_guard_tests.rs"]
196mod tests;