coven_database/store/store_session/

local_blob_cleanup.rs

use tracing::debug;

use super::*;
use crate::local_blob_cleanup_intents::{LocalBlobCleanupIdentity, LocalBlobCleanupIntent};
use crate::BlobDecls;

pub(crate) struct ExactBlobBindings {
    by_row: std::collections::BTreeMap<(String, String), coven_protocol::store_commit::ObjectHash>,
}

pub(crate) fn exact_blob_bindings_on(
    conn: &rusqlite::Connection,
) -> Result<ExactBlobBindings, DbError> {
    let mut statement = conn
        .prepare(
            "SELECT binding.table_name, binding.row_id, locator.locator_hash
             FROM row_blob_locators AS binding
             JOIN blob_locators AS locator
               ON locator.remote_object_id = binding.remote_object_id",
        )
        .map_err(DbError::from)?;
    let rows = statement
        .query_map([], |row| {
            Ok((
                row.get::<_, String>(0)?,
                row.get::<_, String>(1)?,
                row.get::<_, String>(2)?,
            ))
        })
        .map_err(DbError::from)?;
    let mut by_row = std::collections::BTreeMap::new();
    for row in rows {
        let (table, row_id, encoded) = row.map_err(DbError::from)?;
        let locator_hash = encoded
            .parse::<coven_protocol::store_commit::ObjectHash>()
            .map_err(|error| DbError::context("parse local cleanup locator hash", error))?;
        if let Some(existing) = by_row.insert((table.clone(), row_id.clone()), locator_hash) {
            if existing != locator_hash {
                return Err(DbError::Message(format!(
                    "local cleanup for {table}.{row_id} has distinct exact locator bindings"
                )));
            }
        }
    }
    Ok(ExactBlobBindings { by_row })
}

/// Record cleanup obligations for each copy identity no live row needs in this
/// transaction. The transaction mutating the carrying rows must also record the
/// obligation, so the obsolete state and its cleanup commit together.
pub(crate) fn record_obsolete_copy_intents_on(
    conn: &rusqlite::Connection,
    decls: &BlobDecls,
    intent: &LocalBlobCleanupIntent,
) -> Result<(), DbError> {
    match intent.identity() {
        LocalBlobCleanupIdentity::Local => {
            let local_referenced = decls
                .local_copy_is_referenced(conn, intent.namespace(), intent.blob_id())
                .map_err(DbError::from)?;
            if !local_referenced {
                record_durable_intent(conn, intent)?;
            }
        }
        LocalBlobCleanupIdentity::Exact(_) => {
            return Err(DbError::Message(
                "exact local cleanup identity is already durable".to_string(),
            ));
        }
        LocalBlobCleanupIdentity::Row { table, row_id } => {
            record_obsolete_row_copy_intents_on(
                conn,
                decls,
                intent,
                exact_blob_binding_for_row_on(conn, table, row_id)?,
            )?;
        }
    }
    Ok(())
}

fn exact_blob_binding_for_row_on(
    conn: &rusqlite::Connection,
    table: &str,
    row_id: &str,
) -> Result<Option<coven_protocol::store_commit::ObjectHash>, DbError> {
    let mut statement = conn
        .prepare(
            "SELECT DISTINCT locator.locator_hash
             FROM row_blob_locators AS binding
             JOIN blob_locators AS locator
               ON locator.remote_object_id = binding.remote_object_id
             WHERE binding.table_name = ?1 AND binding.row_id = ?2",
        )
        .map_err(DbError::from)?;
    let locator_hashes = statement
        .query_map((table, row_id), |row| row.get::<_, String>(0))
        .map_err(DbError::from)?
        .collect::<Result<Vec<_>, _>>()
        .map_err(DbError::from)?;
    match locator_hashes.as_slice() {
        [] => Ok(None),
        [encoded] => encoded
            .parse()
            .map(Some)
            .map_err(|error| DbError::context("parse local cleanup locator hash", error)),
        _ => Err(DbError::Message(format!(
            "local cleanup for {table}.{row_id} has {} distinct exact locator bindings",
            locator_hashes.len()
        ))),
    }
}

