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,
    connection::{Database, KeystoreDatabaseConnection},
    entities::{MlsPendingMessage, MlsPendingMessagePrimaryKey, PersistedMlsGroupExt},
    traits::{BorrowPrimaryKey, Entity, EntityBase as _, EntityDatabaseMutation, EntityDeleteBorrowed, KeyType},
    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(crate) 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<'a, E>(&self, mut entity: E) -> CryptoKeystoreResult<E::AutoGeneratedFields>
    where
        E: Entity<ConnectionType = KeystoreDatabaseConnection> + EntityDatabaseMutation<'a> + Send + Sync,
    {
        let auto_generated_fields = entity.pre_save().await?;

        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.to_transaction_entity());
        }
        {
            // 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<'a, E>(&self, entity_id: EntityId) -> CryptoKeystoreResult<()>
    where
        E: Entity + EntityDatabaseMutation<'a>,
    {
        // 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<'a, E>(&self, id: &E::PrimaryKey) -> CryptoKeystoreResult<()>
    where
        E: Entity + EntityDatabaseMutation<'a>,
    {
        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<'a, E>(&self, id: &E::BorrowedPrimaryKey) -> CryptoKeystoreResult<()>
    where
        E: EntityDeleteBorrowed<'a> + 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
    }

    pub(crate) async fn child_groups<E>(
        &self,
        entity: E,
        persisted_records: impl IntoIterator<Item = E>,
    ) -> CryptoKeystoreResult<Vec<E>>
    where
        E: Clone + Entity + BorrowPrimaryKey + PersistedMlsGroupExt + Send + Sync,
        for<'pk> &'pk <E as BorrowPrimaryKey>::BorrowedPrimaryKey: KeyType,
    {
        // First get all raw groups from the cache, then filter by their parent id
        let cached_records = self.find_all_in_cache::<E>().await;
        let cached_records = cached_records
            .iter()
            .filter(|maybe_child| {
                maybe_child
                    .parent_id()
                    .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_primary_key::<MlsPendingMessage>(&MlsPendingMessagePrimaryKey::from_conversation_id(
                conversation_id,
            ))
            .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: 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: 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: 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: 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: 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()
    }

    pub(crate) async fn find_all<E>(&self, persisted_records: Vec<E>) -> CryptoKeystoreResult<Vec<E>>
    where
        E: 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)
    }

    /// 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: 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 conn = db.conn().await?;
        let mut conn = conn.conn().await;

        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 mut tables = table_names_with_deletion
            .chain(table_names_with_save)
            .collect::<Vec<_>>();

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

        tables.sort_unstable();
        tables.dedup();

        // open a database transaction
        #[cfg(target_family = "wasm")]
        let tx = conn.new_transaction(&tables).await?;
        #[cfg(not(target_family = "wasm"))]
        let tx = conn.transaction()?.into();

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

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

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

        Ok(())
    }
}