From 0a4b83e93b525cfee83ff4a106c781b573989136 Mon Sep 17 00:00:00 2001
From: Andrew Schran <aschran@mystenlabs.com>
Date: Fri, 23 Feb 2024 15:22:09 -0800
Subject: [PATCH] post-merge fixes

---
 .../src/execution_cache/passthrough_cache.rs  |  2 +-
 .../src/execution_cache/writeback_cache.rs    | 35 +++++++++++--------
 2 files changed, 21 insertions(+), 16 deletions(-)

diff --git a/crates/sui-core/src/execution_cache/passthrough_cache.rs b/crates/sui-core/src/execution_cache/passthrough_cache.rs
index 2be6c792e2bc8..1bd35b44dc965 100644
--- a/crates/sui-core/src/execution_cache/passthrough_cache.rs
+++ b/crates/sui-core/src/execution_cache/passthrough_cache.rs
@@ -279,7 +279,7 @@ impl ExecutionCacheWrite for PassthroughCache {
                 .await?;
 
             self.executed_effects_digests_notify_read
-                .notify(&tx_digest, &effects_digest);
+                .notify(&TransactionKey::Digest(tx_digest), &effects_digest);
             if let Some(key) = tx_secondary_key {
                 self.executed_effects_digests_notify_read
                     .notify(&key, &effects_digest);
diff --git a/crates/sui-core/src/execution_cache/writeback_cache.rs b/crates/sui-core/src/execution_cache/writeback_cache.rs
index 0219c599af294..7b3afe78a80ab 100644
--- a/crates/sui-core/src/execution_cache/writeback_cache.rs
+++ b/crates/sui-core/src/execution_cache/writeback_cache.rs
@@ -78,7 +78,7 @@ use sui_types::messages_checkpoint::CheckpointSequenceNumber;
 use sui_types::object::Object;
 use sui_types::storage::{MarkerValue, ObjectKey, ObjectOrTombstone, ObjectStore, PackageObject};
 use sui_types::sui_system_state::{get_sui_system_state, SuiSystemState};
-use sui_types::transaction::VerifiedTransaction;
+use sui_types::transaction::{TransactionKey, VerifiedTransaction};
 use tracing::{info, instrument};
 
 use super::ExecutionCacheAPI;
@@ -170,7 +170,7 @@ struct UncommittedData {
     transaction_events:
         DashMap<TransactionEventsDigest, (BTreeSet<TransactionDigest>, TransactionEvents)>,
 
-    executed_effects_digests: DashMap<TransactionDigest, TransactionEffectsDigest>,
+    executed_effects_digests: DashMap<TransactionKey, TransactionEffectsDigest>,
 
     // Transaction outputs that have not yet been written to the DB. Items are removed from this
     // table as they are flushed to the db.
@@ -251,7 +251,7 @@ pub struct WritebackCache {
     // - note that we removed any unfinalized packages from the cache during revert_state_update().
     packages: MokaCache<ObjectID, PackageObject>,
 
-    executed_effects_digests_notify_read: NotifyRead<TransactionDigest, TransactionEffectsDigest>,
+    executed_effects_digests_notify_read: NotifyRead<TransactionKey, TransactionEffectsDigest>,
     store: Arc<AuthorityStore>,
     metrics: Arc<ExecutionCacheMetrics>,
 }
@@ -545,6 +545,7 @@ impl WritebackCache {
         }
 
         let tx_digest = *transaction.digest();
+        let tx_key = transaction.key();
         let effects_digest = effects.digest();
 
         self.dirty
@@ -567,7 +568,7 @@ impl WritebackCache {
 
         self.dirty
             .executed_effects_digests
-            .remove(&tx_digest)
+            .remove(&tx_key)
             .expect("executed effects must exist");
 
         Ok(())
@@ -919,12 +920,12 @@ impl ExecutionCacheRead for WritebackCache {
 
     fn multi_get_executed_effects_digests(
         &self,
-        digests: &[TransactionDigest],
+        keys: &[TransactionKey],
     ) -> SuiResult<Vec<Option<TransactionEffectsDigest>>> {
         do_fallback_lookup(
-            digests,
-            |digest| {
-                if let Some(digest) = self.dirty.executed_effects_digests.get(digest) {
+            keys,
+            |key| {
+                if let Some(digest) = self.dirty.executed_effects_digests.get(key) {
                     CacheResult::Hit(Some(*digest))
                 } else {
                     CacheResult::Miss
@@ -953,14 +954,12 @@ impl ExecutionCacheRead for WritebackCache {
 
     fn notify_read_executed_effects_digests<'a>(
         &'a self,
-        digests: &'a [TransactionDigest],
+        keys: &'a [TransactionKey],
     ) -> BoxFuture<'a, SuiResult<Vec<TransactionEffectsDigest>>> {
         async move {
-            let registrations = self
-                .executed_effects_digests_notify_read
-                .register_all(digests);
+            let registrations = self.executed_effects_digests_notify_read.register_all(keys);
 
-            let executed_effects_digests = self.multi_get_executed_effects_digests(digests)?;
+            let executed_effects_digests = self.multi_get_executed_effects_digests(keys)?;
 
             let results = executed_effects_digests
                 .into_iter()
@@ -1106,6 +1105,8 @@ impl ExecutionCacheWrite for WritebackCache {
             }
 
             let tx_digest = *transaction.digest();
+            let tx_key = transaction.key();
+            let tx_secondary_key = transaction.non_digest_key();
             let effects_digest = effects.digest();
 
             self.dirty
@@ -1125,14 +1126,18 @@ impl ExecutionCacheWrite for WritebackCache {
 
             self.dirty
                 .executed_effects_digests
-                .insert(tx_digest, effects_digest);
+                .insert(tx_key, effects_digest);
 
             self.dirty
                 .pending_transaction_writes
                 .insert(tx_digest, tx_outputs);
 
             self.executed_effects_digests_notify_read
-                .notify(&tx_digest, &effects_digest);
+                .notify(&TransactionKey::Digest(tx_digest), &effects_digest);
+            if let Some(key) = tx_secondary_key {
+                self.executed_effects_digests_notify_read
+                    .notify(&key, &effects_digest);
+            }
 
             self.metrics
                 .pending_notify_read