DataDog · duncanista · Aug 27, 2024 · Aug 21, 2024 · Aug 21, 2024 · Aug 22, 2024
diff --git 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,17 +303,21 @@ 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;
         if let Err(e) = res {
             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
@@ -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<provider::Provider>,
+}
+
+impl Listener {
+    pub async fn start(&self) -> Result<(), Box<dyn std::error::Error>> {
+        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<Body>) -> http::Result<Response<Body>> {
+        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<Body>) -> http::Result<Response<Body>> {
+        Response::builder()
+            .status(200)
+            .body(Body::from(json!({}).to_string()))
+    }
+
+    fn end_invocation_handler(_: Request<Body>) -> http::Result<Response<Body>> {
+        Response::builder()
+            .status(200)
+            .body(Body::from(json!({}).to_string()))
+    }
+
+    fn hello_handler() -> http::Result<Response<Body>> {
+        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
@@ -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
diff --git 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;