diff --git a/rs/system_api/src/lib.rs b/rs/system_api/src/lib.rs index 4925252e6f7..aa12d435620 100644 --- a/rs/system_api/src/lib.rs +++ b/rs/system_api/src/lib.rs @@ -3396,15 +3396,18 @@ impl SystemApi for SystemApiImpl { let result = match &self.api_type { ApiType::Start { .. } | ApiType::Init { .. } - | ApiType::ReplyCallback { .. } - | ApiType::RejectCallback { .. } | ApiType::Cleanup { .. } | ApiType::PreUpgrade { .. } - | ApiType::InspectMessage { .. } | ApiType::Update { .. } | ApiType::SystemTask { .. } | ApiType::ReplicatedQuery { .. } => Ok(1), - ApiType::NonReplicatedQuery { .. } => Ok(0), + ApiType::ReplyCallback { .. } | ApiType::RejectCallback { .. } => { + match self.execution_parameters.execution_mode { + ExecutionMode::NonReplicated => Ok(0), + ExecutionMode::Replicated => Ok(1), + } + } + ApiType::InspectMessage { .. } | ApiType::NonReplicatedQuery { .. } => Ok(0), }; trace_syscall!(self, ic0_in_replicated_execution, result); result diff --git a/rs/system_api/tests/common/mod.rs b/rs/system_api/tests/common/mod.rs index 84511a5dc29..043f096a127 100644 --- a/rs/system_api/tests/common/mod.rs +++ b/rs/system_api/tests/common/mod.rs @@ -99,6 +99,10 @@ impl ApiTypeBuilder { ) } + pub fn build_replicated_query_api() -> ApiType { + ApiType::replicated_query(UNIX_EPOCH, vec![], user_test_id(1).get()) + } + pub fn build_non_replicated_query_api() -> ApiType { ApiType::non_replicated_query( UNIX_EPOCH, @@ -175,6 +179,27 @@ impl ApiTypeBuilder { 0.into(), ) } + + pub fn build_inspect_message_api() -> ApiType { + ApiType::inspect_message( + PrincipalId::new_anonymous(), + "test".to_string(), + vec![], + UNIX_EPOCH, + ) + } + + pub fn build_start_api() -> ApiType { + ApiType::start(UNIX_EPOCH) + } + + pub fn build_init_api() -> ApiType { + ApiType::init(UNIX_EPOCH, vec![], user_test_id(1).get()) + } + + pub fn build_pre_upgrade_api() -> ApiType { + ApiType::pre_upgrade(UNIX_EPOCH, user_test_id(1).get()) + } } pub fn get_system_api( diff --git a/rs/system_api/tests/system_api.rs b/rs/system_api/tests/system_api.rs index 5cb55e539c1..8e47ccf54ff 100644 --- a/rs/system_api/tests/system_api.rs +++ b/rs/system_api/tests/system_api.rs @@ -2069,3 +2069,40 @@ fn test_save_log_message_keeps_total_log_size_limited() { assert_eq!(log.records().len(), initial_records_number + 1); assert_le!(log.used_space(), MAX_ALLOWED_CANISTER_LOG_BUFFER_SIZE); } + +#[test] +fn in_replicated_execution_works_correctly() { + // The following should execute in replicated mode. + for api_type in &[ + ApiTypeBuilder::build_update_api(), + ApiTypeBuilder::build_system_task_api(), + ApiTypeBuilder::build_start_api(), + ApiTypeBuilder::build_init_api(), + ApiTypeBuilder::build_pre_upgrade_api(), + ApiTypeBuilder::build_replicated_query_api(), + ApiTypeBuilder::build_reply_api(Cycles::new(0)), + ApiTypeBuilder::build_reject_api(RejectContext::new(RejectCode::CanisterReject, "error")), + ] { + let cycles_account_manager = CyclesAccountManagerBuilder::new().build(); + let system_state = SystemStateBuilder::default().build(); + let api = get_system_api(api_type.clone(), &system_state, cycles_account_manager); + assert_eq!(api.ic0_in_replicated_execution(), Ok(1)); + } + + // The following should execute in non-replicated mode. + for api_type in &[ + ApiTypeBuilder::build_non_replicated_query_api(), + ApiTypeBuilder::build_composite_query_api(), + ApiTypeBuilder::build_composite_reply_api(Cycles::new(0)), + ApiTypeBuilder::build_composite_reject_api(RejectContext::new( + RejectCode::CanisterReject, + "error", + )), + ApiTypeBuilder::build_inspect_message_api(), + ] { + let cycles_account_manager = CyclesAccountManagerBuilder::new().build(); + let system_state = SystemStateBuilder::default().build(); + let api = get_system_api(api_type.clone(), &system_state, cycles_account_manager); + assert_eq!(api.ic0_in_replicated_execution(), Ok(0)); + } +}