core_crypto_keystore/transaction/

mod.rs

use std::{
    borrow::Cow,
    collections::{HashMap, HashSet, hash_map::Entry},
    sync::Arc,
};

use async_lock::{RwLock, SemaphoreGuardArc};
use itertools::Itertools;

use crate::{
    CryptoKeystoreError, CryptoKeystoreResult, Database,
    entities::{MlsPendingMessage, PersistedMlsGroup},
    traits::{BorrowPrimaryKey, Entity, EntityDatabaseMutation, EntityDeleteBorrowed, KeyType, SearchableEntity},
    transaction::dynamic_dispatch::EntityId,
};

pub(crate) mod dynamic_dispatch;

/// table: primary key -> entity reference
type InMemoryTable = HashMap<EntityId, dynamic_dispatch::Entity>;
/// collection: collection name -> table
type InMemoryCollection = Arc<RwLock<HashMap<&'static str, InMemoryTable>>>;

/// This represents a transaction, where all operations will be done in memory and committed at the
/// end
#[derive(Debug, Clone)]
pub struct KeystoreTransaction {
    cache: InMemoryCollection,
    deleted: Arc<RwLock<HashSet<EntityId>>>,
    _semaphore_guard: Arc<SemaphoreGuardArc>,
}

impl KeystoreTransaction {
    /// Instantiate a new transaction.
    ///
    /// Requires a semaphore guard to ensure that only one exists at a time.
    pub(crate) async fn new(semaphore_guard: SemaphoreGuardArc) -> CryptoKeystoreResult<Self> {
        Ok(Self {
            cache: Default::default(),
            deleted: Arc::new(Default::default()),
            _semaphore_guard: Arc::new(semaphore_guard),
        })
    }

    /// Save an entity into this transaction.
    ///
    /// This is a multi-step process:
    ///
    /// - Adjust the entity by calling its [`pre_save()`][Entity::pre_save] method.
    /// - Store the entity in an internal map.
    ///   - Remove the entity from the set of deleted entities, if it was there.
    /// - On [`Self::commit`], actually persist the entity into the supplied database.
    pub(crate) async fn save<E>(&self, mut entity: E) -> CryptoKeystoreResult<E::AutoGeneratedFields>
    where
        E: Entity + EntityDatabaseMutation + Send + Sync,
    {
        let auto_generated_fields = entity.pre_save()?;

        let entity_id =
            EntityId::from_entity(&entity).ok_or(CryptoKeystoreError::UnknownCollectionName(E::COLLECTION_NAME))?;
        {
            // start by adding the entity
            let mut cache_guard = self.cache.write().await;
            let table = cache_guard.entry(E::COLLECTION_NAME).or_default();
            table.insert(entity_id.clone(), entity.into());
        }
        {
            // at this point remove the entity from the set of deleted entities to ensure that
            // this new data gets propagated
            let mut cache_guard = self.deleted.write().await;
            cache_guard.remove(&entity_id);
        }

        Ok(auto_generated_fields)
    }

    async fn remove_by_entity_id<E>(&self, entity_id: EntityId) -> CryptoKeystoreResult<()>
    where
        E: Entity + EntityDatabaseMutation,
    {
        // rm this entity from the set of added/modified items
        // it might never touch the real db at all
        let mut cache_guard = self.cache.write().await;
        if let Entry::Occupied(mut table) = cache_guard.entry(E::COLLECTION_NAME)
            && let Entry::Occupied(cached_record) = table.get_mut().entry(entity_id.clone())
        {
            cached_record.remove_entry();
        };

        // add this entity to the set of items which should be deleted from the persisted db
        let mut deleted_set = self.deleted.write().await;
        deleted_set.insert(entity_id);
        Ok(())
    }

    /// Remove an entity by its primary key.
    ///
    /// Where the primary key has a distinct borrowed form, consider [`Self::remove_borrowed`].
    ///
    /// Note that this doesn't return whether or not anything was actually removed because
    /// that won't happen until the transaction is committed.
    pub(crate) async fn remove<E>(&self, id: &E::PrimaryKey) -> CryptoKeystoreResult<()>
    where
        E: Entity + EntityDatabaseMutation,
    {
        let entity_id = EntityId::from_primary_key::<E>(id)
            .ok_or(CryptoKeystoreError::UnknownCollectionName(E::COLLECTION_NAME))?;
        self.remove_by_entity_id::<E>(entity_id).await
    }

    /// Remove an entity by the borrowed form of its primary key.
    ///
    /// Note that this doesn't return whether or not anything was actually removed because
    /// that won't happen until the transaction is committed.
    pub(crate) async fn remove_borrowed<E>(&self, id: &E::BorrowedPrimaryKey) -> CryptoKeystoreResult<()>
    where
        E: EntityDeleteBorrowed + BorrowPrimaryKey,
    {
        let entity_id = EntityId::from_borrowed_primary_key::<E>(id)
            .ok_or(CryptoKeystoreError::UnknownCollectionName(E::COLLECTION_NAME))?;
        self.remove_by_entity_id::<E>(entity_id).await
    }

