diff --git a/rs/embedders/src/wasm_utils/validation.rs b/rs/embedders/src/wasm_utils/validation.rs index 299695b5c748..b6089a077424 100644 --- a/rs/embedders/src/wasm_utils/validation.rs +++ b/rs/embedders/src/wasm_utils/validation.rs @@ -805,6 +805,16 @@ fn get_valid_system_apis_common( }, )], ), + ( + "subnet_self_node_count", + vec![( + API_VERSION_IC0, + FunctionSignature { + param_types: vec![], + return_type: vec![DataType::I32], + }, + )], + ), ( "cost_call", vec![( diff --git a/rs/embedders/src/wasmtime_embedder/linker.rs b/rs/embedders/src/wasmtime_embedder/linker.rs index 5897f2dd7574..c8d7975e0821 100644 --- a/rs/embedders/src/wasmtime_embedder/linker.rs +++ b/rs/embedders/src/wasmtime_embedder/linker.rs @@ -1022,6 +1022,15 @@ pub fn syscalls< }) .unwrap(); + linker + .func_wrap("ic0", "subnet_self_node_count", { + move |mut caller: Caller<'_, StoreData>| { + charge_for_cpu(&mut caller, overhead::SUBNET_SELF_NODE_COUNT)?; + with_system_api(&mut caller, |s| s.ic0_subnet_self_node_count()) + } + }) + .unwrap(); + match main_memory_type { WasmMemoryType::Wasm32 => { linker diff --git a/rs/embedders/src/wasmtime_embedder/system_api.rs b/rs/embedders/src/wasmtime_embedder/system_api.rs index 6bcb4ee10e13..66a27e583578 100644 --- a/rs/embedders/src/wasmtime_embedder/system_api.rs +++ b/rs/embedders/src/wasmtime_embedder/system_api.rs @@ -4672,6 +4672,32 @@ impl SystemApi for SystemApiImpl { result } + + fn ic0_subnet_self_node_count(&self) -> HypervisorResult { + let result = match &self.api_type { + ApiType::Start { .. } => Err(self.error_for("ic0.subnet_self_node_count")), + ApiType::Init { .. } + | ApiType::SystemTask { .. } + | ApiType::Cleanup { .. } + | ApiType::CompositeCleanup { .. } + | ApiType::Update { .. } + | ApiType::ReplicatedQuery { .. } + | ApiType::NonReplicatedQuery { .. } + | ApiType::CompositeQuery { .. } + | ApiType::ReplyCallback { .. } + | ApiType::CompositeReplyCallback { .. } + | ApiType::RejectCallback { .. } + | ApiType::CompositeRejectCallback { .. } + | ApiType::PreUpgrade { .. } + | ApiType::InspectMessage { .. } => Ok(self + .sandbox_safe_system_state + .subnet_cycles_config + .subnet_size as u32), + }; + + trace_syscall!(self, SubnetSelfNodeCount, result); + result + } } /// The default implementation of the `OutOfInstructionHandler` trait. diff --git a/rs/embedders/src/wasmtime_embedder/system_api_complexity.rs b/rs/embedders/src/wasmtime_embedder/system_api_complexity.rs index 00ac349a6293..fd97965ceeea 100644 --- a/rs/embedders/src/wasmtime_embedder/system_api_complexity.rs +++ b/rs/embedders/src/wasmtime_embedder/system_api_complexity.rs @@ -78,6 +78,7 @@ pub mod overhead { pub const PERFORMANCE_COUNTER: NumInstructions = NumInstructions::new(200); pub const SUBNET_SELF_SIZE: NumInstructions = NumInstructions::new(500); pub const SUBNET_SELF_COPY: NumInstructions = NumInstructions::new(500); + pub const SUBNET_SELF_NODE_COUNT: NumInstructions = NumInstructions::new(500); pub const STABLE_GROW: NumInstructions = NumInstructions::new(500); pub const STABLE_READ: NumInstructions = NumInstructions::new(20); pub const STABLE_SIZE: NumInstructions = NumInstructions::new(20); diff --git a/rs/embedders/tests/system_api.rs b/rs/embedders/tests/system_api.rs index 1e9484d88ed0..da0afe019c47 100644 --- a/rs/embedders/tests/system_api.rs +++ b/rs/embedders/tests/system_api.rs @@ -260,6 +260,7 @@ fn is_supported(api_type: SystemApiCallId, context: &str) -> bool { SystemApiCallId::MintCycles128 => vec!["U", "Ry", "Rt", "T"], SystemApiCallId::SubnetSelfSize => vec!["*"], SystemApiCallId::SubnetSelfCopy => vec!["*"], + SystemApiCallId::SubnetSelfNodeCount => vec!["*"], SystemApiCallId::EnvVarCount => vec!["*"], SystemApiCallId::EnvVarNameSize => vec!["*"], SystemApiCallId::EnvVarNameCopy => vec!["*"], @@ -843,6 +844,16 @@ fn api_availability_test( context, ); } + SystemApiCallId::SubnetSelfNodeCount => { + assert_api_availability( + |api| api.ic0_subnet_self_node_count(), + api_type, + &system_state, + cycles_account_manager, + api_type_enum, + context, + ); + } SystemApiCallId::EnvVarCount => { assert_api_availability( |api| api.ic0_env_var_count(), diff --git a/rs/embedders/tests/wasmtime_embedder.rs b/rs/embedders/tests/wasmtime_embedder.rs index 4bcc659068fd..87d3e5ea3e53 100644 --- a/rs/embedders/tests/wasmtime_embedder.rs +++ b/rs/embedders/tests/wasmtime_embedder.rs @@ -2174,6 +2174,8 @@ fn wasm64_import_system_api_functions() { (func $ic0_subnet_self_size (result i64))) (import "ic0" "subnet_self_copy" (func $ic0_subnet_self_copy (param i64) (param i64) (param i64))) + (import "ic0" "subnet_self_node_count" + (func $ic0_subnet_self_node_count (result i32))) (global $g1 (export "g1") (mut i64) (i64.const 0)) (func $test (export "canister_update test") @@ -2801,6 +2803,59 @@ fn wasm64_subnet_self_size() { ); } +#[test] +fn wasm64_subnet_self_node_count() { + // `ic0.subnet_self_node_count` returns an `i32` irrespective of the main + // memory type, so a Wasm64 canister imports it with an `i32` result too. + let wat = r#" + (module + (import "ic0" "subnet_self_node_count" + (func $ic0_subnet_self_node_count (result i32))) + + (global $g1 (export "g1") (mut i32) (i32.const 0)) + (func $test (export "canister_update test") + (call $ic0_subnet_self_node_count) + global.set $g1 + ) + + (memory (export "memory") i64 1) + )"#; + + let caller = user_test_id(24).get(); + let payload: Vec = vec![1, 3, 5, 7]; + let api = ApiType::update( + UNIX_EPOCH, + payload.clone(), + Cycles::zero(), + caller, + call_context_test_id(13), + None, + ); + + let config = ic_config::embedders::Config::default(); + let mut instance = WasmtimeInstanceBuilder::new() + .with_config(config) + .with_wat(wat) + .with_api_type(api) + .build(); + + let res = instance + .run(FuncRef::Method(WasmMethod::Update("test".to_string()))) + .unwrap(); + + assert_eq!( + res.exported_globals[0], + Global::I32( + instance + .store_data() + .system_api() + .unwrap() + .ic0_subnet_self_node_count() + .unwrap() as i32 + ) + ); +} + #[test] fn wasm64_subnet_self_copy() { let wat = r#" diff --git a/rs/execution_environment/src/query_handler/query_cache/tests.rs b/rs/execution_environment/src/query_handler/query_cache/tests.rs index 191b86ed0a16..622320999292 100644 --- a/rs/execution_environment/src/query_handler/query_cache/tests.rs +++ b/rs/execution_environment/src/query_handler/query_cache/tests.rs @@ -1794,6 +1794,7 @@ fn query_cache_future_proof_test() { | SystemApiCallId::PerformanceCounter | SystemApiCallId::SubnetSelfSize | SystemApiCallId::SubnetSelfCopy + | SystemApiCallId::SubnetSelfNodeCount | SystemApiCallId::Stable64Grow | SystemApiCallId::Stable64Read | SystemApiCallId::Stable64Size diff --git a/rs/execution_environment/tests/hypervisor.rs b/rs/execution_environment/tests/hypervisor.rs index 7448b7689e2b..3c329e8afc82 100644 --- a/rs/execution_environment/tests/hypervisor.rs +++ b/rs/execution_environment/tests/hypervisor.rs @@ -3295,6 +3295,77 @@ fn ic0_subnet_self_copy_works() { ); } +const SUBNET_SELF_NODE_COUNT_WAT: &str = r#" + (module + (import "ic0" "subnet_self_node_count" + (func $subnet_self_node_count (result i32)) + ) + (import "ic0" "msg_reply" (func $msg_reply)) + (import "ic0" "msg_reply_data_append" + (func $msg_reply_data_append (param i32 i32))) + (func (export "canister_update test") + ;; heap[0..4] = $subnet_self_node_count() + (i32.store (i32.const 0) (call $subnet_self_node_count)) + ;; return heap[0..4] + (call $msg_reply_data_append (i32.const 0) (i32.const 4)) + (call $msg_reply) + ) + (memory 1 1) + )"#; + +#[test] +fn ic0_subnet_self_node_count_works() { + let mut test = ExecutionTestBuilder::new().build(); + let canister_id = test.canister_from_wat(SUBNET_SELF_NODE_COUNT_WAT).unwrap(); + let result = test.ingress(canister_id, "test", vec![]).unwrap(); + assert_eq!( + WasmResult::Reply((test.subnet_size() as u32).to_le_bytes().to_vec()), + result + ); +} + +#[test] +fn ic0_subnet_self_node_count_reflects_subnet_size() { + // A size that differs from the default, so that the returned value cannot + // accidentally match a hard-coded default. + const SUBNET_SIZE: usize = 34; + let mut test = ExecutionTestBuilder::new() + .with_subnet_size(SUBNET_SIZE) + .build(); + assert_eq!(test.subnet_size(), SUBNET_SIZE); + let canister_id = test.canister_from_wat(SUBNET_SELF_NODE_COUNT_WAT).unwrap(); + let result = test.ingress(canister_id, "test", vec![]).unwrap(); + assert_eq!( + WasmResult::Reply((SUBNET_SIZE as u32).to_le_bytes().to_vec()), + result + ); +} + +#[test] +fn ic0_subnet_self_node_count_traps_in_start() { + let mut test = ExecutionTestBuilder::new().build(); + let wat = r#" + (module + (import "ic0" "subnet_self_node_count" + (func $subnet_self_node_count (result i32)) + ) + (func $start + (drop (call $subnet_self_node_count)) + ) + (start $start) + (func (export "canister_update test")) + (memory 1 1) + )"#; + let err = test.canister_from_wat(wat).unwrap_err(); + assert_eq!(err.code(), ErrorCode::CanisterContractViolation); + assert!( + err.description() + .contains(r#""ic0.subnet_self_node_count" cannot be executed in start mode"#), + "unexpected error: {}", + err.description() + ); +} + #[test] fn ic0_call_new_must_be_called_first() { let mut test = ExecutionTestBuilder::new().build(); diff --git a/rs/interfaces/src/execution_environment.rs b/rs/interfaces/src/execution_environment.rs index 92af0f7ca271..7ce479aba039 100644 --- a/rs/interfaces/src/execution_environment.rs +++ b/rs/interfaces/src/execution_environment.rs @@ -270,6 +270,8 @@ pub enum SystemApiCallId { SubnetSelfSize, /// Tracker for `ic0.subnet_self_copy()` SubnetSelfCopy, + /// Tracker for `ic0.subnet_self_node_count()` + SubnetSelfNodeCount, /// Tracker for `ic0.stable64_grow()` Stable64Grow, /// Tracker for `ic0.stable64_read()` @@ -1432,6 +1434,10 @@ pub trait SystemApi { size: usize, heap: &mut [u8], ) -> HypervisorResult<()>; + + /// Used to look up the number of nodes currently on the subnet that the + /// calling canister is running on. + fn ic0_subnet_self_node_count(&self) -> HypervisorResult; } #[derive(Copy, Clone, Eq, PartialEq, Debug)] diff --git a/rs/test_utilities/execution_environment/src/lib.rs b/rs/test_utilities/execution_environment/src/lib.rs index a334bb0b1ff8..08d96f86d5c6 100644 --- a/rs/test_utilities/execution_environment/src/lib.rs +++ b/rs/test_utilities/execution_environment/src/lib.rs @@ -2937,6 +2937,15 @@ impl ExecutionTestBuilder { self } + /// Sets the number of nodes on the own subnet. This is reflected both in the + /// registry settings and in the nodes of the own subnet's network topology. + pub fn with_subnet_size(mut self, subnet_size: usize) -> Self { + self.registry_settings.subnet_size = subnet_size; + self.registry_settings.node_ids = + (0..subnet_size).map(|i| node_test_id(i as u64)).collect(); + self + } + pub fn with_resource_limits(mut self, resource_limits: ResourceLimits) -> Self { self.resource_limits = resource_limits; self