pub(crate) fn record_obsolete_copy_intents_from_bindings_on(
    conn: &rusqlite::Connection,
    decls: &BlobDecls,
    intent: &LocalBlobCleanupIntent,
    bindings: &ExactBlobBindings,
) -> Result<(), DbError> {
    match intent.identity() {
        LocalBlobCleanupIdentity::Row { table, row_id } => record_obsolete_row_copy_intents_on(
            conn,
            decls,
            intent,
            bindings
                .by_row
                .get(&(table.clone(), row_id.clone()))
                .copied(),
        ),
        _ => record_obsolete_copy_intents_on(conn, decls, intent),
    }
}

fn record_obsolete_row_copy_intents_on(
    conn: &rusqlite::Connection,
    decls: &BlobDecls,
    intent: &LocalBlobCleanupIntent,
    exact_locator_hash: Option<coven_protocol::store_commit::ObjectHash>,
) -> Result<(), DbError> {
    if let Some(locator_hash) = exact_locator_hash {
        let exact =
            LocalBlobCleanupIntent::exact(intent.namespace(), intent.blob_id(), locator_hash);
        let referenced = decls
            .exact_copy_is_referenced(conn, exact.namespace(), exact.blob_id(), locator_hash)
            .map_err(DbError::from)?;
        if !referenced {
            record_durable_intent(conn, &exact)?;
        }
    }
    let local_referenced = decls
        .local_copy_is_referenced(conn, intent.namespace(), intent.blob_id())
        .map_err(DbError::from)?;
    if !local_referenced {
        record_durable_intent(
            conn,
            &LocalBlobCleanupIntent::local(intent.namespace(), intent.blob_id()),
        )?;
    }
    Ok(())
}

fn record_durable_intent(
    conn: &rusqlite::Connection,
    intent: &LocalBlobCleanupIntent,
) -> Result<(), DbError> {
    let persisted_identity = intent.persisted_identity()?;
    let inserted = crate::with_coven_sql_authority(|| {
        conn.execute(
            "INSERT OR IGNORE INTO local_cleanup_intents (namespace, blob_id, copy_identity)
             VALUES (?1, ?2, ?3)",
            (intent.namespace(), intent.blob_id(), persisted_identity),
        )
        .map_err(DbError::from)
    })?;
    if inserted == 0 {
        debug!(
            namespace = %intent.namespace(),
            blob_id = %intent.blob_id(),
            "local blob cleanup intent already exists"
        );
    }
    Ok(())
}

pub(crate) struct SuspendedBlobCleanup {
    local: Vec<LocalBlobCleanupIntent>,
    published: Vec<super::blob_outbox::PublishedBlobDropIntent>,
}

/// Temporarily remove cleanup obligations for blobs whose bytes a replay lease
/// still owns. The caller must reevaluate every returned intent against the
/// resulting rows before committing its transaction.
pub(crate) fn suspend_leased_blob_cleanup_for_restoration_on(
    conn: &rusqlite::Connection,
    blobs: &[coven_protocol::blob::BlobRef],
) -> Result<SuspendedBlobCleanup, DbError> {
    let blob_keys = blobs
        .iter()
        .map(|blob| (blob.namespace.as_str(), blob.id.as_str()))
        .collect::<std::collections::BTreeSet<_>>();
    let mut taken = Vec::new();
    for (namespace, blob_id) in blob_keys {
        let removed = crate::with_coven_sql_authority(|| {
            conn.execute(
                "DELETE FROM local_cleanup_intents
                 WHERE namespace = ?1 AND blob_id = ?2 AND copy_identity = 'local'
                   AND (
                       EXISTS (
                           SELECT 1 FROM store_write_blob_leases
                           WHERE namespace = ?1 AND blob_id = ?2
                       ) OR EXISTS (
                           SELECT 1 FROM retained_replay_blob_leases
                           WHERE namespace = ?1 AND blob_id = ?2
                       )
                   )",
                (namespace, blob_id),
            )
            .map_err(DbError::from)
        })?;
        match removed {
            0 => {}
            1 => taken.push(LocalBlobCleanupIntent::local(namespace, blob_id)),
            count => {
                return Err(DbError::Message(format!(
                "local cleanup restoration removed {count} obligations for {namespace}/{blob_id}"
            )))
            }
        }
    }
    let published = super::blob_outbox::take_leased_published_blob_drop_intents_for_restoration_on(
        conn, blobs,
    )?;
    Ok(SuspendedBlobCleanup {
        local: taken,
        published,
    })
}

