iota_core/execution_cache/

object_locks.rs

// Copyright (c) Mysten Labs, Inc.
// Modifications Copyright (c) 2024 IOTA Stiftung
// SPDX-License-Identifier: Apache-2.0

use dashmap::{DashMap, mapref::entry::Entry as DashMapEntry};
use iota_common::debug_fatal;
use iota_sdk_types::{ObjectId, ObjectReference};
use iota_types::{
    error::{IotaError, IotaResult, UserInputError},
    object::Object,
    storage::ObjectStore,
    transaction::VerifiedSignedTransaction,
};
use tracing::{debug, info, instrument, trace};

use super::writeback_cache::WritebackCache;
use crate::authority::authority_per_epoch_store::{AuthorityPerEpochStore, LockDetails};

type RefCount = usize;

pub(super) struct ObjectLocks {
    // When acquire transaction locks, lock entries are briefly inserted into this map. The map
    // exists to provide atomic test-and-set operations on the locks. After all locks have been
    // inserted into the map, they are written to the db, and then all locks are removed from
    // the map.
    //
    // After a transaction has been executed, newly created objects are available to be locked.
    // But, because of crash recovery, we cannot rule out that a lock may already exist in the db
    // for those objects. Therefore we do a db read for each object we are locking.
    //
    // TODO: find a strategy to allow us to avoid db reads for each object.
    locked_transactions: DashMap<ObjectReference, (RefCount, LockDetails)>,
}

impl ObjectLocks {
    pub fn new() -> Self {
        Self {
            locked_transactions: DashMap::new(),
        }
    }

    pub(crate) fn get_transaction_lock(
        &self,
        obj_ref: &ObjectReference,
        epoch_store: &AuthorityPerEpochStore,
    ) -> IotaResult<Option<LockDetails>> {
        // We don't consult the in-memory state here. We are only interested in state
        // that has been committed to the db. This is because in memory state is
        // reverted if the transaction is not successfully locked.
        epoch_store.tables()?.get_locked_transaction(obj_ref)
    }

    /// Attempts to atomically test-and-set a transaction lock on an object.
    /// If the lock is already set to a conflicting transaction, an error is
    /// returned and the cached state is left untouched. If the lock is not
    /// set, or is already set to the same transaction, the lock is set and the
    /// caller must release it with `clear_cached_locks`.
    fn try_set_transaction_lock(
        &self,
        obj_ref: &ObjectReference,
        new_lock: LockDetails,
        epoch_store: &AuthorityPerEpochStore,
    ) -> IotaResult {
        // entry holds a lock on the dashmap shard, so this function operates atomically
        let entry = self.locked_transactions.entry(*obj_ref);

        // TODO: currently, the common case for this code is that we will miss the cache
        // and read from the db. It is difficult to implement negative caching, since we
        // may have restarted, in which case there could be locks in the db that we do
        // not have in the cache. We may want to explore strategies for proving there
        // cannot be a lock in the db that we do not know about. Two possibilities are:
        //
        // 1. Read all locks into memory at startup (and keep them there). The lifetime
        //    of locks is relatively short in the common case, so this might be
        //    feasible.
        // 2. Find some strategy to distinguish between the cases where we are
        //    re-executing old transactions after restarting vs executing transactions
        //    that we have never seen before. The output objects of novel transactions
        //    cannot previously have been locked on this validator.
        //
        // Solving this is not terribly important as it is not in the execution path,
        // and hence only improves the latency of transaction signing, not
        // transaction execution
        let conflict = |prev_lock: LockDetails| {
            debug!(
                "lock conflict detected for {:?}: {:?} != {:?}",
                obj_ref, prev_lock, new_lock
            );
            IotaError::ObjectLockConflict {
                obj_ref: *obj_ref,
                pending_transaction: prev_lock,
            }
        };

        // A rejected caller never calls `clear_cached_locks` for this object, so
        // the entry must only be created or counted once the lock is known to be
        // compatible.
        match entry {
            DashMapEntry::Vacant(vacant) => {
                let tables = epoch_store.tables()?;
                match tables.get_locked_transaction(obj_ref)? {
                    Some(prev_lock) if prev_lock != new_lock => return Err(conflict(prev_lock)),
                    Some(prev_lock) => trace!("read lock from db: {:?}", prev_lock),
                    None => trace!("set lock: {:?}", new_lock),
                }
                vacant.insert((1, new_lock));
            }
            DashMapEntry::Occupied(mut occupied) => {
                let prev_lock = occupied.get().1;
                if prev_lock != new_lock {
                    return Err(conflict(prev_lock));
                }
                occupied.get_mut().0 += 1;
            }
        }

        Ok(())
    }

