1use 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
17struct StagingSource {
20 copy: StagingCopy,
21 source: FileSource,
22}
23
24struct 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
34struct Run {
37 bytes: Vec<(StagingCopy, Vec<u8>)>,
38 stream: Option<(StagingCopy, StagingReader)>,
39}
40
41impl StagingCopy {
42 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 reserved: Arc<crate::reserved_names::ReservedNames>,
59 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 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 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, ©);
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 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 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 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(©.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 let file_name = copy.name.clone();
216 let held = folder.take();
217 let mut creating = batch;
218 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 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 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 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
310fn 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
338pub struct StagedFiles<T> {
344 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 pub(crate) fn names(&self) -> &[OwnedFile] {
354 &self.names
355 }
356
357 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;