1use crate::observation::{CommitObserver, RowChange};
12use crate::transfer_watch::{changes_since, Heard, TransferChange, TransferEnd};
13use coven_foundation::id_source::{DeviceId, FileId, ObjectId};
14use coven_foundation::locks::lock_unpoisoned;
15use rusqlite::functions::FunctionFlags;
16use std::collections::BTreeMap;
17use std::sync::{Arc, Mutex};
18use tokio::sync::watch;
19
20#[derive(Clone, Copy, Debug, PartialEq, Eq)]
24pub struct TransferProgress {
25 pub bytes_transferred: u64,
27 pub bytes_total: u64,
29}
30
31struct Direction {
33 read_as: &'static [(&'static str, &'static str)],
36 contains: Option<&'static str>,
39 objects: &'static str,
43}
44
45const SENDING: Direction = Direction {
46 read_as: &[("coven_uploads", "sending"), ("coven_sending", "")],
47 contains: Some("_coven_sending"),
48 objects: "_coven_sending_objects",
49};
50
51const DOWNLOADING: Direction = Direction {
52 read_as: &[("coven_downloading", "")],
53 contains: None,
54 objects: "_coven_downloading_objects",
55};
56
57pub(crate) const FUNCTIONS: [&str; 3] =
60 ["_coven_sending", "_coven_sending_objects", "_coven_downloading_objects"];
61
62#[derive(Clone)]
64pub(crate) struct TransfersInFlight {
65 uploads: InFlight,
66 downloads: InFlight,
67}
68
69impl TransfersInFlight {
70 pub(crate) fn new(observer: CommitObserver) -> Self {
71 Self {
72 uploads: InFlight::new(&SENDING, observer.clone()),
73 downloads: InFlight::new(&DOWNLOADING, observer),
74 }
75 }
76
77 pub(crate) fn uploads(&self) -> InFlight {
79 self.uploads.clone()
80 }
81
82 pub(crate) fn downloads(&self) -> InFlight {
85 self.downloads.clone()
86 }
87
88 pub(crate) fn register(&self, connection: &rusqlite::Connection) -> rusqlite::Result<()> {
91 self.uploads.register(connection)?;
92 self.downloads.register(connection)
93 }
94}
95
96#[derive(Clone)]
100pub struct InFlight {
101 shared: Arc<InFlightShared>,
102}
103
104struct InFlightShared {
105 direction: &'static Direction,
106 state: Mutex<InFlightState>,
107 observer: CommitObserver,
110 changed: watch::Sender<()>,
113}
114
115#[derive(Default)]
116struct InFlightState {
117 objects: BTreeMap<ObjectId, Entry>,
119 next: u64,
122 watches: BTreeMap<u64, Option<BTreeMap<FileId, Heard>>>,
128 next_watch: u64,
130}
131
132struct Entry {
133 device: DeviceId,
134 attempts: Vec<(u64, Vec<FileId>, TransferProgress)>,
139}
140
141impl InFlightState {
142 fn ended(&mut self, files: &[FileId], end: TransferEnd) {
145 let now = self.progress();
146 for file in files.iter().filter(|file| !now.contains_key(*file)) {
147 for heard in self.watches.values_mut().flatten() {
148 if let Some(heard) = heard.get_mut(file) {
149 *heard = Heard::Ended(end);
150 }
151 }
152 }
153 }
154
155 fn progress(&self) -> BTreeMap<FileId, TransferProgress> {
158 let mut progress = BTreeMap::new();
159 for entry in self.objects.values() {
160 for (_, files, moved) in &entry.attempts {
161 for file in files {
162 progress.insert(*file, *moved);
163 }
164 }
165 }
166 progress
167 }
168}
169
170impl InFlight {
171 fn new(direction: &'static Direction, observer: CommitObserver) -> Self {
172 Self {
173 shared: Arc::new(InFlightShared {
174 direction,
175 state: Mutex::new(InFlightState::default()),
176 observer,
177 changed: watch::channel(()).0,
178 }),
179 }
180 }
181
182 fn state(&self) -> std::sync::MutexGuard<'_, InFlightState> {
183 lock_unpoisoned(&self.shared.state)
184 }
185
186 pub fn current(&self) -> std::collections::BTreeSet<ObjectId> {
188 self.state().objects.keys().copied().collect()
189 }
190
191 pub fn start(
194 &self,
195 device: DeviceId,
196 object: ObjectId,
197 files: Vec<FileId>,
198 progress: TransferProgress,
199 ) -> Transferring {
200 let (attempt, started) = {
201 let mut state = self.state();
202 let attempt = state.next;
203 state.next += 1;
204 let entry = state
205 .objects
206 .entry(object)
207 .or_insert_with(|| Entry { device, attempts: Vec::new() });
208 entry.attempts.push((attempt, files, progress));
209 (attempt, entry.attempts.len() == 1)
210 };
211 self.shared.changed.send_replace(());
212 if started {
213 self.changed();
214 }
215 Transferring { owner: self.clone(), object, attempt, end: TransferEnd::Dropped }
216 }
217
218 pub fn subscribe(&self) -> crate::TransferWatch {
221 let number = {
222 let mut state = self.state();
223 let number = state.next_watch;
224 state.next_watch += 1;
225 state.watches.insert(number, None);
226 number
227 };
228 crate::TransferWatch::new(
229 self.clone(),
230 number,
231 self.shared.changed.subscribe(),
232 self.shared.observer.subscribe(),
233 )
234 }
235
236 pub(crate) fn read(&self, number: u64) -> Vec<TransferChange> {
239 let mut state = self.state();
240 let now = state.progress();
241 let heard = state.watches.get_mut(&number).expect("a watch is registered until it drops");
242 match heard {
243 Some(heard) => changes_since(heard, now),
244 None => {
245 *heard = Some(
246 now.iter().map(|(file, progress)| (*file, Heard::Moved(*progress))).collect(),
247 );
248 vec![TransferChange::Reset(now)]
249 }
250 }
251 }
252
253 pub(crate) fn forget(&self, number: u64) {
255 self.state().watches.remove(&number);
256 }
257
258 fn reached(&self, object: ObjectId, attempt: u64, bytes: u64) {
261 let mut state = self.state();
262 let Some(progress) = state.objects.get_mut(&object).and_then(|entry| {
263 entry
264 .attempts
265 .iter_mut()
266 .find(|(number, _, _)| *number == attempt)
267 .map(|(_, _, progress)| progress)
268 }) else {
269 return;
270 };
271 progress.bytes_transferred = bytes;
272 drop(state);
273 self.shared.changed.send_replace(());
274 }
275
276 fn changed(&self) {
279 self.shared.observer.commit(
280 crate::observation::Judged::unjudged(
281 self.shared
282 .direction
283 .read_as
284 .iter()
285 .map(|(table, column)| RowChange {
286 table: (*table).into(),
287 column: (*column).into(),
288 rows: None,
289 })
290 .collect(),
291 ),
292 None,
293 );
294 }
295
296 fn register(&self, connection: &rusqlite::Connection) -> rusqlite::Result<()> {
297 let flags = FunctionFlags::SQLITE_UTF8 | FunctionFlags::SQLITE_INNOCUOUS;
298 if let Some(name) = self.shared.direction.contains {
299 let owner = self.clone();
300 connection.create_scalar_function(name, 1, flags, move |ctx| {
301 Ok(owner.state().objects.contains_key(&ctx.get::<ObjectId>(0)?))
303 })?;
304 }
305 let owner = self.clone();
306 connection.create_scalar_function(self.shared.direction.objects, 0, flags, move |_| {
307 let objects: Vec<String> = owner
308 .state()
309 .objects
310 .iter()
311 .map(|(object, entry)| format!("{:016x}{}", entry.device.0, object.0.simple()))
312 .collect();
313 serde_json::to_string(&objects)
314 .map_err(|error| rusqlite::Error::UserFunctionError(Box::new(error)))
315 })
316 }
317}
318
319pub struct Transferring {
323 owner: InFlight,
324 object: ObjectId,
325 attempt: u64,
326 end: TransferEnd,
328}
329
330impl Transferring {
331 pub fn meter(&self) -> TransferMeter {
333 TransferMeter { owner: self.owner.clone(), object: self.object, attempt: self.attempt }
334 }
335
336 pub fn end(mut self, end: TransferEnd) {
338 self.end = end;
339 }
340}
341
342#[derive(Clone)]
345pub struct TransferMeter {
346 owner: InFlight,
347 object: ObjectId,
348 attempt: u64,
349}
350
351impl TransferMeter {
352 pub fn reached(&self, bytes: u64) {
354 self.owner.reached(self.object, self.attempt, bytes);
355 }
356}
357
358impl Drop for Transferring {
359 fn drop(&mut self) {
360 let ended = {
361 let mut state = self.owner.state();
362 let attempt = state.objects.get_mut(&self.object).and_then(|entry| {
363 let at =
364 entry.attempts.iter().position(|(attempt, _, _)| *attempt == self.attempt)?;
365 let (_, files, _) = entry.attempts.remove(at);
366 Some((files, entry.attempts.is_empty()))
367 });
368 attempt.map(|(files, ended)| {
369 if ended {
370 state.objects.remove(&self.object);
371 }
372 state.ended(&files, self.end);
373 ended
374 })
375 };
376 match ended {
377 Some(ended) => {
378 self.owner.shared.changed.send_replace(());
379 if ended {
380 self.owner.changed();
381 }
382 }
383 None => self.owner.shared.observer.fail_with_bug(coven_foundation::task::Bug::new(
389 "end a transfer in flight",
390 "an attempt in flight lost its place in its object's entry",
391 )),
392 }
393 }
394}
395
396#[cfg(test)]
397#[path = "transfers_in_flight_tests.rs"]
398mod tests;