aj/add hello route (#311)

astuyve · web-flow · commit 5eb5a8868f31 · 2024-07-09T17:12:35.000-04:00
* 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
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;

Original file line number	Diff line number	Diff line change
`@@ -65,6 +65,10 @@ impl DogStatsD {`
`65`	`65`	`continue;`
`66`	`66`	`}`
`67`	`67`	`};`
	`68`	`+ if parsed_metric.name == "aws.lambda.enhanced.invocations" {`
	`69`	`+ debug!("dropping invocation metric from layer, as it's set by agent");`
	`70`	`+ continue;`
	`71`	`+ }`
`68`	`72`	`let first_value = match parsed_metric.first_value() {`
`69`	`73`	`Ok(val) => val,`
`70`	`74`	`Err(e) => {`
`@@ -85,6 +89,7 @@ impl DogStatsD {`
`85`	`89`	`// Don't publish until after validation and adding metric_event to buff`
`86`	`90`	`let _ = self.event_bus.send(Event::Metric(metric_event)).await; // todo check the result`
`87`	`91`	`if self.cancel_token.is_cancelled() {`
	`92`	`+ debug!("closing dogstatsd listener");`
`88`	`93`	`break;`
`89`	`94`	`}`
`90`	`95`	`}`