diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index 13b52cd38..89e97af29 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/bottlecap/src/bin/bottlecap/main.rs @@ -17,6 +17,7 @@ use bottlecap::{ lifecycle::{ flush_control::FlushControl, invocation_context::{InvocationContext, InvocationContextBuffer}, + listener::Listener as LifecycleListener, }, logger, logs::{ @@ -39,7 +40,6 @@ use bottlecap::{ listener::TelemetryListener, }, traces::{ - hello_agent, stats_flusher::{self, StatsFlusher}, stats_processor, trace_agent, trace_flusher::{self, TraceFlusher}, @@ -303,7 +303,7 @@ async fn extension_loop_active( trace_flusher: trace_flusher_clone, stats_processor, stats_flusher: stats_flusher_clone, - tags_provider, + tags_provider: Arc::clone(&tags_provider), }); tokio::spawn(async move { let res = trace_agent.start_trace_agent().await; @@ -311,9 +311,13 @@ async fn extension_loop_active( error!("Error starting trace agent: {e:?}"); } }); + + let lifecycle_listener = LifecycleListener { + tags_provider: Arc::clone(&tags_provider), + }; // TODO(astuyve): deprioritize this task after the first request tokio::spawn(async move { - let res = hello_agent::start_handler().await; + let res = lifecycle_listener.start().await; if let Err(e) = res { error!("Error starting hello agent: {e:?}"); } diff --git a/bottlecap/src/lifecycle/listener.rs b/bottlecap/src/lifecycle/listener.rs new file mode 100644 index 000000000..0afdab41c --- /dev/null +++ b/bottlecap/src/lifecycle/listener.rs @@ -0,0 +1,79 @@ +// Copyright 2024-Present Datadog, Inc. https://www.datadoghq.com/ +// SPDX-License-Identifier: Apache-2.0 + +use std::convert::Infallible; +use std::net::SocketAddr; +use std::sync::Arc; + +use hyper::service::{make_service_fn, service_fn}; +use hyper::{http, Body, Method, Request, Response, StatusCode}; +use serde_json::json; +use tracing::{error, warn}; + +use crate::tags::provider; + +const HELLO_PATH: &str = "/lambda/hello"; +const START_INVOCATION_PATH: &str = "/lambda/start-invocation"; +const END_INVOCATION_PATH: &str = "/lambda/end-invocation"; +const AGENT_PORT: usize = 8124; + +pub struct Listener { + pub tags_provider: Arc, +} + +impl Listener { + pub async fn start(&self) -> Result<(), Box> { + let make_svc = make_service_fn(move |_| { + let service = service_fn(Self::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 = hyper::Server::try_bind(&addr)?; + + let server = server_builder.serve(make_svc); + + // start hyper http server + if let Err(e) = server.await { + error!("Failed to start the Lifecycle Listener {e}"); + return Err(e.into()); + } + + Ok(()) + } + + #[allow(clippy::unused_async)] + async fn handler(req: Request) -> http::Result> { + match (req.method(), req.uri().path()) { + (&Method::POST, START_INVOCATION_PATH) => Self::start_invocation_handler(req), + (&Method::POST, END_INVOCATION_PATH) => Self::end_invocation_handler(req), + (&Method::GET, HELLO_PATH) => Self::hello_handler(), + _ => { + let mut not_found = Response::default(); + *not_found.status_mut() = StatusCode::NOT_FOUND; + Ok(not_found) + } + } + } + + fn start_invocation_handler(_: Request) -> http::Result> { + Response::builder() + .status(200) + .body(Body::from(json!({}).to_string())) + } + + fn end_invocation_handler(_: Request) -> http::Result> { + Response::builder() + .status(200) + .body(Body::from(json!({}).to_string())) + } + + fn hello_handler() -> http::Result> { + warn!("[DEPRECATED] Please upgrade your tracing library, the /hello route is deprecated"); + Response::builder() + .status(200) + .body(Body::from(json!({}).to_string())) + } +} diff --git a/bottlecap/src/lifecycle/mod.rs b/bottlecap/src/lifecycle/mod.rs index 1c0924d84..a35e20622 100644 --- a/bottlecap/src/lifecycle/mod.rs +++ b/bottlecap/src/lifecycle/mod.rs @@ -1,2 +1,3 @@ pub mod flush_control; pub mod invocation_context; +pub mod listener; diff --git a/bottlecap/src/traces/hello_agent.rs b/bottlecap/src/traces/hello_agent.rs deleted file mode 100644 index c3584ccf3..000000000 --- a/bottlecap/src/traces/hello_agent.rs +++ /dev/null @@ -1,51 +0,0 @@ -// 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, warn}; - -const HELLO_PATH: &str = "/lambda/hello"; -const AGENT_PORT: usize = 8124; - -pub async fn start_handler() -> Result<(), Box> { - 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) -> http::Result> { - if let (&Method::GET, HELLO_PATH) = (req.method(), req.uri().path()) { - warn!("[DEPRECATED] Please upgrade your tracing library, the /hello route is deprecated"); - 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 index b70d26a83..8545fbe40 100644 --- a/bottlecap/src/traces/mod.rs +++ b/bottlecap/src/traces/mod.rs @@ -1,7 +1,6 @@ // 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;