Add input object notify in the writeback cache (#21995)

lxfind · web-flow · commit ed70303929d4 · 2025-05-13T00:59:05.000Z
## Description 

This PR adds an object notify in the writeback cache. It will allow us
to replace transaction manager.
A transaction to be scheduled will call the notify wait to wait for
input object to become available.
The notify wait function is very similar to
`multi_input_objects_available` in execution cache.

## Test plan 

Added tests

---

## Release notes

Check each box that your changes affect. If none of the boxes relate to
your changes, release notes aren't required.

For each box you select, include information after the relevant heading
that describes the impact of your changes that a user might notice and
any actions they must take to implement updates.

- [ ] Protocol: 
- [ ] Nodes (Validators and Full nodes): 
- [ ] gRPC:
- [ ] JSON-RPC: 
- [ ] GraphQL: 
- [ ] CLI: 
- [ ] Rust SDK:
diff --git a/crates/sui-core/src/execution_cache.rs b/crates/sui-core/src/execution_cache.rs
@@ -261,7 +261,7 @@ pub trait ObjectCacheRead: Send + Sync {
         ) {
             assert!(
                 input_key.version().is_none() || input_key.version().unwrap().is_valid(),
-                "Shared objects in cancelled transaction should always be available immediately, 
+                "Shared objects in cancelled transaction should always be available immediately,
                  but it appears that transaction manager is waiting for {:?} to become available",
                 input_key
             );
@@ -397,7 +397,7 @@ pub trait ObjectCacheRead: Send + Sync {
         version: SequenceNumber,
         epoch_id: EpochId,
     ) -> bool {
-        let full_id = FullObjectID::Fastpath(object_id); // function explicilty assumes "fastpath"
+        let full_id = FullObjectID::Fastpath(object_id); // function explicitly assumes "fastpath"
         matches!(
             self.get_latest_marker(full_id, epoch_id),
             Some((marker_version, MarkerValue::FastpathStreamEnded)) if marker_version >= version
@@ -406,6 +406,20 @@ pub trait ObjectCacheRead: Send + Sync {
 
     /// Return the watermark for the highest checkpoint for which we've pruned objects.
     fn get_highest_pruned_checkpoint(&self) -> Option<CheckpointSequenceNumber>;
+
+    /// Given a list of input and receiving objects for a transaction,
+    /// wait until all of them become available, so that the transaction
+    /// can start execution.
+    /// `input_and_receiving_keys` contains both input objects and receiving
+    /// input objects, including canceled objects.
+    /// TODO: Eventually this can return the objects read results,
+    /// so that execution does not need to load them again.
+    fn notify_read_input_objects<'a>(
+        &'a self,
+        input_and_receiving_keys: &'a [InputKey],
+        receiving_keys: &'a HashSet<InputKey>,
+        epoch: &'a EpochId,
+    ) -> BoxFuture<'a, Vec<()>>;
 }
 
 pub trait TransactionCacheRead: Send + Sync {
diff --git a/crates/sui-core/src/execution_cache/unit_tests/notify_read_input_objects_tests.rs b/crates/sui-core/src/execution_cache/unit_tests/notify_read_input_objects_tests.rs
@@ -0,0 +1,324 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use crate::authority::authority_store_tables::AuthorityPerpetualTables;
+
+use super::*;
+use futures::FutureExt;
+use std::path::Path;
+use std::time::Duration;
+use sui_framework::BuiltInFramework;
+use sui_move_build::BuildConfig;
+use sui_swarm_config::network_config_builder::ConfigBuilder;
+use sui_types::base_types::{ObjectID, SequenceNumber, SuiAddress};
+use sui_types::object::{Object, Owner};
+use sui_types::storage::InputKey;
+use sui_types::SUI_FRAMEWORK_PACKAGE_ID;
+use tempfile::tempdir;
+use tokio::time::timeout;
+
+async fn create_writeback_cache() -> Arc<WritebackCache> {
+    let path = tempdir().unwrap();
+    let tables = Arc::new(AuthorityPerpetualTables::open(path.path(), None));
+    let config = ConfigBuilder::new_with_temp_dir().build();
+    let store = AuthorityStore::open_with_committee_for_testing(
+        tables,
+        config.committee_with_network().committee(),
+        &config.genesis,
+    )
+    .await
+    .unwrap();
+    Arc::new(WritebackCache::new_for_tests(store))
+}
+
+#[tokio::test]
+async fn test_immediate_return_canceled_shared() {
+    let cache = create_writeback_cache().await;
+
+    let canceled_key = InputKey::VersionedObject {
+        id: FullObjectID::new(ObjectID::random(), Some(SequenceNumber::from(1))),
+        version: SequenceNumber::CANCELLED_READ,
+    };
+    let receiving_keys = HashSet::new();
+    let epoch = &0;
+
+    // Should return immediately since canceled shared objects are always available
+    let result = cache
+        .notify_read_input_objects(&[canceled_key], &receiving_keys, epoch)
+        .now_or_never()
+        .unwrap();
+    assert_eq!(result.len(), 1);
+
+    let congested_key = InputKey::VersionedObject {
+        id: FullObjectID::new(ObjectID::random(), Some(SequenceNumber::from(1))),
+        version: SequenceNumber::CONGESTED,
+    };
+
+    let result = cache
+        .notify_read_input_objects(&[congested_key], &receiving_keys, epoch)
+        .now_or_never()
+        .unwrap();
+    assert_eq!(result.len(), 1);
+
+    let randomness_unavailable_key = InputKey::VersionedObject {
+        id: FullObjectID::new(ObjectID::random(), Some(SequenceNumber::from(1))),
+        version: SequenceNumber::RANDOMNESS_UNAVAILABLE,
+    };
+
+    let result = cache
+        .notify_read_input_objects(&[randomness_unavailable_key], &receiving_keys, epoch)
+        .now_or_never()
+        .unwrap();
+    assert_eq!(result.len(), 1);
+}
+
+#[tokio::test]
+async fn test_immediate_return_cached_object() {
+    let cache = create_writeback_cache().await;
+
+    let object_id = ObjectID::random();
+    let version = SequenceNumber::from(1);
+    let object = Object::with_id_owner_version_for_testing(object_id, version, Owner::Immutable);
+
+    cache.write_object_entry(&object_id, version, ObjectEntry::Object(object));
+
+    let input_keys = vec![InputKey::VersionedObject {
+        id: FullObjectID::new(object_id, None),
+        version,
+    }];
+    let receiving_keys = HashSet::new();
+    let epoch = &0;
+
+    // Should return immediately since object is in cache
+    let result = cache
+        .notify_read_input_objects(&input_keys, &receiving_keys, epoch)
+        .now_or_never()
+        .unwrap();
+
+    assert_eq!(result.len(), 1);
+}
+
+#[tokio::test]
+async fn test_immediate_return_cached_package() {
+    let cache = create_writeback_cache().await;
+
+    let input_keys = vec![InputKey::Package {
+        id: SUI_FRAMEWORK_PACKAGE_ID,
+    }];
+    let receiving_keys = HashSet::new();
+    let epoch = &0;
+
+    // Should return immediately since system package is available by default.
+    let result = cache
+        .notify_read_input_objects(&input_keys, &receiving_keys, epoch)
+        .now_or_never()
+        .unwrap();
+
+    assert_eq!(result.len(), 1);
+}
+
+#[tokio::test]
+async fn test_immediate_return_consensus_stream_ended() {
+    let cache = create_writeback_cache().await;
+
+    let object_id = ObjectID::random();
+    let version = SequenceNumber::from(1);
+    let epoch = 0;
+
+    // Write consensus stream ended marker
+    cache.write_marker_value(
+        epoch,
+        FullObjectKey::new(FullObjectID::new(object_id, Some(version)), version),
+        MarkerValue::ConsensusStreamEnded(TransactionDigest::random()),
+    );
+
+    let input_keys = vec![InputKey::VersionedObject {
+        id: FullObjectID::new(object_id, Some(version)),
+        version,
+    }];
+    let receiving_keys = HashSet::new();
+
+    // Should return immediately since object is marked as consensus stream ended
+    let result = cache
+        .notify_read_input_objects(&input_keys, &receiving_keys, &epoch)
+        .now_or_never()
+        .unwrap();
+
+    assert_eq!(result.len(), 1);
+}
+
+#[tokio::test]
+async fn test_wait_for_object() {
+    let cache = create_writeback_cache().await;
+
+    let object_id = ObjectID::random();
+    let version = SequenceNumber::from(1);
+
+    let input_keys = vec![InputKey::VersionedObject {
+        id: FullObjectID::new(object_id, Some(version)),
+        version,
+    }];
+    let receiving_keys = HashSet::new();
+    let epoch = &0;
+
+    let result = timeout(
+        Duration::from_secs(3),
+        cache.notify_read_input_objects(&input_keys, &receiving_keys, epoch),
+    )
+    .await;
+    assert!(result.is_err());
+
+    // Write an older version of the object.
+    tokio::spawn({
+        let cache = cache.clone();
+        async move {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            let object = Object::with_id_owner_version_for_testing(
+                object_id,
+                SequenceNumber::from(0),
+                Owner::Shared {
+                    initial_shared_version: version,
+                },
+            );
+            cache.write_object_entry(&object_id, version, ObjectEntry::Object(object));
+        }
+    });
+    let result = timeout(
+        Duration::from_secs(3),
+        cache.notify_read_input_objects(&input_keys, &receiving_keys, epoch),
+    )
+    .await;
+    assert!(result.is_err());
+
+    // Write the correct version of the object.
+    tokio::spawn({
+        let cache = cache.clone();
+        async move {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            let object = Object::with_id_owner_version_for_testing(
+                object_id,
+                version,
+                Owner::Shared {
+                    initial_shared_version: version,
+                },
+            );
+            cache.write_object_entry(&object_id, version, ObjectEntry::Object(object));
+        }
+    });
+    let result = timeout(
+        Duration::from_secs(3),
+        cache.notify_read_input_objects(&input_keys, &receiving_keys, epoch),
+    )
+    .await
+    .unwrap();
+    assert_eq!(result.len(), 1);
+}
+
+#[tokio::test]
+async fn test_wait_for_package() {
+    let cache = create_writeback_cache().await;
+
+    let path = Path::new(env!("CARGO_MANIFEST_DIR")).join("../../examples/move/basics");
+    let compiled_modules = BuildConfig::new_for_testing()
+        .build(&path)
+        .unwrap()
+        .into_modules();
+    let package = Object::new_package_for_testing(
+        &compiled_modules,
+        TransactionDigest::genesis_marker(),
+        BuiltInFramework::genesis_move_packages(),
+    )
+    .unwrap();
+    let package_id = package.id();
+    let version = package.version();
+
+    let input_keys = vec![InputKey::Package { id: package_id }];
+    let receiving_keys = HashSet::new();
+    let epoch = &0;
+
+    // Start notification future
+    let notification = cache.notify_read_input_objects(&input_keys, &receiving_keys, epoch);
+
+    // Write package after small delay
+    tokio::spawn({
+        let cache = cache.clone();
+        async move {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            cache.write_object_entry(&package_id, version, ObjectEntry::Object(package));
+        }
+    });
+
+    // Should complete once package is written
+    let result = timeout(Duration::from_secs(1), notification).await.unwrap();
+
+    assert_eq!(result.len(), 1);
+}
+
+#[tokio::test]
+async fn test_wait_for_consensus_stream_end() {
+    let cache = create_writeback_cache().await;
+
+    let object_id = ObjectID::random();
+    let version = SequenceNumber::from(1);
+    let epoch = &0;
+
+    let input_keys = vec![InputKey::VersionedObject {
+        id: FullObjectID::new(object_id, Some(version)),
+        version,
+    }];
+    let receiving_keys = HashSet::new();
+
+    // Start notification future
+    let notification = cache.notify_read_input_objects(&input_keys, &receiving_keys, epoch);
+
+    // Write consensus stream ended marker after small delay
+    tokio::spawn({
+        let cache = cache.clone();
+        async move {
+            tokio::time::sleep(Duration::from_millis(100)).await;
+            cache.write_marker_value(
+                *epoch,
+                FullObjectKey::new(FullObjectID::new(object_id, Some(version)), version),
+                MarkerValue::ConsensusStreamEnded(TransactionDigest::random()),
+            );
+        }
+    });
+
+    // Should complete once marker is written
+    let result = timeout(Duration::from_secs(1), notification).await.unwrap();
+
+    assert_eq!(result.len(), 1);
+}
+
+#[tokio::test]
+async fn test_receiving_object_higher_version() {
+    let cache = create_writeback_cache().await;
+
+    let object_id = ObjectID::random();
+    let requested_version = SequenceNumber::from(1);
+    let higher_version = SequenceNumber::from(2);
+    let object = Object::with_id_owner_version_for_testing(
+        object_id,
+        higher_version,
+        Owner::AddressOwner(SuiAddress::default()),
+    );
+
+    // Write higher version to cache
+    cache.write_object_entry(&object_id, higher_version, ObjectEntry::Object(object));
+
+    let input_keys = vec![InputKey::VersionedObject {
+        id: FullObjectID::new(object_id, None),
+        version: requested_version,
+    }];
+    let mut receiving_keys = HashSet::new();
+    receiving_keys.insert(input_keys[0]);
+    let epoch = &0;
+
+    // Should return immediately since a higher version exists for receiving object
+    let result = cache
+        .notify_read_input_objects(&input_keys, &receiving_keys, epoch)
+        .now_or_never()
+        .unwrap();
+
+    assert_eq!(result.len(), 1);
+}
diff --git a/crates/sui-core/src/execution_cache/writeback_cache.rs b/crates/sui-core/src/execution_cache/writeback_cache.rs