/// Cancel a suspended obligation when the restored rows need its local source,
/// or put it back when the completed replay still leaves the source obsolete.
pub(crate) fn reevaluate_suspended_blob_cleanup_on(
    conn: &rusqlite::Connection,
    decls: &BlobDecls,
    cleanup: &SuspendedBlobCleanup,
) -> Result<(), DbError> {
    for intent in &cleanup.local {
        record_obsolete_copy_intents_on(conn, decls, intent)?;
    }
    for intent in &cleanup.published {
        let local_referenced = decls
            .local_copy_is_referenced(conn, &intent.drop.namespace, &intent.drop.id)
            .map_err(DbError::from)?;
        if !local_referenced {
            super::blob_outbox::reinsert_published_blob_drop_intent_on(conn, intent)?;
        }
    }
    Ok(())
}

pub(crate) fn local_blob_cleanup_intents_on(
    conn: &rusqlite::Connection,
) -> Result<Vec<(LocalBlobCleanupIntent, bool)>, DbError> {
    let mut statement = conn
        .prepare(
            "SELECT intent.namespace, intent.blob_id, intent.copy_identity, EXISTS (
                     SELECT 1 FROM store_write_blob_leases lease
                     WHERE lease.namespace = intent.namespace
                       AND lease.blob_id = intent.blob_id
                       AND intent.copy_identity = 'local'
                 )
                     OR EXISTS (
                         SELECT 1 FROM retained_replay_blob_leases baseline
                         WHERE baseline.namespace = intent.namespace
                           AND baseline.blob_id = intent.blob_id
                           AND intent.copy_identity = 'local'
                     )
                 FROM local_cleanup_intents intent
                 ORDER BY namespace, blob_id,
                          CASE WHEN copy_identity = 'local' THEN 1 ELSE 0 END,
                          copy_identity",
        )
        .map_err(DbError::from)?;
    let rows = statement
        .query_map([], |row| {
            Ok((
                LocalBlobCleanupIntent::from_persisted(
                    row.get::<_, String>(0)?,
                    row.get::<_, String>(1)?,
                    row.get::<_, String>(2)?,
                )
                .map_err(|error| {
                    rusqlite::Error::FromSqlConversionFailure(
                        2,
                        rusqlite::types::Type::Text,
                        Box::new(error),
                    )
                })?,
                row.get::<_, bool>(3)?,
            ))
        })
        .map_err(DbError::from)?;
    rows.collect::<Result<Vec<_>, _>>().map_err(DbError::from)
}

pub(crate) fn complete_local_blob_cleanup_on(
    conn: &rusqlite::Connection,
    namespace: &str,
    blob_id: &str,
    persisted_identity: &str,
) -> Result<(), DbError> {
    conn.execute(
        "DELETE FROM local_cleanup_intents
                 WHERE namespace = ?1 AND blob_id = ?2 AND copy_identity = ?3",
        (namespace, blob_id, persisted_identity),
    )
    .map(|_| ())
    .map_err(DbError::from)
}

pub struct LocalBlobCleanup<'operation> {
    database: &'operation StoreDatabase,
}

impl<'operation> LocalBlobCleanup<'operation> {
    pub fn new(database: &'operation StoreDatabase) -> Self {
        Self { database }
    }

