feat: Migrating sns downloader to canister calls (#451)

* migrating to canister call instead of api calls

* adding back sns name

Cargo.Bazel.lock

@@ -1,5 +1,5 @@
 {
-  "checksum": "582a531639bbb97df204fdb6b1988fe1e94551b0bd30d13d773fb82fb866451b",
+  "checksum": "9fc786f423c08cb67bfd86f1b1da73c6101907a6c2b6c7880874b41bbf3b4190",
   "crates": {
     "actix-codec 0.5.2": {
       "name": "actix-codec",
@@ -17725,6 +17725,10 @@
               "id": "ic-registry-transport 0.9.0",
               "target": "ic_registry_transport"
             },
+            {
+              "id": "ic-sns-wasm 1.0.0",
+              "target": "ic_sns_wasm"
+            },
             {
               "id": "ic-sys 0.9.0",
               "target": "ic_sys"
@@ -30245,6 +30249,10 @@
               "id": "erased-serde 0.4.5",
               "target": "erased_serde"
             },
+            {
+              "id": "ic-sns-wasm 1.0.0",
+              "target": "ic_sns_wasm"
+            },
             {
               "id": "ic-types 0.9.0",
               "target": "ic_types"

rs/ic-canisters/Cargo.toml

@@ -32,3 +32,4 @@ sha2 = { workspace = true }
 simple_asn1 = { workspace = true }
 thiserror = { workspace = true }
 url = { workspace = true }
+ic-sns-wasm = { workspace = true }

rs/ic-canisters/src/lib.rs

@@ -22,6 +22,7 @@ pub mod governance;
 pub mod management;
 pub mod parallel_hardware_identity;
 pub mod registry;
+pub mod sns_wasm;
 
 pub struct CanisterClient {
     pub agent: CanisterClientAgent,

rs/ic-canisters/src/sns_wasm.rs

@@ -0,0 +1,43 @@
+use std::str::FromStr;
+
+use crate::IcAgentCanisterClient;
+use ic_agent::Agent;
+use ic_base_types::CanisterId;
+use ic_sns_wasm::pb::v1::{ListDeployedSnsesRequest, ListDeployedSnsesResponse};
+
+const SNS_WASM_CANISTER: &str = "qaa6y-5yaaa-aaaaa-aaafa-cai";
+
+#[derive(Clone)]
+pub struct SnsWasmCanister {
+    agent: Agent,
+    sns_wasm_canister: CanisterId,
+}
+
+impl From<IcAgentCanisterClient> for SnsWasmCanister {
+    fn from(value: IcAgentCanisterClient) -> Self {
+        Self {
+            agent: value.agent,
+            sns_wasm_canister: CanisterId::from_str(SNS_WASM_CANISTER).unwrap(),
+        }
+    }
+}
+
+impl SnsWasmCanister {
+    pub fn new(agent: Agent, sns_canister_id: Option<&str>) -> Self {
+        Self {
+            agent,
+            sns_wasm_canister: CanisterId::from_str(sns_canister_id.unwrap_or(SNS_WASM_CANISTER)).unwrap(),
+        }
+    }
+
+    pub async fn list_deployed_snses(&self) -> anyhow::Result<ListDeployedSnsesResponse> {
+        let arg = candid::encode_one(ListDeployedSnsesRequest {})?;
+        let response = self
+            .agent
+            .query(&self.sns_wasm_canister.into(), "list_deployed_snses")
+            .with_arg(arg)
+            .call()
+            .await?;
+        Ok(candid::decode_one(&response)?)
+    }
+}

rs/ic-observability/multiservice-discovery-shared/Cargo.toml

@@ -12,3 +12,4 @@ serde_json = { workspace = true }
 ic-types = { workspace = true }
 erased-serde = { workspace = true }
 regex = { workspace = true }
+ic-sns-wasm = { workspace = true }

rs/ic-observability/multiservice-discovery-shared/src/builders/sns_canister_config_structure.rs

@@ -2,7 +2,7 @@ use std::collections::HashMap;
 
 use serde::Serialize;
 
-use crate::contracts::sns::Sns;
+use crate::contracts::deployed_sns::Sns;
 
 use super::{
     log_vector_config_structure::VectorRemapTransform,
@@ -17,62 +17,62 @@ pub struct SnsCanisterConfigStructure {
     pub include_stderr: bool,
 }
 
-static SCRAPABLE_TYPES: [&str; 3] = ["root", "swap", "governance"];
-
+// Scrapable types are: root, swap, governance
 impl SnsCanisterConfigStructure {
     pub fn build(&self, snses: Vec<Sns>) -> String {
         let mut config = VectorConfigEnriched::new();
 
         for sns in snses {
-            for canister in sns.canisters {
-                if !SCRAPABLE_TYPES.contains(&canister.canister_type.as_str()) {
-                    continue;
-                }
-
-                let key = canister.canister_id.to_string();
-                let source = VectorScriptSource {
-                    _type: "exec".to_string(),
-                    command: [
-                        self.script_path.as_str(),
-                        "--url",
-                        format!("https://{}.raw.icp0.io/logs", canister.canister_id).as_str(),
-                    ]
-                    .iter()
-                    .map(|s| s.to_string())
-                    .collect(),
-                    mode: "streaming".to_string(),
-                    streaming: SourceStreamingWrapper {
-                        respawn_on_exit: self.restart_on_exit,
-                    },
-                    include_stderr: self.include_stderr,
-                };
-
-                let transform = VectorRemapTransform {
-                    _type: "remap".to_string(),
-                    inputs: vec![key.clone()],
-                    source: vec![
-                        ("canister_id", canister.canister_id),
-                        ("canister_type", canister.canister_type),
-                        ("module_hash", canister.module_hash),
-                    ]
-                    .into_iter()
-                    .map(|(k, v)| format!(".{} = \"{}\"", k, v))
-                    .collect::<Vec<String>>()
-                    .join("\n"),
-                };
-
-                let mut sources = HashMap::new();
-                sources.insert(key.clone(), Box::new(source) as Box<dyn VectorSource>);
-
-                let mut transforms = HashMap::new();
-                transforms.insert(format!("{}-transform", key), Box::new(transform) as Box<dyn VectorTransform>);
-
-                config.add_target_group(sources, transforms)
+            if sns.root_canister_id != String::default() {
+                self.insert_into_config(&mut config, &sns.root_canister_id, "root", &sns.name);
+            }
+            if sns.swap_canister_id != String::default() {
+                self.insert_into_config(&mut config, &sns.swap_canister_id, "swap", &sns.name);
+            }
+            if sns.governance_canister_id != String::default() {
+                self.insert_into_config(&mut config, &sns.governance_canister_id, "governance", &sns.name);
             }
         }
 
         serde_json::to_string_pretty(&config).unwrap()
     }
+
+    fn insert_into_config(&self, config: &mut VectorConfigEnriched, canister_id: &str, canister_type: &str, sns_name: &str) {
+        let source = VectorScriptSource {
+            _type: "exec".to_string(),
+            command: [
+                self.script_path.as_str(),
+                "--url",
+                format!("https://{}.raw.icp0.io/logs", canister_id).as_str(),
+            ]
+            .iter()
+            .map(|s| s.to_string())
+            .collect(),
+            mode: "streaming".to_string(),
+            streaming: SourceStreamingWrapper {
+                respawn_on_exit: self.restart_on_exit,
+            },
+            include_stderr: self.include_stderr,
+        };
+
+        let transform = VectorRemapTransform {
+            _type: "remap".to_string(),
+            inputs: vec![canister_id.to_string()],
+            source: vec![("canister_id", canister_id), ("canister_type", canister_type), ("sns_name", sns_name)]
+                .into_iter()
+                .map(|(k, v)| format!(".{} = \"{}\"", k, v))
+                .collect::<Vec<String>>()
+                .join("\n"),
+        };
+
+        let mut sources = HashMap::new();
+        sources.insert(canister_id.to_string(), Box::new(source) as Box<dyn VectorSource>);
+
+        let mut transforms = HashMap::new();
+        transforms.insert(format!("{}-transform", canister_id), Box::new(transform) as Box<dyn VectorTransform>);
+
+        config.add_target_group(sources, transforms)
+    }
 }
 
 #[derive(Debug, Clone, Serialize)]

rs/ic-observability/multiservice-discovery-shared/src/contracts/deployed_sns.rs

@@ -0,0 +1,51 @@
+use std::hash::Hash;
+
+use ic_sns_wasm::pb::v1::DeployedSns;
+
+pub struct Sns {
+    pub root_canister_id: String,
+    pub governance_canister_id: String,
+    pub ledger_canister_id: String,
+    pub swap_canister_id: String,
+    pub index_canister_id: String,
+    pub name: String,
+}
+
+impl From<DeployedSns> for Sns {
+    fn from(value: DeployedSns) -> Self {
+        Self {
+            governance_canister_id: match value.governance_canister_id {
+                None => "".to_string(),
+                Some(val) => val.to_string(),
+            },
+            index_canister_id: match value.index_canister_id {
+                None => "".to_string(),
+                Some(val) => val.to_string(),
+            },
+            ledger_canister_id: match value.ledger_canister_id {
+                None => "".to_string(),
+                Some(val) => val.to_string(),
+            },
+            root_canister_id: match value.root_canister_id {
+                None => "".to_string(),
+                Some(val) => val.to_string(),
+            },
+            swap_canister_id: match value.swap_canister_id {
+                None => "".to_string(),
+                Some(val) => val.to_string(),
+            },
+            name: "".to_string(),
+        }
+    }
+}
+
+impl Hash for Sns {
+    fn hash<H: std::hash::Hasher>(&self, state: &mut H) {
+        self.root_canister_id.hash(state);
+        self.governance_canister_id.hash(state);
+        self.ledger_canister_id.hash(state);
+        self.swap_canister_id.hash(state);
+        self.index_canister_id.hash(state);
+        self.name.hash(state);
+    }
+}

rs/ic-observability/multiservice-discovery-shared/src/contracts/mod.rs

@@ -1,4 +1,4 @@
-pub mod sns;
+pub mod deployed_sns;
 pub mod target;
 
 pub trait DataContract {

rs/ic-observability/multiservice-discovery-shared/src/filters/mod.rs

@@ -4,7 +4,6 @@ use crate::contracts::DataContract;
 
 pub mod ic_name_regex_filter;
 pub mod node_regex_id_filter;
-pub mod sns_name_regex_filter;
 
 pub trait TargetGroupFilter: Send + Sync + Debug {
     fn filter(&self, target_groups: &dyn DataContract) -> bool;

rs/ic-observability/sns-downloader/BUILD.bazel

@@ -4,7 +4,8 @@ load("@//rs:oci_images.bzl", "rust_binary_oci_image_rules")
 
 DEPS = [
     "//rs/ic-observability/service-discovery",
-    "//rs/ic-observability/multiservice-discovery-shared"
+    "//rs/ic-observability/multiservice-discovery-shared",
+    "//rs/ic-canisters"
 ]
 
 rust_binary(

rs/ic-observability/sns-downloader/Cargo.toml

@@ -21,3 +21,4 @@ slog-async = { workspace = true }
 slog-term = { workspace = true }
 tokio = { workspace = true }
 url = { workspace = true }
+ic-canisters = { path = "../../ic-canisters" }