chore: EXC-1735: Apply priority credit at the round start (#1736)

Since priority credit is now serialized in the state, applying it at the
beginning of the round offers functional equivalence while simplifying
the code.

rs/execution_environment/src/execution_environment.rs

@@ -72,7 +72,7 @@ use ic_types::{
     },
     methods::SystemMethod,
     nominal_cycles::NominalCycles,
-    CanisterId, Cycles, LongExecutionMode, NumBytes, NumInstructions, SubnetId, Time,
+    CanisterId, Cycles, NumBytes, NumInstructions, SubnetId, Time,
 };
 use ic_types::{messages::MessageId, methods::WasmMethod};
 use ic_wasm_types::WasmHash;
@@ -3205,10 +3205,6 @@ impl ExecutionEnvironment {
                 })
                 .collect();
             canister.apply_priority_credit();
-            // Aborting a long-running execution moves the canister to the
-            // default execution mode because the canister does not have a
-            // pending execution anymore.
-            canister.scheduler_state.long_execution_mode = LongExecutionMode::default();
             let canister_id = canister.canister_id();
             canister.system_state.apply_ingress_induction_cycles_debit(
                 canister_id,

rs/execution_environment/src/scheduler.rs

@@ -228,7 +228,7 @@ impl SchedulerImpl {
             let has_aborted_or_paused_execution =
                 canister.has_aborted_execution() || canister.has_paused_execution();
             if !has_aborted_or_paused_execution {
-                canister.scheduler_state.long_execution_mode = Default::default();
+                canister.apply_priority_credit();
             }
 
             let compute_allocation = canister.scheduler_state.compute_allocation;
@@ -567,26 +567,16 @@ impl SchedulerImpl {
     /// Invoked in the first iteration of the inner round to add the `Heartbeat`
     /// and `GlobalTimer` tasks that are carried out prior to processing
     /// any input messages.
-    /// It also returns the list of canisters that have non-zero priority credit.
-    fn initialize_inner_round(
-        &self,
-        state: &mut ReplicatedState,
-    ) -> (BTreeSet<CanisterId>, BTreeSet<CanisterId>) {
+    fn initialize_inner_round(&self, state: &mut ReplicatedState) -> BTreeSet<CanisterId> {
         let _timer = self
             .metrics
             .round_inner_heartbeat_overhead_duration
             .start_timer();
 
         let mut heartbeat_and_timer_canister_ids = BTreeSet::new();
-        let mut non_zero_priority_credit_canister_ids = BTreeSet::new();
 
         let now = state.time();
         for canister in state.canisters_iter_mut() {
-            // Remember all non-zero priority_credit canisters to apply it after the round.
-            if canister.scheduler_state.priority_credit != AccumulatedPriority::default() {
-                non_zero_priority_credit_canister_ids.insert(canister.system_state.canister_id);
-            }
-
             // Add `Heartbeat` or `GlobalTimer` for running canisters only.
             match canister.system_state.status {
                 CanisterStatus::Running { .. } => {}
@@ -627,10 +617,7 @@ impl SchedulerImpl {
                 }
             }
         }
-        (
-            heartbeat_and_timer_canister_ids,
-            non_zero_priority_credit_canister_ids,
-        )
+        heartbeat_and_timer_canister_ids
     }
 
     /// Performs multiple iterations of canister execution until the instruction
@@ -657,7 +644,6 @@ impl SchedulerImpl {
         let mut total_heap_delta = NumBytes::from(0);
 
         let mut heartbeat_and_timer_canister_ids = BTreeSet::new();
-        let mut non_zero_priority_credit_canister_ids = BTreeSet::new();
         let mut round_executed_canister_ids = BTreeSet::new();
 
         // Start iteration loop:
@@ -709,10 +695,7 @@ impl SchedulerImpl {
 
             // Add `Heartbeat` and `GlobalTimer` tasks to be executed before input messages.
             if is_first_iteration {
-                (
-                    heartbeat_and_timer_canister_ids,
-                    non_zero_priority_credit_canister_ids,
-                ) = self.initialize_inner_round(&mut state)
+                heartbeat_and_timer_canister_ids = self.initialize_inner_round(&mut state);
             }
 
             // Update subnet available memory before taking out the canisters.
@@ -830,16 +813,6 @@ impl SchedulerImpl {
                     | ExecutionTask::AbortedInstallCode { .. } => true,
                 });
             }
-            // Apply priority credit for all the finished executions.
-            for canister_id in &non_zero_priority_credit_canister_ids {
-                let canister = state.canister_state_mut(canister_id).unwrap();
-                match canister.next_execution() {
-                    NextExecution::None
-                    | NextExecution::StartNew
-                    | NextExecution::ContinueInstallCode => canister.apply_priority_credit(),
-                    NextExecution::ContinueLong => {}
-                }
-            }
         }
 
         // We only export metrics for "executable" canisters to ensure that the metrics

rs/replicated_state/src/canister_state.rs

@@ -144,10 +144,14 @@ impl CanisterState {
         }
     }
 
-    /// Apply priority credit
+    /// Applies priority credit and resets long execution mode.
     pub fn apply_priority_credit(&mut self) {
         self.scheduler_state.accumulated_priority -=
             std::mem::take(&mut self.scheduler_state.priority_credit);
+        // Aborting a long-running execution moves the canister to the
+        // default execution mode because the canister does not have a
+        // pending execution anymore.
+        self.scheduler_state.long_execution_mode = LongExecutionMode::default();
     }
 
     pub fn canister_id(&self) -> CanisterId {

rs/state_manager/src/tip.rs

@@ -1003,8 +1003,6 @@ fn serialize_canister_to_tip(
             metrics,
         )?;
 
-    // Priority credit must be zero at this point
-    assert_eq!(canister_state.scheduler_state.priority_credit.get(), 0);
     canister_layout.canister().serialize(
         CanisterStateBits {
             controllers: canister_state.system_state.controllers.clone(),