Implement object balance withdraw

Author: lxfind

crates/sui-core/src/authority.rs

@@ -12,6 +12,7 @@ use crate::execution_cache::ExecutionCacheTraitPointers;
 use crate::execution_cache::TransactionCacheRead;
 use crate::execution_scheduler::ExecutionScheduler;
 use crate::execution_scheduler::SchedulingSource;
+use crate::execution_scheduler::balance_withdraw_scheduler::BalanceSettlement;
 use crate::jsonrpc_index::CoinIndexKey2;
 use crate::rpc_index::RpcIndexStore;
 use crate::traffic_controller::TrafficController;
@@ -802,8 +803,9 @@ pub struct ExecutionEnv {
     pub expected_effects_digest: Option<TransactionEffectsDigest>,
     /// The source of the scheduling of the transaction.
     pub scheduling_source: SchedulingSource,
-    /// Status of the balance withdraw scheduling of the transaction.
-    pub withdraw_status: BalanceWithdrawStatus,
+    /// Status of the balance withdraw scheduling of the transaction,
+    /// including both address and object balance withdraws.
+    pub balance_withdraw_status: BalanceWithdrawStatus,
     /// Transactions that must finish before this transaction can be executed.
     /// Used to schedule barrier transactions after non-exclusive writes.
     pub barrier_dependencies: Vec<TransactionDigest>,
@@ -815,7 +817,7 @@ impl Default for ExecutionEnv {
             assigned_versions: Default::default(),
             expected_effects_digest: None,
             scheduling_source: SchedulingSource::NonFastPath,
-            withdraw_status: BalanceWithdrawStatus::NoWithdraw,
+            balance_withdraw_status: BalanceWithdrawStatus::Unknown,
             barrier_dependencies: Default::default(),
         }
     }
@@ -844,13 +846,8 @@ impl ExecutionEnv {
         self
     }
 
-    pub fn with_sufficient_balance(mut self) -> Self {
-        self.withdraw_status = BalanceWithdrawStatus::SufficientBalance;
-        self
-    }
-
     pub fn with_insufficient_balance(mut self) -> Self {
-        self.withdraw_status = BalanceWithdrawStatus::InsufficientBalance;
+        self.balance_withdraw_status = BalanceWithdrawStatus::InsufficientBalance;
         self
     }
 
@@ -1763,7 +1760,7 @@ impl AuthorityState {
         &self,
         tx_lock: &CertLockGuard,
         certificate: &VerifiedExecutableTransaction,
-        assigned_shared_object_versions: AssignedVersions,
+        assigned_shared_object_versions: &AssignedVersions,
         epoch_store: &Arc<AuthorityPerEpochStore>,
     ) -> SuiResult<InputObjects> {
         let _scope = monitored_scope("Execution::load_input_objects");
@@ -1776,7 +1773,7 @@ impl AuthorityState {
             &certificate.key(),
             tx_lock,
             input_objects,
-            &assigned_shared_object_versions,
+            assigned_shared_object_versions,
             epoch_store.epoch(),
         )
     }
