Renames application insights keys to be more clear (#587)

* renames `telemetry_key` to `microsoft_telemetry_key`
* renames `instrumentation_key` to `instance_telemetry_key`
* renames `can_share` to `can_share_with_microsoft`
* renames the `applicationinsights-rs` instances to `internal` and `microsoft` respective of the keys used during construction.

This clarifies the underlying use of Application Insights keys and uses struct tuple to ensure the keys are used correctly via rust's type checker.

docs/telemetry.md

@@ -35,8 +35,8 @@ nodes and the service API running in the Azure Functions instance.
 
 1. The rust library [onefuzz::telemetry](../src/agent/onefuzz-telemetry/src/lib.rs)
    provides a detailed set of telemetry types, as well as the function
-   `can_share`, which gates if a given telemetry field should be sent to the
-   Microsoft central telemetry instance.
+   `can_share_with_microsoft`, which gates if a given telemetry field should be
+   sent to the Microsoft central telemetry instance.
 1. The Python library
    [onefuzzlib.telemetry](../src/api-service/__app__/onefuzzlib/telemetry.py)
    provides a filtering mechanism to identify a per-object set of filtering

src/agent/onefuzz-agent/src/local/common.rs

@@ -166,13 +166,11 @@ pub fn build_common_config(args: &ArgMatches<'_>) -> Result<CommonConfig> {
     };
 
     let config = CommonConfig {
-        heartbeat_queue: None,
-        instrumentation_key: None,
-        telemetry_key: None,
         job_id,
         task_id,
         instance_id,
         setup_dir,
+        ..Default::default()
     };
     Ok(config)
 }

src/agent/onefuzz-agent/src/managed/cmd.rs