    /// Restore an entity that was deleted in this transaction by removing it from the deleted list. This is
    /// idempotent: if the entity doesn't exist in the deleted list, do nothing.
    ///
    /// NOTE: This will only work if the entity has been added in an earlier transaction, because otherwise,
    /// removing its id from the deleted list wouldn't suffice: we'd need to replay its insertion.
    pub(crate) async fn restore<E>(&self, id: &E::BorrowedPrimaryKey) -> CryptoKeystoreResult<()>
    where
        E: EntityDeleteBorrowed + BorrowPrimaryKey,
    {
        let entity_id = EntityId::from_borrowed_primary_key::<E>(id)
            .ok_or(CryptoKeystoreError::UnknownCollectionName(E::COLLECTION_NAME))?;
        let mut deleted = self.deleted.write().await;
        deleted.remove(&entity_id);
        Ok(())
    }

    pub(crate) async fn child_groups(
        &self,
        entity: PersistedMlsGroup,
        persisted_records: impl IntoIterator<Item = PersistedMlsGroup>,
    ) -> CryptoKeystoreResult<Vec<PersistedMlsGroup>> {
        // First get all raw groups from the cache, then filter by their parent id
        let cached_records = self.find_all_in_cache::<PersistedMlsGroup>().await;
        let cached_records = cached_records
            .iter()
            .filter(|maybe_child| {
                maybe_child
                    .parent_id
                    .as_deref()
                    .map(|parent_id| parent_id == entity.borrow_primary_key().bytes().as_ref())
                    .unwrap_or_default()
            })
            .map(Arc::as_ref)
            .map(Cow::Borrowed);

        let persisted_records = persisted_records.into_iter().map(Cow::Owned);

        Ok(self.merge_records(cached_records, persisted_records).await)
    }

    pub(crate) async fn remove_pending_messages_by_conversation_id(&self, conversation_id: impl AsRef<[u8]> + Send) {
        let conversation_id = conversation_id.as_ref();

        let mut cache_guard = self.cache.write().await;
        if let Entry::Occupied(mut table) = cache_guard.entry(MlsPendingMessage::COLLECTION_NAME) {
            table.get_mut().retain(|_key, entity| {
                let pending_message = entity
                    .downcast::<MlsPendingMessage>()
                    .expect("table for MlsPendingMessage contains only that type");
                pending_message.foreign_id != conversation_id
            });
        }
        drop(cache_guard);

        let mut deleted_set = self.deleted.write().await;
        deleted_set.insert(
            EntityId::from_key::<MlsPendingMessage>(conversation_id.into())
                .expect("mls pending messages are proper entities which can be parsed"),
        );
    }

    pub(crate) async fn find_pending_messages_by_conversation_id(
        &self,
        conversation_id: &[u8],
        persisted_records: impl IntoIterator<Item = MlsPendingMessage>,
    ) -> CryptoKeystoreResult<Vec<MlsPendingMessage>> {
        let persisted_records = persisted_records.into_iter().map(Cow::Owned);

        let cached_records = self.find_all_in_cache::<MlsPendingMessage>().await;
        let cached_records = cached_records
            .iter()
            .filter(|pending_message| pending_message.foreign_id == conversation_id)
            .map(Arc::as_ref)
            .map(Cow::Borrowed);

        let merged_records = self.merge_records(cached_records, persisted_records).await;
        Ok(merged_records)
    }

    async fn find_in_cache<E>(&self, entity_id: &EntityId) -> Option<Arc<E>>
    where
        E: 'static + Entity + Send + Sync,
    {
        let cache_guard = self.cache.read().await;
        cache_guard
            .get(E::COLLECTION_NAME)
            .and_then(|table| table.get(entity_id).and_then(|entity| entity.downcast()))
    }

    /// The result of this function will have different contents for different scenarios:
    /// * `Some(Some(E))` - the transaction cache contains the record
    /// * `Some(None)` - the deletion of the record has been cached
    /// * `None` - there is no information about the record in the cache
    async fn get_by_entity_id<E>(&self, entity_id: &EntityId) -> Option<Option<Arc<E>>>
    where
        E: 'static + Entity + Send + Sync,
    {
        // when applying our transaction to the real database, we delete after inserting,
        // so here we have to check for deletion before we check for existing values
        let deleted_list = self.deleted.read().await;
        if deleted_list.contains(entity_id) {
            return Some(None);
        }

        self.find_in_cache::<E>(entity_id).await.map(Some)
    }

    /// The result of this function will have different contents for different scenarios:
    /// * `Some(Some(E))` - the transaction cache contains the record
    /// * `Some(None)` - the deletion of the record has been cached
    /// * `None` - there is no information about the record in the cache
    pub(crate) async fn get<E>(&self, id: &E::PrimaryKey) -> Option<Option<Arc<E>>>
    where
        E: 'static + Entity + Send + Sync,
    {
        let entity_id = EntityId::from_primary_key::<E>(id)?;
        self.get_by_entity_id(&entity_id).await
    }