@@ -1871,7 +1868,7 @@ impl AuthorityState {
         let input_objects = match self.read_objects_for_execution(
             tx_guard.as_lock_guard(),
             certificate,
-            execution_env.assigned_versions,
+            &execution_env.assigned_versions,
             epoch_store,
         ) {
             Ok(objects) => objects,
@@ -1907,7 +1904,7 @@ impl AuthorityState {
             certificate,
             input_objects,
             expected_effects_digest,
-            execution_env.withdraw_status,
+            execution_env,
             epoch_store,
         )
     }
@@ -2020,7 +2017,7 @@ impl AuthorityState {
         certificate: &VerifiedExecutableTransaction,
         input_objects: InputObjects,
         expected_effects_digest: Option<TransactionEffectsDigest>,
-        withdraw_status: BalanceWithdrawStatus,
+        execution_env: ExecutionEnv,
         epoch_store: &Arc<AuthorityPerEpochStore>,
     ) -> ExecutionOutput<(
         TransactionOutputs,
@@ -2062,7 +2059,7 @@ impl AuthorityState {
             &tx_digest,
             &input_objects,
             self.config.certificate_deny_config.certificate_deny_set(),
-            &withdraw_status,
+            &execution_env.balance_withdraw_status,
         );
         let execution_params = match early_execution_error {
             Some(error) => ExecutionOrEarlyError::Err(error),
@@ -2097,6 +2094,13 @@ impl AuthorityState {
                 &mut None,
             );
 
+        if !self
+            .execution_scheduler
+            .should_commit_object_balance_withdraws(certificate, &effects, &execution_env)
+        {
+            return ExecutionOutput::RetryLater;
+        }
+
         if let Some(expected_effects_digest) = expected_effects_digest
             && effects.digest() != expected_effects_digest
         {
@@ -2232,7 +2236,7 @@ impl AuthorityState {
                 certificate,
                 input_objects,
                 None,
-                BalanceWithdrawStatus::NoWithdraw,
+                ExecutionEnv::default(),
                 epoch_store,
             )
             .unwrap();
@@ -2376,7 +2380,7 @@ impl AuthorityState {
             &checked_input_objects,
             self.config.certificate_deny_config.certificate_deny_set(),
             // TODO(address-balances): Mimic withdraw scheduling and pass the result.
-            &BalanceWithdrawStatus::NoWithdraw,
+            &BalanceWithdrawStatus::Unknown,
         );
         let execution_params = match early_execution_error {
             Some(error) => ExecutionOrEarlyError::Err(error),
@@ -2584,7 +2588,7 @@ impl AuthorityState {
             &checked_input_objects,
             self.config.certificate_deny_config.certificate_deny_set(),
             // TODO(address-balances): Mimic withdraw scheduling and pass the result.
-            &BalanceWithdrawStatus::NoWithdraw,
+            &BalanceWithdrawStatus::Unknown,
         );
         let execution_params = match early_execution_error {
             Some(error) => ExecutionOrEarlyError::Err(error),
@@ -2829,7 +2833,7 @@ impl AuthorityState {
             &checked_input_objects,
             self.config.certificate_deny_config.certificate_deny_set(),
             // TODO(address-balances): Mimic withdraw scheduling and pass the result.
-            &BalanceWithdrawStatus::NoWithdraw,
+            &BalanceWithdrawStatus::Unknown,
         );
         let execution_params = match early_execution_error {
             Some(error) => ExecutionOrEarlyError::Err(error),
@@ -3910,7 +3914,7 @@ impl AuthorityState {
         Ok(new_epoch_store)
     }
 
-    pub async fn settle_transactions_for_testing(
+    pub async fn settle_accumulator_for_testing(
         &self,
         ckpt_seq: CheckpointSequenceNumber,
         effects: &[TransactionEffects],
@@ -3921,6 +3925,7 @@ impl AuthorityState {
             ckpt_seq,
             0,
         );
+        let balance_changes = builder.collect_accumulator_changes();
         let epoch_store = self.epoch_store_for_testing();
         let epoch = epoch_store.epoch();
         let (settlements, barrier) = builder.build_tx(
@@ -3978,12 +3983,17 @@ impl AuthorityState {
             .next()
             .unwrap()
             .1;
+        let next_accumulator_version = assigned_versions.accumulator_version.unwrap().next();
         let env = ExecutionEnv::new().with_assigned_versions(assigned_versions);
         let (effects, _) = self
             .try_execute_immediately(&barrier, env, &epoch_store)
             .await
             .unwrap();
         assert!(effects.status().is_ok());
+        self.execution_scheduler.settle_balances(BalanceSettlement {
+            balance_changes,
+            next_accumulator_version,
+        });
     }
 
     /// Advance the epoch store to the next epoch for testing only.
@@ -6045,7 +6055,7 @@ impl AuthorityState {
         let input_objects = self.read_objects_for_execution(
             &tx_lock,
             &executable_tx,
-            assigned_versions,
+            &assigned_versions,
             epoch_store,
         )?;
 
@@ -6055,7 +6065,7 @@ impl AuthorityState {
                 &executable_tx,
                 input_objects,
                 None,
-                BalanceWithdrawStatus::NoWithdraw,
+                ExecutionEnv::default(),
                 epoch_store,
             )
             .unwrap();

crates/sui-core/src/checkpoints/checkpoint_executor/mod.rs

@@ -876,8 +876,8 @@ impl CheckpointExecutor {
                             .with_barrier_dependencies(barrier_deps);
 
                         // Check if the expected effects indicate insufficient balance
-                        if let ExecutionStatus::Failure {
-                            error: ExecutionFailureStatus::InsufficientBalanceForWithdraw,
+                        if let &ExecutionStatus::Failure {
+                            error: ExecutionFailureStatus::InsufficientFunds,
                             ..
                         } = effects.status()
                         {

crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/e2e_tests.rs renamed to crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/address_balance/e2e_tests.rs

@@ -28,7 +28,7 @@ use sui_types::{
 use tokio::sync::mpsc::{self, unbounded_channel};
 use tokio::time::timeout;
 
-use crate::execution_scheduler::balance_withdraw_scheduler::BalanceSettlement;
+use super::BalanceSettlement;
 use crate::{
     authority::{
         AuthorityState, ExecutionEnv, shared_object_version_manager::Schedulable,
@@ -158,7 +158,7 @@ impl TestEnv {
             let cert = self.receive_certificate().await.unwrap();
             results.insert(
                 *cert.certificate.digest(),
-                cert.execution_env.withdraw_status,
+                cert.execution_env.balance_withdraw_status,
             );
         }
         assert_eq!(results, expected_results);
@@ -212,14 +212,8 @@ async fn test_withdraw_schedule_e2e() {
     test_env.enqueue_transactions(transactions.clone());
     test_env
         .expect_withdraw_results(BTreeMap::from([
-            (
-                *transactions[0].digest(),
-                BalanceWithdrawStatus::SufficientBalance,
-            ),
-            (
-                *transactions[1].digest(),
-                BalanceWithdrawStatus::SufficientBalance,
-            ),
+            (*transactions[0].digest(), BalanceWithdrawStatus::Unknown),
+            (*transactions[1].digest(), BalanceWithdrawStatus::Unknown),
             (
                 *transactions[2].digest(),
                 BalanceWithdrawStatus::InsufficientBalance,
@@ -235,10 +229,7 @@ async fn test_withdraw_schedule_e2e() {
     test_env.settle_balances(BTreeMap::from([(test_env.account_objects[0], -500)]));
     test_env
         .expect_withdraw_results(BTreeMap::from([
-            (
-                *transactions[0].digest(),
-                BalanceWithdrawStatus::SufficientBalance,
-            ),
+            (*transactions[0].digest(), BalanceWithdrawStatus::Unknown),
             (
                 *transactions[1].digest(),
                 BalanceWithdrawStatus::InsufficientBalance,

crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/eager_scheduler.rs renamed to crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/address_balance/eager_scheduler.rs

@@ -14,13 +14,11 @@ use sui_types::{
 use tokio::sync::oneshot::Sender;
 use tracing::{debug, instrument};
 
-use crate::{
-    accumulators::balance_read::AccountBalanceRead,
-    execution_scheduler::balance_withdraw_scheduler::{
-        BalanceSettlement, ScheduleResult, ScheduleStatus, TxBalanceWithdraw,
-        scheduler::{BalanceWithdrawSchedulerTrait, WithdrawReservations},
-    },
+use super::{
+    BalanceSettlement, ScheduleResult, ScheduleStatus, TxBalanceWithdraw,
+    scheduler::{BalanceWithdrawSchedulerTrait, WithdrawReservations},
 };
+use crate::accumulators::balance_read::AccountBalanceRead;
 
 pub(crate) struct EagerBalanceWithdrawScheduler {
     balance_read: Arc<dyn AccountBalanceRead>,

crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/address_balance/mod.rs

@@ -0,0 +1,61 @@
+// Copyright (c) Mysten Labs, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+use std::collections::BTreeMap;
+
+use sui_types::{
+    accumulator_root::AccumulatorObjId, base_types::SequenceNumber, digests::TransactionDigest,
+};
+
+mod eager_scheduler;
+mod naive_scheduler;
+pub(crate) mod scheduler;
+#[cfg(test)]
+mod tests;
+
+#[cfg(test)]
+mod e2e_tests;
+
+/// The status of scheduling the withdraw reservations for a transaction.
+#[derive(Debug, Copy, Clone, Eq, PartialEq)]
+pub(crate) enum ScheduleStatus {
+    /// We know for sure that the withdraw reservations in this transactions all have enough balance.
+    /// This transaction can be executed normally as soon as its object dependencies are ready.
+    SufficientBalance,
+    /// We know for sure that the withdraw reservations in this transactions do not all have enough balance.
+    /// This transaction should result in an execution failure without actually executing it, similar to
+    /// how transaction cancellation works.
+    InsufficientBalance,
+    /// We can skip scheduling this transaction, due to one of the following reasons:
+    /// 1. The accumulator version for this transaction has already been settled.
+    /// 2. We are observing some account objects bumping to the next version, indicating
+    ///    that the withdraw transactions in this commit have already been executed and are
+    ///    being settled.
+    SkipSchedule,
+}
+
+/// The result of scheduling the withdraw reservations for a transaction.
+#[derive(Debug, Clone, Eq, PartialEq)]
+pub(crate) struct ScheduleResult {
+    pub tx_digest: TransactionDigest,
+    pub status: ScheduleStatus,
+}
+
+/// Details regarding a balance settlement, generated when a settlement transaction has been executed
+/// and committed to the writeback cache.
+#[derive(Debug, Clone)]
+pub struct BalanceSettlement {
+    // After this settlement, the accumulator object will be at this version.
+    // This means that all transactions that read `next_accumulator_version - 1`
+    // are settled as part of this settlement.
+    pub next_accumulator_version: SequenceNumber,
+    /// The balance changes for each account object ID.
+    pub balance_changes: BTreeMap<AccumulatorObjId, i128>,
+}
+
+/// Details regarding all balance withdraw reservations in a transaction.
+#[derive(Clone, Debug)]
+pub(crate) struct TxBalanceWithdraw {
+    pub tx_digest: TransactionDigest,
+    pub reservations: BTreeMap<AccumulatorObjId, u64>,
+}

crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/naive_scheduler.rs renamed to crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/address_balance/naive_scheduler.rs

@@ -7,13 +7,11 @@ use sui_types::base_types::SequenceNumber;
 use tokio::sync::watch;
 use tracing::{debug, instrument};
 
-use crate::{
-    accumulators::balance_read::AccountBalanceRead,
-    execution_scheduler::balance_withdraw_scheduler::{
-        BalanceSettlement, ScheduleResult, ScheduleStatus,
-        scheduler::{BalanceWithdrawSchedulerTrait, WithdrawReservations},
-    },
+use super::{
+    BalanceSettlement, ScheduleResult, ScheduleStatus,
+    scheduler::{BalanceWithdrawSchedulerTrait, WithdrawReservations},
 };
+use crate::accumulators::balance_read::AccountBalanceRead;
 
 /// A naive implementation of the balance withdraw scheduler that does not attempt to optimize the scheduling.
 /// For each withdraw reservation, it will always wait until the dependent accumulator object is available,

crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/scheduler.rs renamed to crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/address_balance/scheduler.rs

@@ -3,14 +3,11 @@
 
 use std::{collections::BTreeMap, sync::Arc};
 
-use crate::{
-    accumulators::balance_read::AccountBalanceRead,
-    execution_scheduler::balance_withdraw_scheduler::{
-        BalanceSettlement, ScheduleResult, ScheduleStatus, TxBalanceWithdraw,
-        eager_scheduler::EagerBalanceWithdrawScheduler,
-        naive_scheduler::NaiveBalanceWithdrawScheduler,
-    },
+use super::{
+    BalanceSettlement, ScheduleResult, ScheduleStatus, TxBalanceWithdraw,
+    eager_scheduler::EagerBalanceWithdrawScheduler, naive_scheduler::NaiveBalanceWithdrawScheduler,
 };
+use crate::accumulators::balance_read::AccountBalanceRead;
 use futures::stream::FuturesUnordered;
 use mysten_metrics::monitored_mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
 use sui_types::base_types::SequenceNumber;

crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/tests.rs renamed to crates/sui-core/src/execution_scheduler/balance_withdraw_scheduler/address_balance/tests.rs

@@ -1,9 +1,10 @@
 // Copyright (c) Mysten Labs, Inc.
 // SPDX-License-Identifier: Apache-2.0
 
-use crate::execution_scheduler::balance_withdraw_scheduler::ScheduleResult;
-use crate::execution_scheduler::balance_withdraw_scheduler::{
-    BalanceSettlement, ScheduleStatus, TxBalanceWithdraw, mock_balance_read::MockBalanceRead,
+use crate::execution_scheduler::balance_withdraw_scheduler::mock_balance_read::MockBalanceRead;
+
+use super::{
+    BalanceSettlement, ScheduleResult, ScheduleStatus, TxBalanceWithdraw,
     scheduler::BalanceWithdrawScheduler,
 };
 use futures::StreamExt;