Update input to start function

rs/messaging/tests/queue_tests.rs

@@ -17,7 +17,7 @@ use ic_types::{
 use maplit::btreemap;
 use std::collections::BTreeSet;
 use std::sync::Arc;
-use xnet_test::Metrics;
+use xnet_test::{Metrics, StartArgs};
 
 const MAX_TICKS: u64 = 100;
 
@@ -96,14 +96,14 @@ impl SubnetPairProxy {
         payload_size_bytes: u64,
     ) -> Result<Vec<u8>, candid::Error> {
         let network_topology = vec![
-            vec![self.local_canister_id.get().to_vec()],
-            vec![self.remote_canister_id.get().to_vec()],
+            vec![self.local_canister_id.get().into()],
+            vec![self.remote_canister_id.get().into()],
         ];
-        Encode!(
-            &network_topology,
-            &canister_to_subnet_rate,
-            &payload_size_bytes
-        )
+        Encode!(&StartArgs {
+            network_topology,
+            canister_to_subnet_rate,
+            payload_size_bytes,
+        })
     }
 
     /// Calls the 'start' method on the local canister.

rs/rust_canisters/xnet_test/src/lib.rs

@@ -1,16 +1,22 @@
 use candid::{CandidType, Deserialize};
+use ic_cdk::api::management_canister::provisional::CanisterId;
 use serde::Serialize;
 use std::time::Duration;
 
-/// Id of a canister is an opaque blob.
-pub type CanisterId = Vec<u8>;
-
 /// Configuration of the network: the outer vector enumerates canisters
 /// installed on the same subnet.
 ///
 /// This message is used as request payload for "start" call.
 pub type NetworkTopology = Vec<Vec<CanisterId>>;
 
+/// Arguments for the "start" call of this canister.
+#[derive(Default, Clone, CandidType, Deserialize, Debug)]
+pub struct StartArgs {
+    pub network_topology: NetworkTopology,
+    pub canister_to_subnet_rate: u64,
+    pub payload_size_bytes: u64,
+}
+
 /// Metrics observed by this canister.
 ///
 /// This message is used as reply payload for "metrics" query.

rs/rust_canisters/xnet_test/src/main.rs

@@ -6,6 +6,7 @@
 //! cargo build --target wasm32-unknown-unknown --release
 //! ```
 use candid::{CandidType, Deserialize, Principal};
+use ic_cdk::api::management_canister::provisional::CanisterId;
 use ic_cdk::{
     api::{
         call::{call, call_with_payment},
@@ -21,7 +22,7 @@ use std::collections::BTreeMap;
 use std::convert::TryFrom;
 use std::str::FromStr;
 use std::time::Duration;
-use xnet_test::{CanisterId, Metrics, NetworkTopology};
+use xnet_test::{Metrics, NetworkTopology, StartArgs};
 
 thread_local! {
     /// Whether this canister is generating traffic.
@@ -144,13 +145,14 @@ fn log(message: &str) {
 /// Initializes network topology and instructs this canister to start sending
 /// requests to other canisters.
 #[update]
-fn start(network_topology: NetworkTopology, rate: u64, payload_size: u64) -> String {
+fn start(start_args: StartArgs) -> String {
+    ic_cdk::setup();
+
     NETWORK_TOPOLOGY.with(move |canisters| {
-        *canisters.borrow_mut() = network_topology;
+        *canisters.borrow_mut() = start_args.network_topology;
     });
-
-    PER_SUBNET_RATE.with(|r| *r.borrow_mut() = rate);
-    PAYLOAD_SIZE.with(|r| *r.borrow_mut() = payload_size);
+    PER_SUBNET_RATE.with(|r| *r.borrow_mut() = start_args.canister_to_subnet_rate);
+    PAYLOAD_SIZE.with(|r| *r.borrow_mut() = start_args.payload_size_bytes);
 
     RUNNING.with(|r| *r.borrow_mut() = true);
 
@@ -178,7 +180,7 @@ async fn fanout() {
             continue;
         }
 
-        if canisters.contains(&self_id.as_slice().to_vec()) {
+        if canisters.contains(&self_id) {
             // Same subnet
             continue;
         }
@@ -225,10 +227,7 @@ async fn fanout() {
 #[update]
 fn handle_request(req: Request) -> Reply {
     let caller = caller();
-    let in_seq_no = STATE.with(|s| {
-        s.borrow_mut()
-            .set_in_seq_no(caller.as_slice().to_vec(), req.seq_no)
-    });
+    let in_seq_no = STATE.with(|s| s.borrow_mut().set_in_seq_no(caller, req.seq_no));
 
     if req.seq_no <= in_seq_no {
         METRICS.with(|m| m.borrow_mut().seq_errors += 1);

rs/rust_canisters/xnet_test/xnet_test.did

@@ -1,5 +1,11 @@
+type StartArgs = record {
+  network_topology : vec vec principal;
+  canister_to_subnet_rate : nat64;
+  payload_size_bytes : nat64;
+};
+
 service : {
-  start : (vec vec blob, nat64, nat64) -> (text);
+  start : (StartArgs) -> (text);
   stop : () -> (text);
   return_cycles : () -> (text);
-}
+};