    pub(crate) fn clear(&self) {
        info!("clearing old transaction locks");
        self.locked_transactions.clear();
    }

    fn verify_live_object(obj_ref: &ObjectReference, live_object: &Object) -> IotaResult {
        debug_assert_eq!(obj_ref.object_id, live_object.id());
        if obj_ref.version != live_object.version() {
            debug!(
                "object version unavailable for consumption: {:?} (current: {})",
                obj_ref,
                live_object.version()
            );
            return Err(IotaError::UserInput {
                error: UserInputError::ObjectVersionUnavailableForConsumption {
                    provided_obj_ref: *obj_ref,
                    current_version: live_object.version(),
                },
            });
        }

        let live_digest = live_object.digest();
        if obj_ref.digest != live_digest {
            return Err(IotaError::UserInput {
                error: UserInputError::InvalidObjectDigest {
                    object_id: obj_ref.object_id,
                    expected_digest: live_digest,
                },
            });
        }

        Ok(())
    }

    fn clear_cached_locks(&self, locks: &[(ObjectReference, LockDetails)]) {
        for (obj_ref, lock) in locks {
            let entry = self.locked_transactions.entry(*obj_ref);
            let mut occupied = match entry {
                DashMapEntry::Vacant(_) => {
                    debug_fatal!("lock must exist for object: {:?}", obj_ref);
                    continue;
                }
                DashMapEntry::Occupied(occupied) => occupied,
            };

            if occupied.get().1 == *lock {
                occupied.get_mut().0 -= 1;
                if occupied.get().0 == 0 {
                    trace!("clearing lock: {:?}", lock);
                    occupied.remove();
                }
            } else {
                // this is impossible because the only case in which we overwrite a
                // lock is when the lock is from a previous epoch. but we are holding
                // execution_lock, so the epoch cannot have changed.
                panic!("lock was changed since we set it");
            }
        }
    }

    fn multi_get_objects_must_exist(
        cache: &WritebackCache,
        object_ids: &[ObjectId],
    ) -> IotaResult<Vec<Object>> {
        let objects = cache.try_multi_get_objects(object_ids)?;
        let mut result = Vec::with_capacity(objects.len());
        for (i, object) in objects.into_iter().enumerate() {
            if let Some(object) = object {
                result.push(object);
            } else {
                return Err(IotaError::UserInput {
                    error: UserInputError::ObjectNotFound {
                        object_id: object_ids[i],
                        version: None,
                    },
                });
            }
        }
        Ok(result)
    }

    /// Validates that all owned input objects exist and their versions/digests
    /// match the live objects. Does not acquire any locks.
    pub(crate) fn validate_owned_object_versions(
        cache: &WritebackCache,
        owned_input_objects: &[ObjectReference],
    ) -> IotaResult {
        if owned_input_objects.is_empty() {
            return Ok(());
        }
        let object_ids: Vec<_> = owned_input_objects.iter().map(|o| o.object_id).collect();
        let live_objects = Self::multi_get_objects_must_exist(cache, &object_ids)?;
        for (obj_ref, live_object) in owned_input_objects.iter().zip(live_objects.iter()) {
            Self::verify_live_object(obj_ref, live_object)?;
        }
        Ok(())
    }