    /// The result of this function will have different contents for different scenarios:
    /// * `Some(Some(E))` - the transaction cache contains the record
    /// * `Some(None)` - the deletion of the record has been cached
    /// * `None` - there is no information about the record in the cache
    pub(crate) async fn get_borrowed<E>(&self, id: &E::BorrowedPrimaryKey) -> Option<Option<Arc<E>>>
    where
        E: 'static + Entity + BorrowPrimaryKey + Send + Sync,
    {
        let entity_id = EntityId::from_borrowed_primary_key::<E>(id)?;
        self.get_by_entity_id(&entity_id).await
    }

    async fn find_all_in_cache<E>(&self) -> Vec<Arc<E>>
    where
        E: 'static + Entity + Send + Sync,
    {
        let cache_guard = self.cache.read().await;
        cache_guard
            .get(E::COLLECTION_NAME)
            .map(|table| {
                table
                    .values()
                    .map(|record: &dynamic_dispatch::Entity| {
                        record
                            .downcast::<E>()
                            .expect("all entries in this table are of this type")
                            .clone()
                    })
                    .collect::<Vec<_>>()
            })
            .unwrap_or_default()
    }

    async fn search_in_cache<E, SearchKey>(&self, search_key: &SearchKey) -> Vec<Arc<E>>
    where
        E: 'static + Entity + SearchableEntity<SearchKey> + Send + Sync,
        SearchKey: KeyType,
    {
        let cache_guard = self.cache.read().await;
        cache_guard
            .get(E::COLLECTION_NAME)
            .map(|table| {
                table
                    .values()
                    .filter_map(|record: &dynamic_dispatch::Entity| {
                        let entity = record
                            .downcast::<E>()
                            .expect("all entries in this table are of this type")
                            .clone();
                        entity.matches(search_key).then_some(entity)
                    })
                    .collect()
            })
            .unwrap_or_default()
    }

    pub(crate) async fn find_all<E>(&self, persisted_records: Vec<E>) -> CryptoKeystoreResult<Vec<E>>
    where
        E: 'static + Clone + Entity + Send + Sync,
    {
        let cached_records = self.find_all_in_cache().await;
        let merged_records = self
            .merge_records(
                cached_records.iter().map(Arc::as_ref).map(Cow::Borrowed),
                persisted_records.into_iter().map(Cow::Owned),
            )
            .await;
        Ok(merged_records)
    }

    pub(crate) async fn search<E, SearchKey>(
        &self,
        persisted_records: Vec<E>,
        search_key: &SearchKey,
    ) -> CryptoKeystoreResult<Vec<E>>
    where
        E: 'static + Clone + Entity + SearchableEntity<SearchKey> + Send + Sync,
        SearchKey: KeyType,
    {
        let cached_records = self.search_in_cache(search_key).await;
        let merged_records = self
            .merge_records(
                cached_records.iter().map(Arc::as_ref).map(Cow::Borrowed),
                persisted_records.into_iter().map(Cow::Owned),
            )
            .await;
        Ok(merged_records)
    }

    /// Build a single list of unique records from two potentially overlapping lists.
    /// In case of overlap, records in `records_a` are prioritized.
    /// Identity from the perspective of this function is determined by the output of
    /// [Entity::merge_key].
    ///
    /// Further, the output list of records is built with respect to the provided [EntityFindParams]
    /// and the deleted records cached in this [Self] instance.
    async fn merge_records<'a, E>(
        &self,
        records_a: impl IntoIterator<Item = Cow<'a, E>>,
        records_b: impl IntoIterator<Item = Cow<'a, E>>,
    ) -> Vec<E>
    where
        E: 'static + Clone + Entity,
    {
        let deleted_records = self.deleted.read().await;

        records_a
            .into_iter()
            .chain(records_b)
            .unique_by(|e| e.primary_key().bytes().into_owned())
            .filter_map(|record| {
                let id = EntityId::from_entity(record.as_ref())?;
                (!deleted_records.contains(&id)).then_some(record.into_owned())
            })
            .collect()
    }

    /// Persists all the operations in the database. It will effectively open a transaction
    /// internally, perform all the buffered operations and commit.
    pub(crate) async fn commit(self, db: &Database) -> Result<(), CryptoKeystoreError> {
        let cache = self.cache.read().await;
        let deleted_ids = self.deleted.read().await;

        let table_names_with_deletion = deleted_ids.iter().map(|entity_id| entity_id.collection_name());
        let table_names_with_save = cache
            .values()
            .flat_map(|table| table.keys())
            .map(|entity_id| entity_id.collection_name());
        let any_tables_are_modified = table_names_with_deletion.chain(table_names_with_save).next().is_some();

        if !any_tables_are_modified {
            log::debug!("Empty transaction was committed.");
            return Ok(());
        }

        // open a database transaction
        // Because `rusqlite::Transaction: !Send + !Sync`, it's critical that
        // we don't hold this transaction over any `.await` points.
        let mut conn = db.conn().await;
        let tx = conn.transaction()?;

        for entity in cache.values().flat_map(|table| table.values()) {
            entity.execute_save(&tx)?;
        }

        for deleted_id in deleted_ids.iter() {
            deleted_id.execute_delete(&tx)?;
        }

        // and commit everything
        tx.commit()?;

        Ok(())
    }
}