coven_database/store/store_session/
host_write_operation.rs

1use crate::DbError;
2use coven_foundation::store_dir::{PathTokenError, StoreDir};
3use coven_keys::encryption::EncryptionService;
4use coven_protocol::blob::BlobRef;
5use coven_protocol::write::WriteReceipt;
6
7use super::{SqlContext, StoreDatabase, StoreSession};
8
9pub struct WriteBatch {
10    new_blobs: Vec<NewBlob>,
11    deleted_blobs: Vec<BlobRef>,
12}
13
14impl WriteBatch {
15    pub fn new() -> Self {
16        Self {
17            new_blobs: Vec::new(),
18            deleted_blobs: Vec::new(),
19        }
20    }
21
22    pub fn put_blob(
23        &mut self,
24        namespace: impl Into<String>,
25        id: impl Into<String>,
26        bytes: impl Into<Vec<u8>>,
27    ) {
28        self.new_blobs.push(NewBlob {
29            namespace: namespace.into(),
30            id: id.into(),
31            bytes: bytes.into(),
32        });
33    }
34
35    pub fn delete_blob(&mut self, blob: BlobRef) {
36        self.deleted_blobs.push(blob);
37    }
38}
39
40impl Default for WriteBatch {
41    fn default() -> Self {
42        Self::new()
43    }
44}
45
46/// A local blob file that failed to unwind after a host write failed. Names the
47/// blob so a host learns which files are left behind, not just that some were.
48#[derive(Debug)]
49pub struct BlobFileFailure {
50    pub namespace: String,
51    pub id: String,
52    pub reason: coven_foundation::atomic_file::FileError,
53}
54
55impl std::fmt::Display for BlobFileFailure {
56    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
57        write!(formatter, "{}/{}: {}", self.namespace, self.id, self.reason)
58    }
59}
60
61/// Every blob that failed to unwind, in the order they were attempted.
62#[derive(Debug)]
63pub struct BlobFileFailures(pub Vec<BlobFileFailure>);
64
65impl std::fmt::Display for BlobFileFailures {
66    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
67        for (index, failure) in self.0.iter().enumerate() {
68            if index > 0 {
69                formatter.write_str("; ")?;
70            }
71            write!(formatter, "{failure}")?;
72        }
73        Ok(())
74    }
75}
76
77pub(crate) struct NewBlob {
78    namespace: String,
79    id: String,
80    bytes: Vec<u8>,
81}
82
83struct StagedBlob {
84    namespace: String,
85    id: String,
86    staged: Option<coven_foundation::local_file::AtomicStagedFile>,
87    published: Option<coven_foundation::local_file::PublishedAtomicFile>,
88}
89
90pub(crate) struct StagedBlobBatch {
91    blobs: Vec<StagedBlob>,
92}
93
94impl StagedBlob {
95    async fn stage<E>(store_dir: &StoreDir, blob: NewBlob) -> Result<Self, HostWriteError<E>> {
96        let destination = store_dir.local_blob_path(&blob.namespace, &blob.id)?;
97        let staged = store_dir
98            .stage_atomic_file(&destination)
99            .await
100            .map_err(HostWriteError::Blob)?;
101        let mut staged_blob = Self {
102            namespace: blob.namespace,
103            id: blob.id,
104            staged: Some(staged),
105            published: None,
106        };
107        if let Err(operation) = staged_blob.staged_mut().write_bytes(&blob.bytes).await {
108            let namespace = staged_blob.namespace.clone();
109            let id = staged_blob.id.clone();
110            return match staged_blob.discard().await {
111                Ok(()) => Err(HostWriteError::Blob(operation)),
112                Err(reason) => Err(HostWriteError::BlobCleanupFailed {
113                    operation: Box::new(HostWriteError::Blob(operation)),
114                    cleanup: BlobFileFailures(vec![BlobFileFailure {
115                        namespace,
116                        id,
117                        reason,
118                    }]),
119                }),
120            };
121        }
122        Ok(staged_blob)
123    }
124
125    fn staged_mut(&mut self) -> &mut coven_foundation::local_file::AtomicStagedFile {
126        self.staged.as_mut().expect("blob is staged")
127    }
128
129    fn publish(&mut self) -> Result<(), coven_foundation::atomic_file::FileError> {
130        let staged = self.staged.take().expect("blob is staged");
131        self.published = Some(staged.publish_for_transaction()?);
132        Ok(())
133    }
134
135    async fn discard(mut self) -> Result<(), coven_foundation::atomic_file::FileError> {
136        match self.staged.take() {
137            Some(staged) => staged.discard().await,
138            None => Ok(()),
139        }
140    }
141
142    fn rollback(mut self) -> Vec<coven_foundation::atomic_file::FileError> {
143        let mut failures = Vec::new();
144        if let Some(published) = self.published.take() {
145            if let Err(error) = published.rollback() {
146                failures.push(error);
147            }
148        }
149        if let Some(staged) = self.staged.take() {
150            if let Err(error) = staged.discard_blocking() {
151                failures.push(error);
152            }
153        }
154        failures
155    }
156
157    fn commit(mut self) {
158        assert!(self.staged.is_none(), "committed blob remains staged");
159        assert!(self.published.take().is_some(), "blob was not published");
160    }
161}
162
163impl StagedBlobBatch {
164    pub(crate) async fn stage<E>(
165        store_dir: &StoreDir,
166        blobs: Vec<NewBlob>,
167    ) -> Result<Self, HostWriteError<E>> {
168        let mut staged = Vec::new();
169        for blob in blobs {
170            match StagedBlob::stage(store_dir, blob).await {
171                Ok(blob) => staged.push(blob),
172                Err(error) => {
173                    return Err(Self { blobs: staged }
174                        .discard_after_stage_failure(error)
175                        .await);
176                }
177            }
178        }
179        Ok(Self { blobs: staged })
180    }
181
182    async fn discard_after_stage_failure<E>(
183        self,
184        operation: HostWriteError<E>,
185    ) -> HostWriteError<E> {
186        let mut failures = Vec::new();
187        for blob in self.blobs {
188            let (namespace, id) = (blob.namespace.clone(), blob.id.clone());
189            if let Err(reason) = blob.discard().await {
190                failures.push(BlobFileFailure {
191                    namespace: namespace.clone(),
192                    id,
193                    reason,
194                });
195            }
196        }
197        if failures.is_empty() {
198            operation
199        } else {
200            HostWriteError::BlobCleanupFailed {
201                operation: Box::new(operation),
202                cleanup: BlobFileFailures(failures),
203            }
204        }
205    }
206
207    pub(super) fn publish<E>(
208        &mut self,
209        mut validate: impl FnMut(&str, &str) -> Result<(), HostWriteError<E>>,
210    ) -> Result<(), HostWriteError<E>> {
211        for blob in &mut self.blobs {
212            validate(&blob.namespace, &blob.id)?;
213            blob.publish().map_err(HostWriteError::Blob)?;
214        }
215        Ok(())
216    }
217
218    pub(super) fn commit(self) {
219        for blob in self.blobs {
220            blob.commit();
221        }
222    }
223
224    pub(super) fn rollback<E>(self, write: HostWriteError<E>) -> HostWriteError<E> {
225        let mut failures = Vec::new();
226        for blob in self.blobs.into_iter().rev() {
227            let (namespace, id) = (blob.namespace.clone(), blob.id.clone());
228            for reason in blob.rollback() {
229                failures.push(BlobFileFailure {
230                    namespace: namespace.clone(),
231                    id: id.clone(),
232                    reason,
233                });
234            }
235        }
236        if failures.is_empty() {
237            write
238        } else {
239            HostWriteError::WriteRollbackFailed {
240                write: Box::new(write),
241                rollback: BlobFileFailures(failures),
242            }
243        }
244    }
245}
246
247pub(super) type HostSql<R, E> = Box<
248    dyn for<'context, 'connection> FnOnce(SqlContext<'context, 'connection>) -> Result<R, E> + Send,
249>;
250
251pub struct HostWriteOperation<R, E> {
252    batch: WriteBatch,
253    sql: HostSql<R, E>,
254}
255
256impl<R, E> HostWriteOperation<R, E> {
257    pub fn new(
258        batch: WriteBatch,
259        sql: impl for<'context, 'connection> FnOnce(SqlContext<'context, 'connection>) -> Result<R, E>
260            + Send
261            + 'static,
262    ) -> Self {
263        Self {
264            batch,
265            sql: Box::new(sql),
266        }
267    }
268}
269
270#[derive(Debug, thiserror::Error)]
271pub enum HostWriteError<E> {
272    #[error("host write closure failed: {0}")]
273    Host(#[source] E),
274    #[error("database write failed: {0}")]
275    Database(#[source] DbError),
276    #[error("local blob write failed: {0}")]
277    Blob(#[source] coven_foundation::atomic_file::FileError),
278    #[error("blob declaration failed: {0}")]
279    BlobDeclaration(#[source] crate::BlobDeclError),
280    #[error("unsafe blob path: {0}")]
281    UnsafeBlobPath(#[source] PathTokenError),
282    #[error("the host write closure panicked")]
283    WriteClosurePanicked,
284    #[error(
285        "write failed: {write}; failed to remove installed local blobs during rollback: {rollback}"
286    )]
287    WriteRollbackFailed {
288        #[source]
289        write: Box<Self>,
290        rollback: BlobFileFailures,
291    },
292    #[error("write failed: {operation}; failed to remove unpublished local blobs: {cleanup}")]
293    BlobCleanupFailed {
294        #[source]
295        operation: Box<Self>,
296        cleanup: BlobFileFailures,
297    },
298    #[error("blob {namespace}/{id} is still referenced by a row after the write")]
299    BlobStillReferenced { namespace: String, id: String },
300    #[error("blob {namespace}/{id} is already referenced by a row")]
301    BlobAlreadyReferenced { namespace: String, id: String },
302    #[error("blob {namespace}/{id} is owned by an unpublished write")]
303    BlobOwnedByPendingWrite { namespace: String, id: String },
304    #[error("host write I/O failed: {0}")]
305    Io(#[source] std::io::Error),
306}
307
308impl<E> From<DbError> for HostWriteError<E> {
309    fn from(value: DbError) -> Self {
310        Self::Database(value)
311    }
312}
313
314impl<E> From<PathTokenError> for HostWriteError<E> {
315    fn from(value: PathTokenError) -> Self {
316        Self::UnsafeBlobPath(value)
317    }
318}
319
320impl<E> From<std::io::Error> for HostWriteError<E> {
321    fn from(value: std::io::Error) -> Self {
322        Self::Io(value)
323    }
324}
325
326#[derive(Clone)]
327pub struct StoreRowWrites {
328    database: StoreDatabase,
329}
330
331impl StoreSession<'_> {
332    fn execute_host_write<R, E>(
333        &mut self,
334        staged: StagedBlobBatch,
335        deleted: Vec<BlobRef>,
336        sql: HostSql<R, E>,
337        routing_encryption: Option<EncryptionService>,
338        blob_staging: Option<Box<dyn crate::AudienceBlobMoveStaging>>,
339        write_id: coven_protocol::write::WriteId,
340    ) -> Result<Result<WriteReceipt<R>, HostWriteError<E>>, DbError> {
341        let verified_authority = &mut *self.verified_store_authority;
342        let stamper = coven_protocol::hlc::UpdatedAtStamper::new(self.hlc.clone());
343        let result = super::host_write_capture::CapturedStoreWriteTransaction::begin_host(
344            self.conn,
345            self.store_dir,
346            self.synced_tables,
347            self.gates,
348            self.blob_decls,
349            routing_encryption.as_ref(),
350            blob_staging.as_deref(),
351            verified_authority,
352            write_id,
353        )
354        .map_err(HostWriteError::from)
355        .and_then(|transaction| transaction.execute_host(staged, deleted, sql, stamper));
356        Ok(result)
357    }
358}
359
360impl StoreRowWrites {
361    pub fn new(database: StoreDatabase) -> Self {
362        Self { database }
363    }
364
365    pub fn requires_routing_encryption(&self) -> bool {
366        self.database.has_scoped_graph()
367    }
368
369    pub fn subscribe_committed_changes(
370        &self,
371    ) -> tokio::sync::broadcast::Receiver<std::sync::Arc<crate::CommittedChanges>> {
372        self.database.subscribe_committed_changes()
373    }
374
375    pub async fn pending_writes(
376        &self,
377    ) -> Result<Vec<coven_protocol::write::PendingWrite>, DbError> {
378        self.database.pending_writes().await
379    }
380
381    pub async fn blocked_writes(
382        &self,
383    ) -> Result<Vec<coven_protocol::write::PendingWrite>, DbError> {
384        self.database.blocked_writes().await
385    }
386
387    pub async fn retry_blocked_write(
388        &self,
389        write_id: &coven_protocol::write::WriteId,
390    ) -> Result<Vec<coven_protocol::write::WriteId>, DbError> {
391        self.database.retry_blocked_write(write_id).await
392    }
393
394    pub async fn discard_blocked_write(
395        &self,
396        write_id: &coven_protocol::write::WriteId,
397    ) -> Result<super::BlockedWriteDiscard, DbError> {
398        self.database.discard_blocked_write(write_id).await
399    }
400
401    pub async fn write_status(
402        &self,
403        write_id: &coven_protocol::write::WriteId,
404    ) -> Result<coven_protocol::write::WriteStatus, DbError> {
405        self.database.write_status(write_id).await
406    }
407
408    pub async fn subscribe_write_status(
409        &self,
410        write_id: &coven_protocol::write::WriteId,
411    ) -> Result<tokio::sync::watch::Receiver<coven_protocol::write::WriteStatus>, DbError> {
412        self.database.subscribe_write_status(write_id).await
413    }
414
415    pub async fn execute<R, E>(
416        &self,
417        operation: HostWriteOperation<R, E>,
418        routing_encryption: Option<EncryptionService>,
419        blob_staging: Option<Box<dyn crate::AudienceBlobMoveStaging>>,
420    ) -> Result<WriteReceipt<R>, HostWriteError<E>>
421    where
422        R: Send + 'static,
423        E: Send + 'static,
424    {
425        let database = &self.database;
426        let HostWriteOperation { batch, sql } = operation;
427        let staged = database.stage_host_write_blobs(batch.new_blobs).await?;
428        let write_id = database.new_store_write_id();
429        let deleted = batch.deleted_blobs;
430
431        let outcome = database
432            .call_store(move |session| {
433                session.execute_host_write(
434                    staged,
435                    deleted,
436                    sql,
437                    routing_encryption,
438                    blob_staging,
439                    write_id,
440                )
441            })
442            .await;
443
444        let receipt = match outcome {
445            Ok(Ok(receipt)) => receipt,
446            Ok(Err(error)) => return Err(error),
447            Err(error) => return Err(HostWriteError::Database(error)),
448        };
449
450        if let Err(error) = super::local_blob_cleanup::LocalBlobCleanup::new(database)
451            .drain()
452            .await
453        {
454            tracing::warn!(
455                error = %error,
456                "failed to drain local blob cleanup intents after write commit"
457            );
458        }
459        Ok(receipt)
460    }
461
462    #[cfg(any(test, feature = "test-utils"))]
463    pub async fn store_write_partition_for_test(
464        &self,
465        write_id: &coven_protocol::write::WriteId,
466    ) -> Result<Vec<u8>, DbError> {
467        self.database.store_write_partition_for_test(write_id).await
468    }
469
470    #[cfg(any(test, feature = "test-utils"))]
471    pub async fn write_blob_lease_count_for_test(
472        &self,
473        write_id: &coven_protocol::write::WriteId,
474    ) -> Result<i64, DbError> {
475        self.database
476            .write_blob_lease_count_for_test(write_id)
477            .await
478    }
479
480    #[cfg(any(test, feature = "test-utils"))]
481    pub async fn store_write_journal_counts_for_test(&self) -> Result<(i64, i64), DbError> {
482        self.database.store_write_journal_counts_for_test().await
483    }
484
485    #[cfg(any(test, feature = "test-utils"))]
486    pub async fn cleanup_intent_count_for_test(
487        &self,
488        namespace: &str,
489        blob_id: &str,
490    ) -> Result<i64, DbError> {
491        self.database
492            .cleanup_intent_count_for_test(namespace, blob_id)
493            .await
494    }
495
496    #[cfg(any(test, feature = "test-utils"))]
497    pub async fn coven_table_exists_for_test(
498        &self,
499        table: crate::DatabaseTestTable,
500    ) -> Result<bool, DbError> {
501        self.database.coven_table_exists_for_test(table).await
502    }
503
504    #[cfg(any(test, feature = "test-utils"))]
505    pub async fn install_store_write_failure_trigger_for_test(&self) -> Result<(), DbError> {
506        self.database
507            .install_store_write_failure_trigger_for_test()
508            .await
509    }
510
511    #[cfg(any(test, feature = "test-utils"))]
512    pub async fn remove_store_write_failure_trigger_for_test(&self) -> Result<(), DbError> {
513        self.database
514            .remove_store_write_failure_trigger_for_test()
515            .await
516    }
517
518    #[cfg(any(test, feature = "test-utils"))]
519    pub async fn write_blob_facts_for_test(
520        &self,
521        write_id: coven_protocol::write::WriteId,
522    ) -> Result<String, DbError> {
523        self.database.write_blob_facts_for_test(write_id).await
524    }
525}
coven_database/store/store_session/host_write_operation.rs

coven_database/store/store_session/
host_write_operation.rs
