feat: ic0_subnet_self_node_count

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![(

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

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<u32> {
+        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.

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);

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(),

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<u8> = 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#"

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

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();

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<u32>;
 }
 
 #[derive(Copy, Clone, Eq, PartialEq, Debug)]

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