fix: Skip log message processing when logs are disabled (#831)

# What?

Refactors telemetry client to be a method which receives parameters and
subscribes accordingly

Also gets events based on logging configuration

# Why?

We needed to ensure that we don't subscribe to logs events when customer
is not enabling `DD_SERVERLESS_LOGS_ENABLED`

---------

Co-authored-by: jordan gonzález <30836115+duncanista@users.noreply.github.com>

Author: litianningdatadog

bottlecap/src/bin/bottlecap/main.rs

@@ -17,7 +17,7 @@ use tikv_jemallocator::Jemalloc;
 static GLOBAL: Jemalloc = Jemalloc;
 
 use bottlecap::{
-    DOGSTATSD_PORT, LAMBDA_RUNTIME_SLUG, TELEMETRY_PORT,
+    DOGSTATSD_PORT, LAMBDA_RUNTIME_SLUG,
     appsec::processor::{
         Error::FeatureDisabled as AppSecFeatureDisabled, Processor as AppSecProcessor,
     },
@@ -48,7 +48,7 @@ use bottlecap::{
         provider::Provider as TagProvider,
     },
     telemetry::{
-        client::TelemetryApiClient,
+        self, TELEMETRY_PORT,
         events::{TelemetryEvent, TelemetryRecord},
         listener::TelemetryListener,
     },
@@ -474,9 +474,14 @@ async fn extension_loop_active(
 
     let dogstatsd_cancel_token = start_dogstatsd(metrics_aggr_handle.clone()).await;
 
-    let telemetry_listener_cancel_token =
-        setup_telemetry_client(&r.extension_id, &aws_config.runtime_api, logs_agent_channel)
-            .await?;
+    let telemetry_listener_cancel_token = setup_telemetry_client(
+        client,
+        &r.extension_id,
+        &aws_config.runtime_api,
+        logs_agent_channel,
+        config.serverless_logs_enabled,
+    )
+    .await?;
 
     let otlp_cancel_token = start_otlp_agent(
         config,
@@ -1107,9 +1112,11 @@ async fn start_dogstatsd(metrics_aggr_handle: MetricsAggregatorHandle) -> Cancel
 }
 
 async fn setup_telemetry_client(
+    client: &Client,
     extension_id: &str,
     runtime_api: &str,
     logs_agent_channel: Sender<TelemetryEvent>,
+    logs_enabled: bool,
 ) -> anyhow::Result<CancellationToken> {
     let telemetry_listener =
         TelemetryListener::new(EXTENSION_HOST_IP, TELEMETRY_PORT, logs_agent_channel);
@@ -1121,15 +1128,16 @@ async fn setup_telemetry_client(
         }
     });
 
-    let telemetry_client = TelemetryApiClient::new(
-        extension_id.to_string(),
+    telemetry::subscribe(
+        client,
+        runtime_api,
+        extension_id,
         TELEMETRY_PORT,
-        runtime_api.to_string(),
-    );
-    telemetry_client
-        .subscribe()
-        .await
-        .map_err(|e| anyhow::anyhow!("Failed to subscribe to telemetry: {e:?}"))?;
+        logs_enabled,
+    )
+    .await
+    .map_err(|e| anyhow::anyhow!("Failed to subscribe to telemetry: {e:?}"))?;
+
     Ok(cancel_token)
 }
 

bottlecap/src/lib.rs

@@ -45,9 +45,3 @@ pub const FLUSH_RETRY_COUNT: usize = 3;
 
 // todo: make sure we can override those with environment variables
 pub const DOGSTATSD_PORT: u16 = 8125;
-
-pub const TELEMETRY_SUBSCRIPTION_ROUTE: &str = "2022-07-01/telemetry";
-// todo(astuyve) should be 8124 on /lambda/logs but
-// telemetry is implemented on a raw socket now and
-// does not multiplex routes on the same port.
-pub const TELEMETRY_PORT: u16 = 8999;

bottlecap/src/telemetry/client.rs

@@ -1,3 +1,63 @@
-pub mod client;
+use reqwest::{Client, Response};
+use tracing::debug;
+
+use crate::extension::{EXTENSION_ID_HEADER, base_url};
+
 pub mod events;
 pub mod listener;
+
+pub const TELEMETRY_SUBSCRIPTION_ROUTE: &str = "2022-07-01/telemetry";
+// todo(astuyve) should be 8124 on /lambda/logs but
+// telemetry is implemented on a raw socket now and
+// does not multiplex routes on the same port.
+pub const TELEMETRY_PORT: u16 = 8999;
+
+const PLATFORM_ONLY_EVENTS: &[&str] = &["platform"];
+const ALL_EVENTS: &[&str] = &["platform", "extension", "function"];
+
+/// Error conditions that can arise from extension operations
+#[derive(thiserror::Error, Debug)]
+pub enum ExtensionSubscriptionError {
+    #[error("Subscription request failed: {0}")]
+    HttpError(#[from] reqwest::Error),
+}
+
+fn get_subscription_event_types(logs_enabled: bool) -> Vec<&'static str> {
+    (if logs_enabled {
+        ALL_EVENTS
+    } else {
+        PLATFORM_ONLY_EVENTS
+    })
+    .to_vec()
+}
+
+pub async fn subscribe(
+    client: &Client,
+    runtime_api: &str,
+    extension_id: &str,
+    destination_port: u16,
+    logs_enabled: bool,
+) -> Result<Response, ExtensionSubscriptionError> {
+    let url = base_url(TELEMETRY_SUBSCRIPTION_ROUTE, runtime_api);
+    let response = client
+        .put(&url)
+        .header(EXTENSION_ID_HEADER, extension_id)
+        .json(&serde_json::json!({
+            "schemaVersion": "2022-12-13",
+            "destination": {
+                "protocol": "HTTP",
+                "URI": format!("http://sandbox:{}/", destination_port),
+            },
+            "types": get_subscription_event_types(logs_enabled),
+            "buffering": { // TODO: re evaluate using default values
+                "maxItems": 1000,
+                "maxBytes": 256 * 1024,
+                "timeoutMs": 25
+            }
+        }))
+        .send()
+        .await?;
+
+    debug!("Subscribed to Telemetry API: {:?}", response);
+    Ok(response)
+}