    #[instrument(level = "debug", skip_all)]
    pub(crate) fn acquire_transaction_locks(
        &self,
        cache: &WritebackCache,
        epoch_store: &AuthorityPerEpochStore,
        owned_input_objects: &[ObjectReference],
        transaction: VerifiedSignedTransaction,
    ) -> IotaResult {
        let tx_digest = *transaction.digest();

        let object_ids = owned_input_objects
            .iter()
            .map(|o| o.object_id)
            .collect::<Vec<_>>();
        let live_objects = Self::multi_get_objects_must_exist(cache, &object_ids)?;

        // Only live objects can be locked
        for (obj_ref, live_object) in owned_input_objects.iter().zip(live_objects.iter()) {
            Self::verify_live_object(obj_ref, live_object)?;
        }

        let mut locks_to_write: Vec<(_, LockDetails)> =
            Vec::with_capacity(owned_input_objects.len());

        // Sort the objects before locking. This is not required by the protocol (since
        // it's okay to reject any equivocating tx). However, this does prevent
        // a confusing error on the client. Consider the case:
        //   TX1: [o1, o2];
        //   TX2: [o2, o1];
        // If two threads race to acquire these locks, they might both acquire the first
        // object, then error when trying to acquire the second. The error
        // returned to the client would say that there is a conflicting tx on
        // that object, but in fact neither object was locked and the tx was never
        // signed. If one client then retries, they will succeed (counterintuitively).
        let owned_input_objects = {
            let mut o = owned_input_objects.to_vec();
            o.sort_by_key(|o| o.object_id);
            o
        };

        // Note that this function does not have to operate atomically. If there are two
        // racing threads, then they are either trying to lock the same
        // transaction (in which case both will succeed), or they are trying to
        // lock the same object in two different transactions, in which case the
        // sender has equivocated, and we are under no obligation to help them form a
        // cert.
        for obj_ref in owned_input_objects.iter() {
            match self.try_set_transaction_lock(obj_ref, tx_digest, epoch_store) {
                Ok(()) => locks_to_write.push((*obj_ref, tx_digest)),
                Err(e) => {
                    // revert all pending writes and return error
                    // Note that reverting is not required for liveness, since a well formed and
                    // un-equivocating txn cannot fail to acquire locks.
                    // However, reverting is easy enough to do in this implementation that we do it
                    // anyway.
                    self.clear_cached_locks(&locks_to_write);
                    return Err(e);
                }
            }
        }

        // commit all writes to DB
        epoch_store
            .tables()?
            .write_transaction_locks(transaction, locks_to_write.iter().cloned())?;

        // remove pending locks from unbounded storage
        self.clear_cached_locks(&locks_to_write);

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use iota_sdk_types::TransactionDigest;

    use super::ObjectLocks;
    use crate::execution_cache::{
        ExecutionCacheWrite, writeback_cache::writeback_cache_tests::Scenario,
    };

    #[tokio::test]
    async fn test_lock_conflict_does_not_count_against_pending_lock() {
        telemetry_subscribers::init_for_testing();
        Scenario::iterate(|mut s| async move {
            s.with_created(&[1]);
            s.do_tx().await;
            let obj = s.obj_ref(1);

            let holder = TransactionDigest::random();
            let contender = TransactionDigest::random();

            // A signer of `holder` is mid-flight: its entry is set but not yet
            // written to the db and released.
            let locks = ObjectLocks::new();
            locks
                .try_set_transaction_lock(&obj, holder, &s.epoch_store)
                .expect("first lock on a fresh object");
            assert_eq!(*locks.locked_transactions.get(&obj).unwrap(), (1, holder));

            locks
                .try_set_transaction_lock(&obj, contender, &s.epoch_store)
                .unwrap_err();
            assert_eq!(
                *locks.locked_transactions.get(&obj).unwrap(),
                (1, holder),
                "a rejected contender must not be counted as a holder"
            );

            // Another signer of the same transaction still shares the entry.
            locks
                .try_set_transaction_lock(&obj, holder, &s.epoch_store)
                .expect("same transaction may lock again");
            assert_eq!(*locks.locked_transactions.get(&obj).unwrap(), (2, holder));
        })
        .await;
    }

    #[tokio::test]
    async fn test_transaction_locks_are_exclusive() {
        telemetry_subscribers::init_for_testing();
        Scenario::iterate(|mut s| async move {
            s.with_created(&[1, 2, 3]);
            s.do_tx().await;

            s.with_mutated(&[1, 2, 3]);
            s.do_tx().await;

            let new1 = s.obj_ref(1);
            let new2 = s.obj_ref(2);
            let new3 = s.obj_ref(3);

            s.with_mutated(&[1, 2, 3]); // begin forming a tx but never execute it
            let outputs = s.take_outputs();

            let tx1 = s.make_signed_transaction(&outputs.transaction);

            s.cache
                .try_acquire_transaction_locks(&s.epoch_store, &[new1, new2], tx1)
                .expect("locks should be available");
            // Entries are staged only until the locks are written, and a
            // rejected conflict must leave the map exactly as it found it.
            let cache = s.cache.clone();
            let assert_no_staged_locks = || {
                assert!(
                    cache
                        .object_locks_for_testing()
                        .locked_transactions
                        .is_empty()
                )
            };
            assert_no_staged_locks();

            // this tx doesn't use the actual objects in question, but we just need
            // something to insert into the table.
            s.with_created(&[4, 5]);
            let tx2 = s.take_outputs().transaction.clone();
            let tx2 = s.make_signed_transaction(&tx2);

            // both locks are held by tx1, so this should fail
            s.cache
                .try_acquire_transaction_locks(&s.epoch_store, &[new1, new2], tx2.clone())
                .unwrap_err();
            assert_no_staged_locks();

            // new3 is lockable, but new2 is not, so this should fail
            s.cache
                .try_acquire_transaction_locks(&s.epoch_store, &[new3, new2], tx2.clone())
                .unwrap_err();
            assert_no_staged_locks();

            // new3 is unlocked
            s.cache
                .try_acquire_transaction_locks(&s.epoch_store, &[new3], tx2)
                .expect("new3 should be unlocked");
            assert_no_staged_locks();
        })
        .await;
    }

    #[tokio::test]
    async fn test_transaction_locks_are_durable() {
        telemetry_subscribers::init_for_testing();
        Scenario::iterate(|mut s| async move {
            s.with_created(&[1, 2]);
            s.do_tx().await;

            let old2 = s.obj_ref(2);

            s.with_mutated(&[1, 2]);
            s.do_tx().await;

            let new1 = s.obj_ref(1);
            let new2 = s.obj_ref(2);

            s.with_mutated(&[1, 2]); // begin forming a tx but never execute it
            let outputs = s.take_outputs();

            let tx = s.make_signed_transaction(&outputs.transaction);

            // fails because we are referring to an old object
            s.cache
                .try_acquire_transaction_locks(&s.epoch_store, &[new1, old2], tx.clone())
                .unwrap_err();

            // succeeds because the above call releases the lock on new1 after failing
            // to get the lock on old2
            s.cache
                .try_acquire_transaction_locks(&s.epoch_store, &[new1, new2], tx)
                .expect("new1 should be unlocked after revert");
        })
        .await;
    }

    #[tokio::test]
    async fn test_acquire_transaction_locks_revert() {
        telemetry_subscribers::init_for_testing();
        Scenario::iterate(|mut s| async move {
            s.with_created(&[1, 2]);
            s.do_tx().await;

            let old2 = s.obj_ref(2);

            s.with_mutated(&[1, 2]);
            s.do_tx().await;

            let new1 = s.obj_ref(1);
            let new2 = s.obj_ref(2);

            s.with_mutated(&[1, 2]); // begin forming a tx but never execute it
            let outputs = s.take_outputs();

            let tx = s.make_signed_transaction(&outputs.transaction);

            // fails because we are referring to an old object
            s.cache
                .try_acquire_transaction_locks(&s.epoch_store, &[new1, old2], tx)
                .unwrap_err();

            // this tx doesn't use the actual objects in question, but we just need
            // something to insert into the table.
            s.with_created(&[4, 5]);
            let tx2 = s.take_outputs().transaction.clone();
            let tx2 = s.make_signed_transaction(&tx2);

            // succeeds because the above call releases the lock on new1 after failing
            // to get the lock on old2
            s.cache
                .try_acquire_transaction_locks(&s.epoch_store, &[new1, new2], tx2)
                .expect("new1 should be unlocked after revert");
        })
        .await;
    }

    #[tokio::test]
    async fn test_acquire_transaction_locks_is_sync() {
        telemetry_subscribers::init_for_testing();
        Scenario::iterate(|mut s| async move {
            s.with_created(&[1, 2]);
            s.do_tx().await;

            let objects: Vec<_> = vec![s.object(1), s.object(2)]
                .into_iter()
                .map(|o| o.object_ref())
                .collect();

            s.with_mutated(&[1, 2]);
            let outputs = s.take_outputs();

            let tx2 = s.make_signed_transaction(&outputs.transaction);
            // assert that acquire_transaction_locks is sync in non-simtest, which causes
            // the fail_point_async! macros above to be elided
            s.cache
                .try_acquire_transaction_locks(&s.epoch_store, &objects, tx2)
                .unwrap();
        })
        .await;
    }
}