@@ -23,7 +23,10 @@ pub async fn run(args: &clap::ArgMatches<'_>) -> Result<()> {
 }
 
 fn init_telemetry(config: &CommonConfig) {
-    onefuzz_telemetry::set_appinsights_clients(config.instrumentation_key, config.telemetry_key);
+    onefuzz_telemetry::set_appinsights_clients(
+        config.instance_telemetry_key.clone(),
+        config.microsoft_telemetry_key.clone(),
+    );
 }
 
 pub fn args(name: &str) -> App<'static, 'static> {

src/agent/onefuzz-agent/src/tasks/config.rs

@@ -5,7 +5,10 @@
 use crate::tasks::{analysis, coverage, fuzz, heartbeat::*, merge, report};
 use anyhow::Result;
 use onefuzz::machine_id::{get_machine_id, get_scaleset_name};
-use onefuzz_telemetry::{self as telemetry, Event::task_start, EventData, Role};
+use onefuzz_telemetry::{
+    self as telemetry, Event::task_start, EventData, InstanceTelemetryKey, MicrosoftTelemetryKey,
+    Role,
+};
 use reqwest::Url;
 use serde::{self, Deserialize};
 use std::path::PathBuf;
@@ -26,11 +29,15 @@ pub struct CommonConfig {
 
     pub instance_id: Uuid,
 
-    pub instrumentation_key: Option<Uuid>,
-
     pub heartbeat_queue: Option<Url>,
 
-    pub telemetry_key: Option<Uuid>,
+    // TODO: remove the alias once the service has been updated to match
+    #[serde(alias = "instrumentation_key")]
+    pub instance_telemetry_key: Option<InstanceTelemetryKey>,
+
+    // TODO: remove the alias once the service has been updated to match
+    #[serde(alias = "telemetry_key")]
+    pub microsoft_telemetry_key: Option<MicrosoftTelemetryKey>,
 
     #[serde(default)]
     pub setup_dir: PathBuf,

src/agent/onefuzz-supervisor/src/config.rs

@@ -6,6 +6,7 @@ use onefuzz::{
     http::{is_auth_error_code, ResponseExt},
     jitter::delay_with_jitter,
 };
+use onefuzz_telemetry::{InstanceTelemetryKey, MicrosoftTelemetryKey};
 use reqwest_retry::SendRetry;
 use std::{
     path::{Path, PathBuf},
@@ -25,9 +26,13 @@ pub struct StaticConfig {
 
     pub onefuzz_url: Url,
 
-    pub instrumentation_key: Option<Uuid>,
+    // TODO: remove the alias once the service has been updated to match
+    #[serde(alias = "instrumentation_key")]
+    pub instance_telemetry_key: Option<InstanceTelemetryKey>,
 
-    pub telemetry_key: Option<Uuid>,
+    // TODO: remove the alias once the service has been updated to match
+    #[serde(alias = "telemetry_key")]
+    pub microsoft_telemetry_key: Option<MicrosoftTelemetryKey>,
 
     pub heartbeat_queue: Option<Url>,
 
@@ -43,9 +48,13 @@ struct RawStaticConfig {
 
     pub onefuzz_url: Url,
 
-    pub instrumentation_key: Option<Uuid>,
+    // TODO: remove the alias once the service has been updated to match
+    #[serde(alias = "instrumentation_key")]
+    pub instance_telemetry_key: Option<InstanceTelemetryKey>,
 
-    pub telemetry_key: Option<Uuid>,
+    // TODO: remove the alias once the service has been updated to match
+    #[serde(alias = "telemetry_key")]
+    pub microsoft_telemetry_key: Option<MicrosoftTelemetryKey>,
 
     pub heartbeat_queue: Option<Url>,
 
@@ -73,8 +82,8 @@ impl StaticConfig {
             credentials,
             pool_name: config.pool_name,
             onefuzz_url: config.onefuzz_url,
-            instrumentation_key: config.instrumentation_key,
-            telemetry_key: config.telemetry_key,
+            microsoft_telemetry_key: config.microsoft_telemetry_key,
+            instance_telemetry_key: config.instance_telemetry_key,
             heartbeat_queue: config.heartbeat_queue,
             instance_id: config.instance_id,
         };
@@ -103,17 +112,19 @@ impl StaticConfig {
             None
         };
 
-        let instrumentation_key = if let Ok(key) = std::env::var("ONEFUZZ_INSTRUMENTATION_KEY") {
-            Some(Uuid::parse_str(&key)?)
-        } else {
-            None
-        };
+        let instance_telemetry_key =
+            if let Ok(key) = std::env::var("ONEFUZZ_INSTANCE_TELEMETRY_KEY") {
+                Some(InstanceTelemetryKey::new(Uuid::parse_str(&key)?))
+            } else {
+                None
+            };
 
-        let telemetry_key = if let Ok(key) = std::env::var("ONEFUZZ_TELEMETRY_KEY") {
-            Some(Uuid::parse_str(&key)?)
-        } else {
-            None
-        };
+        let microsoft_telemetry_key =
+            if let Ok(key) = std::env::var("ONEFUZZ_MICROSOFT_TELEMETRY_KEY") {
+                Some(MicrosoftTelemetryKey::new(Uuid::parse_str(&key)?))
+            } else {
+                None
+            };
 
         let credentials =
             ClientCredentials::new(client_id, client_secret, onefuzz_url.to_string(), tenant)
@@ -124,8 +135,8 @@ impl StaticConfig {
             credentials,
             pool_name,
             onefuzz_url,
-            instrumentation_key,
-            telemetry_key,
+            instance_telemetry_key,
+            microsoft_telemetry_key,
             heartbeat_queue,
         })
     }

src/agent/onefuzz-supervisor/src/main.rs

@@ -6,8 +6,6 @@ extern crate async_trait;
 #[macro_use]
 extern crate downcast_rs;
 #[macro_use]
-extern crate serde;
-#[macro_use]
 extern crate clap;
 #[macro_use]
 extern crate anyhow;
@@ -226,7 +224,8 @@ async fn run_agent(config: StaticConfig) -> Result<()> {
 }
 
 fn init_telemetry(config: &StaticConfig) {
-    let inst_key = config.instrumentation_key;
-    let tele_key = config.telemetry_key;
-    telemetry::set_appinsights_clients(inst_key, tele_key);
+    telemetry::set_appinsights_clients(
+        config.instance_telemetry_key.clone(),
+        config.microsoft_telemetry_key.clone(),
+    );
 }

src/agent/onefuzz-telemetry/Cargo.toml

@@ -9,6 +9,7 @@ license = "MIT"
 appinsights = "0.1"
 log = "0.4"
 uuid = { version = "0.8", features = ["serde", "v4"] }
+serde = { version = "1.0", features = ["derive"] }
 
 [dev-dependencies]
 tokio = { version = "0.2" }

src/agent/onefuzz-telemetry/src/lib.rs

@@ -1,15 +1,46 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+use serde::{Deserialize, Serialize};
+use std::fmt;
 use std::sync::{LockResult, RwLockReadGuard, RwLockWriteGuard};
 use uuid::Uuid;
 
 pub use appinsights::telemetry::SeverityLevel::{Critical, Error, Information, Verbose, Warning};
 
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
+#[serde(transparent)]
+pub struct MicrosoftTelemetryKey(Uuid);
+impl MicrosoftTelemetryKey {
+    pub fn new(value: Uuid) -> Self {
+        Self(value)
+    }
+}
+
+impl fmt::Display for MicrosoftTelemetryKey {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize, PartialEq, Eq)]
+pub struct InstanceTelemetryKey(Uuid);
+impl InstanceTelemetryKey {
+    pub fn new(value: Uuid) -> Self {
+        Self(value)
+    }
+}
+
+impl fmt::Display for InstanceTelemetryKey {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "{}", self.0)
+    }
+}
+
 pub type TelemetryClient = appinsights::TelemetryClient<appinsights::InMemoryChannel>;
 pub enum ClientType {
     Instance,
-    Shared,
+    Microsoft,
 }
 
 #[derive(Clone, Debug, PartialEq)]
