aj/add hello route (#311)

* feat: Add hello route * feat: call it * feat: fmt * fix: clippy * fix: 8124 not 8126 for hello * move telemetry listener temporarily * feat: Fix an issue where dogstatsd port is wrong. Drop enhanced invocation metric from clients because we always set it * fix: fmt
DataDog · Jul 9, 2024 · 5eb5a88 · 5eb5a88
1 parent 1e59f15
commit 5eb5a88
Show file tree

Hide file tree

Showing 5 changed files with 70 additions and 2 deletions.
diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs
@@ -37,6 +37,7 @@ use bottlecap::{
         listener::TelemetryListener,
     },
     traces::{
+        hello_agent,
         stats_flusher::{self, StatsFlusher},
         stats_processor, trace_agent,
         trace_flusher::{self, TraceFlusher},
@@ -313,6 +314,14 @@ async fn extension_loop_active(
             error!("Error starting trace agent: {e:?}");
         }
     });
+    // TODO(astuyve): deprioritize this task after the first request
+    tokio::spawn(async move {
+        let res = hello_agent::start_handler().await;
+        if let Err(e) = res {
+            error!("Error starting hello agent: {e:?}");
+        }
+    });
+
     let lambda_enhanced_metrics =
         enhanced_metrics::new(Arc::clone(&metrics_aggr), Arc::clone(config));
     let dogstatsd_cancel_token = start_dogstatsd(event_bus.get_sender_copy(), &metrics_aggr).await;

diff --git a/bottlecap/src/lib.rs b/bottlecap/src/lib.rs
@@ -41,10 +41,13 @@ pub const EXTENSION_ROUTE: &str = "2020-01-01/extension";
 pub const LAMBDA_RUNTIME_SLUG: &str = "lambda";
 
 // todo: make sure we can override those with environment variables
-pub const DOGSTATSD_PORT: u16 = 8185;
+pub const DOGSTATSD_PORT: u16 = 8125;
 
 pub const TELEMETRY_SUBSCRIPTION_ROUTE: &str = "2022-07-01/telemetry";
-pub const TELEMETRY_PORT: u16 = 8124;
+// 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;
 
 /// Return the base URL for the lambda runtime API
 ///

diff --git a/bottlecap/src/metrics/dogstatsd.rs b/bottlecap/src/metrics/dogstatsd.rs
@@ -65,6 +65,10 @@ impl DogStatsD {
                     continue;
                 }
             };
+            if parsed_metric.name == "aws.lambda.enhanced.invocations" {
+                debug!("dropping invocation metric from layer, as it's set by agent");
+                continue;
+            }
             let first_value = match parsed_metric.first_value() {
                 Ok(val) => val,
                 Err(e) => {
@@ -85,6 +89,7 @@ impl DogStatsD {
             // Don't publish until after validation and adding metric_event to buff
             let _ = self.event_bus.send(Event::Metric(metric_event)).await; // todo check the result
             if self.cancel_token.is_cancelled() {
+                debug!("closing dogstatsd listener");
                 break;
             }
         }

diff --git a/bottlecap/src/traces/hello_agent.rs b/bottlecap/src/traces/hello_agent.rs
@@ -0,0 +1,50 @@
+// Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
+// SPDX-License-Identifier: Apache-2.0
+
+// TODO(Astuyve): Deprecate.
+// older clients require the 127.0.0.1:8126/lambda/hello route
+// to identify the presence of the extension.
+
+use hyper::service::{make_service_fn, service_fn};
+use hyper::{http, Body, Method, Request, Response, Server, StatusCode};
+use serde_json::json;
+use std::convert::Infallible;
+use std::net::SocketAddr;
+use tracing::error;
+
+const HELLO_PATH: &str = "/lambda/hello";
+const AGENT_PORT: usize = 8124;
+
+pub async fn start_handler() -> Result<(), Box<dyn std::error::Error>> {
+    let make_svc = make_service_fn(move |_| {
+        let service = service_fn(hello_handler);
+
+        async move { Ok::<_, Infallible>(service) }
+    });
+
+    let port = u16::try_from(AGENT_PORT).expect("AGENT_PORT is too large");
+    let addr = SocketAddr::from(([127, 0, 0, 1], port));
+    let server_builder = Server::try_bind(&addr)?;
+
+    let server = server_builder.serve(make_svc);
+
+    // start hyper http server
+    if let Err(e) = server.await {
+        error!("Server error: {e}");
+        return Err(e.into());
+    }
+
+    Ok(())
+}
+
+async fn hello_handler(req: Request<Body>) -> http::Result<Response<Body>> {
+    if let (&Method::GET, HELLO_PATH) = (req.method(), req.uri().path()) {
+        Response::builder()
+            .status(200)
+            .body(Body::from(json!({}).to_string()))
+    } else {
+        let mut not_found = Response::default();
+        *not_found.status_mut() = StatusCode::NOT_FOUND;
+        Ok(not_found)
+    }
+}
diff --git a/bottlecap/src/traces/mod.rs b/bottlecap/src/traces/mod.rs
@@ -1,6 +1,7 @@
 // Copyright 2023-Present Datadog, Inc. https://www.datadoghq.com/
 // SPDX-License-Identifier: Apache-2.0
 
+pub mod hello_agent;
 pub mod stats_flusher;
 pub mod stats_processor;
 pub mod trace_agent;