Skip to main content

soma_fleet/
transfer_guard.rs

1use std::sync::{Arc, Mutex, MutexGuard};
2
3use crate::{FleetError, FleetResult, HostId, TransferReceipt, TransferRequest};
4
5/// Observable terminal or in-flight transfer state.
6#[derive(Debug, Clone, PartialEq, Eq)]
7pub enum TransferGuardState {
8    /// Transfer is active and has observed bounded bytes.
9    InProgress {
10        /// Bytes observed so far.
11        bytes: u64,
12        /// Maximum bytes admitted by the request.
13        max_bytes: u64,
14    },
15    /// Transfer completed with optional digest verification.
16    Completed {
17        /// Completed byte count.
18        bytes: u64,
19        /// Whether source and destination digests matched.
20        verified: bool,
21    },
22    /// Transfer was explicitly cancelled.
23    Cancelled {
24        /// Bytes observed before cancellation.
25        bytes: u64,
26    },
27    /// Transfer failed with a bounded diagnostic.
28    Failed {
29        /// Bytes observed before failure.
30        bytes: u64,
31        /// Bounded failure detail.
32        message: String,
33    },
34    /// Guard was dropped before a terminal method was called.
35    Abandoned {
36        /// Bytes observed before abandonment.
37        bytes: u64,
38    },
39}
40
41/// Cloneable observer for one transfer guard.
42#[derive(Debug, Clone)]
43pub struct TransferLifecycle {
44    source: HostId,
45    destination: HostId,
46    state: Arc<Mutex<TransferGuardState>>,
47}
48
49impl TransferLifecycle {
50    /// Starts a lifecycle and returns its mutable RAII guard.
51    #[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    /// Returns the source host.
70    #[must_use]
71    pub fn source(&self) -> &HostId {
72        &self.source
73    }
74    /// Returns the destination host.
75    #[must_use]
76    pub fn destination(&self) -> &HostId {
77        &self.destination
78    }
79    /// Returns a consistent state snapshot.
80    #[must_use]
81    pub fn snapshot(&self) -> TransferGuardState {
82        lock(&self.state).clone()
83    }
84}
85
86/// RAII transfer accounting guard.
87pub struct TransferGuard {
88    state: Arc<Mutex<TransferGuardState>>,
89    terminal: bool,
90}
91
92impl TransferGuard {
93    /// Records a completed chunk and rejects overflow or bound violations.
94    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    /// Marks the transfer complete after checking receipt byte parity.
118    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    /// Marks the transfer cancelled.
142    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    /// Marks the transfer failed with bounded detail.
150    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;