    /// Drain every committed cleanup obligation. A filesystem or database
    /// failure leaves the intent durable and fails the operation. `true` means
    /// every remaining intent is blocked by an active Store-write lease.
    pub async fn drain(&self) -> Result<bool, DbError> {
        let database = self.database;
        #[cfg(any(test, feature = "test-utils"))]
        database
            .reach_test_point(crate::DatabaseTestPoint::LocalBlobCleanupRequested)
            .await;
        let _cleanup_guard = database.local_blob_cleanup_permit().await;
        #[cfg(any(test, feature = "test-utils"))]
        database
            .reach_test_point(crate::DatabaseTestPoint::LocalBlobCleanupAcquired)
            .await;

        let intents = database
            .call_database(|session| session.local_blob_cleanup_intents())
            .await?;

        let mut pending = false;
        for (intent, leased) in intents {
            if leased {
                pending = true;
                debug!(
                    namespace = %intent.namespace(),
                    blob_id = %intent.blob_id(),
                    "local blob cleanup is blocked by an active Store-write lease"
                );
                continue;
            }
            #[cfg(any(test, feature = "test-utils"))]
            database
                .reach_test_point(crate::DatabaseTestPoint::LocalBlobCleanupBeforeFilesystem {
                    namespace: intent.namespace().to_string(),
                    blob_id: intent.blob_id().to_string(),
                })
                .await;
            let persisted_identity = intent.persisted_identity()?;
            database.apply_local_blob_cleanup_intent(&intent).await?;

            let namespace = intent.namespace().to_string();
            let blob_id = intent.blob_id().to_string();
            database
                .call_database(move |session| {
                    session.complete_local_blob_cleanup(&namespace, &blob_id, &persisted_identity)
                })
                .await?;
        }
        #[cfg(any(test, feature = "test-utils"))]
        database
            .reach_test_point(crate::DatabaseTestPoint::LocalBlobCleanupFinished)
            .await;
        Ok(pending)
    }
}

#[cfg(test)]
impl StoreSession<'_> {
    fn record_obsolete_copy_intent_for_test(
        &self,
        intent: &LocalBlobCleanupIntent,
    ) -> Result<(), DbError> {
        record_obsolete_copy_intents_on(self.conn, self.blob_decls, intent)
    }
}

#[cfg(test)]
impl StoreDatabase {
    async fn record_obsolete_copy_intent_for_test(
        &self,
        intent: LocalBlobCleanupIntent,
    ) -> Result<(), DbError> {
        self.call_store(move |session| session.record_obsolete_copy_intent_for_test(&intent))
            .await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::synthetic_store::open_test_db_with_blob;
    use coven_protocol::blob::{CacheFill, Provenance};
    use coven_protocol::store_commit::ObjectHash;
    use coven_protocol::synced_schema::BlobDecl;

    #[tokio::test]
    async fn a_live_same_id_row_with_another_locator_does_not_suppress_exact_cleanup() {
        let store_dir = crate::synthetic_store::test_store_dir();
        let db = open_test_db_with_blob(
            store_dir,
            BlobDecl::new("photos", Provenance::HostProvided, CacheFill::CacheEager)
                .with_id_column("blob_id"),
        );
        let removed_locator = ObjectHash::digest(b"removed locator");
        let live_locator = ObjectHash::digest(b"live locator");
        let removed_object = ObjectHash::digest(b"removed object");
        let live_object = ObjectHash::digest(b"live object");
        let database = StoreDatabase::new(&db);

        db.seed_distinct_cleanup_bindings_for_test(
            removed_locator,
            live_locator,
            removed_object,
            live_object,
        )
        .await
        .expect("seed removed and live blob bindings");

        database
            .record_obsolete_copy_intent_for_test(LocalBlobCleanupIntent::for_row(
                "photos",
                "shared-id",
                "note_photos",
                "removed-row",
            ))
            .await
            .expect("record obsolete row copies");

        db.cleanup_intent_copy_identities_for_test()
            .await
            .map(|identities| {
                assert_eq!(
                    identities,
                    [removed_locator.to_string(), "local".to_string()]
                );
            })
            .expect("record exact cleanup despite a live same-id row");
    }
}