@@ -133,7 +164,7 @@ impl EventData {
         }
     }
 
-    pub fn can_share(&self) -> bool {
+    pub fn can_share_with_microsoft(&self) -> bool {
         match self {
             // TODO: Request CELA review of Version, as having this for central stats
             //       would be useful to track uptake of new releases
@@ -184,20 +215,20 @@ mod global {
     #[derive(Default)]
     pub struct Clients {
         instance: Option<RwLock<TelemetryClient>>,
-        shared: Option<RwLock<TelemetryClient>>,
+        microsoft: Option<RwLock<TelemetryClient>>,
     }
 
     pub static mut CLIENTS: Clients = Clients {
         instance: None,
-        shared: None,
+        microsoft: None,
     };
     const UNSET: usize = 0;
     const SETTING: usize = 1;
     const SET: usize = 2;
 
     static STATE: AtomicUsize = AtomicUsize::new(UNSET);
 
-    pub fn set_clients(instance: Option<TelemetryClient>, shared: Option<TelemetryClient>) {
+    pub fn set_clients(instance: Option<TelemetryClient>, microsoft: Option<TelemetryClient>) {
         use Ordering::SeqCst;
 
         let result = STATE.compare_exchange(UNSET, SETTING, SeqCst, SeqCst);
@@ -212,7 +243,7 @@ mod global {
 
         unsafe {
             CLIENTS.instance = instance.map(RwLock::new);
-            CLIENTS.shared = shared.map(RwLock::new);
+            CLIENTS.microsoft = microsoft.map(RwLock::new);
         }
 
         STATE.store(SET, SeqCst);
@@ -221,7 +252,7 @@ mod global {
     pub fn client_lock(client_type: ClientType) -> Option<&'static RwLock<TelemetryClient>> {
         match client_type {
             ClientType::Instance => unsafe { CLIENTS.instance.as_ref() },
-            ClientType::Shared => unsafe { CLIENTS.shared.as_ref() },
+            ClientType::Microsoft => unsafe { CLIENTS.microsoft.as_ref() },
         }
     }
 
@@ -239,13 +270,13 @@ mod global {
         }
 
         let instance = unsafe { CLIENTS.instance.take() };
-        let shared = unsafe { CLIENTS.shared.take() };
+        let microsoft = unsafe { CLIENTS.microsoft.take() };
 
         STATE.store(UNSET, SeqCst);
 
         let mut clients = Vec::new();
 
-        for client in vec![instance, shared] {
+        for client in vec![instance, microsoft] {
             if let Some(client) = client {
                 match client.into_inner() {
                     Ok(c) => clients.push(c),
@@ -257,10 +288,13 @@ mod global {
     }
 }
 
-pub fn set_appinsights_clients(ikey: Option<Uuid>, tkey: Option<Uuid>) {
-    let instance_client = ikey.map(|k| TelemetryClient::new(k.to_string()));
-    let shared_client = tkey.map(|k| TelemetryClient::new(k.to_string()));
-    global::set_clients(instance_client, shared_client);
+pub fn set_appinsights_clients(
+    instance_key: Option<InstanceTelemetryKey>,
+    microsoft_key: Option<MicrosoftTelemetryKey>,
+) {
+    let instance_client = instance_key.map(|k| TelemetryClient::new(k.to_string()));
+    let microsoft_client = microsoft_key.map(|k| TelemetryClient::new(k.to_string()));
+    global::set_clients(instance_client, microsoft_client);
 }
 
 /// Try to submit any pending telemetry with a blocking call.
@@ -313,8 +347,8 @@ pub fn property(client_type: ClientType, key: impl AsRef<str>) -> Option<String>
 pub fn set_property(entry: EventData) {
     let (key, value) = entry.as_values();
 
-    if entry.can_share() {
-        if let Some(mut client) = client_mut(ClientType::Shared) {
+    if entry.can_share_with_microsoft() {
+        if let Some(mut client) = client_mut(ClientType::Microsoft) {
             client
                 .context_mut()
                 .properties_mut()
@@ -353,12 +387,12 @@ pub fn track_event(event: Event, properties: Vec<EventData>) {
         client.track(evt);
     }
 
-    if let Some(client) = client(ClientType::Shared) {
+    if let Some(client) = client(ClientType::Microsoft) {
         let mut evt = appinsights::telemetry::EventTelemetry::new(event.as_str());
         let props = evt.properties_mut();
 
         for property in &properties {
-            if property.can_share() {
+            if property.can_share_with_microsoft() {
                 let (name, val) = property.as_values();
                 props.insert(name.to_string(), val);
             }

src/agent/onefuzz/src/lib.rs

@@ -7,9 +7,6 @@ extern crate anyhow;
 #[macro_use]
 extern crate lazy_static;
 
-#[macro_use]
-extern crate serde;
-
 #[macro_use]
 extern crate onefuzz_telemetry;