Skip to main content

coven_database/
file_staging.rs

1//! Staging app bytes before the write that puts them in rows (§17.2): async
2//! sources are consumed without holding SQLite or a blocking worker.
3
4use super::DatabaseInner;
5use crate::file_hashes::FileHasher;
6use crate::file_removals::OwnedFile;
7use crate::file_write::{StagedFile, StagedSource};
8use crate::write_batch::BatchFile;
9use crate::{DbError, FileSource, Place, Provenance, WriteBatch};
10use coven_foundation::files::StoreDirAccess;
11use coven_foundation::files::{FileName, FlushBatch, Flushed, StagedName, WriterFolder};
12use coven_foundation::locks::read_unpoisoned;
13use std::collections::VecDeque;
14use std::sync::{Arc, RwLock};
15use tokio::sync::OwnedRwLockReadGuard;
16
17/// App bytes still to write to coven's copy of them, and where their file is
18/// to be.
19struct StagingSource {
20    copy: StagingCopy,
21    source: FileSource,
22}
23
24/// Where a staged copy is to be and what it is for, apart from its bytes.
25struct StagingCopy {
26    namespace: String,
27    id: coven_foundation::id_source::FileId,
28    name: StagedName,
29    place: Place,
30}
31
32type StagingReader = std::pin::Pin<Box<dyn tokio::io::AsyncRead + Send>>;
33
34/// Sources in the order the batch put them: the copies of bytes the app
35/// holds, then the stream that ends them, if the batch has one.
36struct Run {
37    bytes: Vec<(StagingCopy, Vec<u8>)>,
38    stream: Option<(StagingCopy, StagingReader)>,
39}
40
41impl StagingCopy {
42    /// The copy as staged, once its `size` bytes were written and hashed.
43    fn staged(self, size: u64, hash: FileHasher) -> StagedFile {
44        StagedFile {
45            id: self.id,
46            namespace: self.namespace,
47            size,
48            source: StagedSource::Copy(self.name),
49            hashes: hash.finish(),
50            place: self.place,
51        }
52    }
53}
54
55pub(super) struct FileStaging {
56    database: Arc<RwLock<Option<DatabaseInner>>>,
57    /// The names the store reserves, these among them until they go.
58    reserved: Arc<crate::reserved_names::ReservedNames>,
59    // Close waits until publication or cancellation releases the reservation.
60    lease: Option<OwnedRwLockReadGuard<()>>,
61    names: Vec<OwnedFile>,
62    sources: VecDeque<StagingSource>,
63    staged: Vec<StagedFile>,
64}
65
66impl FileStaging {
67    pub(super) fn new<T, E: crate::WriteFailure>(
68        database: Arc<RwLock<Option<DatabaseInner>>>,
69        lease: OwnedRwLockReadGuard<()>,
70        build: impl FnOnce(&mut WriteBatch<'_>) -> Result<T, E>,
71    ) -> Result<(Self, T), E> {
72        let (sources, value, reserved) = {
73            let slot = read_unpoisoned(&database);
74            let inner = slot.as_ref().ok_or(DbError::StoreClosed)?;
75            if let Some(error) = inner.access.observer.error() {
76                return Err(error.into());
77            }
78            // A panic in the batch fails the staging with `DbError::Bug`, as
79            // a refused batch fails it.
80            let built = coven_foundation::task::caught_sync("stage files", || {
81                let mut batch = WriteBatch::new(inner.ids.as_ref());
82                let value = build(&mut batch)?;
83                Ok::<_, E>((batch.into_files(), value))
84            })
85            .unwrap_or_else(|bug| Err(E::from(DbError::Bug(bug))));
86            let prepared = inner.access.with_writer(|writer| {
87                let prepared = built.and_then(|(files, value)| {
88                    if files.is_empty() {
89                        return Ok((VecDeque::new(), value));
90                    }
91                    writer
92                        .transaction(|db| {
93                            let mut sources = VecDeque::new();
94                            // The staging's own folder, which its copies'
95                            // names name (Appendix C.3).
96                            let folder = FileName::new(inner.ids.new_id().to_string())
97                                .expect("UUID is a portable filename");
98                            for BatchFile { id, namespace, source, place } in files {
99                                inner
100                                    .access
101                                    .write_schema
102                                    .namespace_of(&namespace, Provenance::AppProvided)?;
103                                let copy = FileName::new(inner.ids.new_id().to_string())
104                                    .expect("UUID is a portable filename");
105                                let name = StagedName::new(&folder, &copy);
106                                if !crate::file_removals::record_unnamed(db, name.as_str())? {
107                                    return Err(DbError::FileNameReused {
108                                        name: name.as_str().to_owned(),
109                                    });
110                                }
111                                sources.push_back(StagingSource {
112                                    copy: StagingCopy { namespace, id, name, place },
113                                    source,
114                                });
115                            }
116                            Ok((sources, value))
117                        })
118                        .map_err(E::from)
119                });
120                match prepared {
121                    Ok((sources, value)) => {
122                        inner.staging.reserve(
123                            sources.iter().map(|s| OwnedFile::Staged(s.copy.name.clone())),
124                        );
125                        Ok::<_, E>(Ok((sources, value)))
126                    }
127                    Err(error) => Ok(Err((error, super::released(writer, inner)))),
128                }
129            })?;
130            // A refused batch's recorded names go once the writer is free.
131            match prepared {
132                Ok((sources, value)) => (sources, value, inner.staging.clone()),
133                Err((error, chosen)) => return inner.finish_removals(Err(error), chosen),
134            }
135        };
136        Ok((
137            Self {
138                reserved,
139                database,
140                lease: Some(lease),
141                names: sources.iter().map(|s| OwnedFile::Staged(s.copy.name.clone())).collect(),
142                sources,
143                staged: Vec::new(),
144            },
145            value,
146        ))
147    }
148
149    /// Write each app source to coven's copy of it in the staging's own
150    /// folder, made with the first copy, each synced with a plain sync, then
151    /// flush the batch: the folders and the drive's cache, once, so every
152    /// copy is on the drive before the commit that records it
153    /// (Appendix C.3). A failed flush fails the write. The bytes the app
154    /// already holds are written in one blocking task per run of them; a
155    /// stream's copy is created on the blocking pool and then read and
156    /// written as the stream yields. Only the creation of a copy failing on
157    /// the blocking pool ends the staging here: a panic in it, or the pool
158    /// dropping it unstarted as its runtime shuts down. The staging went
159    /// with it, and its drop released its names; their recorded removals
160    /// take any bytes written before (E15).
161    pub(super) async fn write(
162        mut self,
163    ) -> Result<(Self, Result<Flushed, DbError>), coven_foundation::task::TaskFailed> {
164        let (bugs, mut batch) = {
165            let slot = read_unpoisoned(&self.database);
166            let inner = slot.as_ref().expect("staging holds close guard");
167            let bugs = inner.access.directory.bugs().clone();
168            let batch = FlushBatch::new(bugs.clone());
169            #[cfg(test)]
170            let batch = batch
171                .with_hook(coven_foundation::locks::lock_unpoisoned(&inner.flush_hook).clone());
172            (bugs, batch)
173        };
174        let mut folder = None;
175        loop {
176            let Run { bytes: run, stream } = self.next_run();
177            if !run.is_empty() {
178                // The staging guard goes with the task, so cancellation
179                // cannot remove a name before this call creates its bytes.
180                let held = folder.take();
181                let (next, written) = coven_foundation::task::blocking_completing(
182                    "stage copies",
183                    bugs.clone(),
184                    crate::commit_scope::bind(move || {
185                        let written = (|| -> Result<_, DbError> {
186                            let held =
187                                held_folder(&self.database, held, &run[0].0.name, &mut batch)?;
188                            let mut staged = Vec::with_capacity(run.len());
189                            for (copy, bytes) in run {
190                                held.write_whole(&copy.name.copy(), &bytes, &mut batch)?;
191                                let mut hash = FileHasher::new();
192                                hash.update(&bytes);
193                                staged.push(copy.staged(bytes.len() as u64, hash));
194                            }
195                            Ok((held, batch, staged))
196                        })();
197                        Ok::<_, std::convert::Infallible>((self, written))
198                    }),
199                )
200                .await?
201                .unwrap_or_else(|never| match never {});
202                self = next;
203                match written {
204                    Ok((held, returned, staged)) => {
205                        folder = Some(held);
206                        batch = returned;
207                        self.staged.extend(staged);
208                    }
209                    Err(error) => return Ok((self, Err(error))),
210                }
211            }
212            let Some((copy, mut reader)) = stream else { break };
213            // Creation owns the staging guard until it finishes. Cancellation
214            // cannot remove a name before this blocking call creates its bytes.
215            let file_name = copy.name.clone();
216            let held = folder.take();
217            let mut creating = batch;
218            // The staging comes back with the writer, or with why it wasn't
219            // made, so no error reaches the sink unread: a dropped caller's
220            // staging releases its names as it drops.
221            let (next, writer, returned) = coven_foundation::task::blocking_completing(
222                "create a staged copy",
223                bugs.clone(),
224                crate::commit_scope::bind(move || {
225                    let writer = (|| -> Result<_, DbError> {
226                        let held = held_folder(&self.database, held, &file_name, &mut creating)?;
227                        let writer = held.create(&file_name.copy())?;
228                        Ok((held, writer))
229                    })();
230                    Ok::<_, std::convert::Infallible>((self, writer, creating))
231                }),
232            )
233            .await?
234            .unwrap_or_else(|never| match never {});
235            self = next;
236            batch = returned;
237            let writer = match writer {
238                Ok((held, writer)) => {
239                    folder = Some(held);
240                    writer
241                }
242                Err(error) => return Ok((self, Err(error))),
243            };
244            let mut hash = FileHasher::new();
245            let mut size = 0;
246            let written = |bytes: &[u8]| {
247                hash.update(bytes);
248                size += bytes.len() as u64;
249            };
250            if let Err(error) = writer.write_from(&mut reader, written, &mut batch).await {
251                return Ok((self, Err(error.into())));
252            }
253            self.staged.push(copy.staged(size, hash));
254        }
255        let flushed = batch.flush().await.map_err(DbError::from);
256        Ok((self, flushed))
257    }
258
259    /// The sources up to the next stream, in order, as the bytes the app
260    /// already holds, and that stream with its reader.
261    fn next_run(&mut self) -> Run {
262        let mut bytes = Vec::new();
263        while let Some(StagingSource { copy, source }) = self.sources.pop_front() {
264            match source {
265                FileSource::Bytes(held) => bytes.push((copy, held)),
266                FileSource::Stream(reader) => return Run { bytes, stream: Some((copy, reader)) },
267            }
268        }
269        Run { bytes, stream: None }
270    }
271
272    /// Fail the staging with `error`: its names go free and the bytes it
273    /// wrote are removed now, with any other released bytes, as a failed
274    /// write's are; a removal that fails joins the error and keeps its
275    /// record for the next write or open (Appendix C.3).
276    pub(super) fn discard<E: crate::WriteFailure>(mut self, error: DbError) -> E {
277        let database = self.database.clone();
278        let slot = read_unpoisoned(&database);
279        let inner = slot.as_ref().expect("staging holds close guard");
280        inner.release_file_reservations(&self.names);
281        self.names.clear();
282        let chosen =
283            inner.access.with_writer(|writer| Ok::<_, DbError>(super::released(writer, inner)));
284        let failures = match chosen {
285            Ok(chosen) => inner.remove_chosen(chosen),
286            Err(cleanup) => vec![cleanup],
287        };
288        let _lease = self.lease.take();
289        if failures.is_empty() {
290            E::from(error)
291        } else {
292            E::with_cleanup(E::from(error), failures)
293        }
294    }
295
296    /// The staged files, with what the batch returned, for the write that
297    /// puts them in rows. The names stay reserved until that write has
298    /// decided or the files are dropped; the lease close waits for ends here.
299    pub(super) fn into_staged<T>(mut self, value: T) -> StagedFiles<T> {
300        let _lease = self.lease.take();
301        StagedFiles {
302            reserved: self.reserved.clone(),
303            names: std::mem::take(&mut self.names),
304            files: std::mem::take(&mut self.staged),
305            value: Some(value),
306        }
307    }
308}
309
310/// The staging's folder, which its copies are made in: `held` once the
311/// staging has it, else made for `first`, its first copy, in the store
312/// `database` holds open, for `batch` to flush with the copies.
313fn held_folder(
314    database: &RwLock<Option<DatabaseInner>>,
315    held: Option<WriterFolder>,
316    first: &StagedName,
317    batch: &mut FlushBatch,
318) -> Result<WriterFolder, DbError> {
319    match held {
320        Some(held) => Ok(held),
321        None => {
322            let slot = read_unpoisoned(database);
323            let inner = slot.as_ref().expect("staging holds close guard");
324            Ok(inner.access.directory.staging_folder(&first.folder(), batch)?)
325        }
326    }
327}
328
329impl Drop for FileStaging {
330    fn drop(&mut self) {
331        let Some(_lease) = self.lease.take() else {
332            return;
333        };
334        self.reserved.release(&self.names);
335    }
336}
337
338/// App bytes `stage_files` wrote to coven's copies of them, under names
339/// reserved in this store, with what its batch returned, until a write puts
340/// them in rows (§17.2). It holds nothing close waits for. Dropped
341/// uncommitted, or after the store closes, its names go free, and the next
342/// write or open removes the bytes.
343pub struct StagedFiles<T> {
344    /// The names the store that staged these files reserves.
345    reserved: Arc<crate::reserved_names::ReservedNames>,
346    names: Vec<OwnedFile>,
347    files: Vec<StagedFile>,
348    value: Option<T>,
349}
350
351impl<T> StagedFiles<T> {
352    /// The names reserved for these files' copies.
353    pub(crate) fn names(&self) -> &[OwnedFile] {
354        &self.names
355    }
356
357    /// The files, their reserved names and the batch's value, for the write
358    /// taking them; dropping what is left releases nothing.
359    pub(crate) fn take(mut self) -> (Vec<StagedFile>, Vec<OwnedFile>, T) {
360        let value = self.value.take().expect("staged files are taken once");
361        (std::mem::take(&mut self.files), std::mem::take(&mut self.names), value)
362    }
363}
364
365impl<T> std::fmt::Debug for StagedFiles<T> {
366    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
367        f.debug_struct("StagedFiles").field("files", &self.files.len()).finish_non_exhaustive()
368    }
369}
370
371impl<T> Drop for StagedFiles<T> {
372    fn drop(&mut self) {
373        self.reserved.release(&self.names);
374    }
375}
376
377#[cfg(test)]
378#[path = "file_staging_tests.rs"]
379mod tests;