diff --git a/bottlecap/src/bin/bottlecap/main.rs b/bottlecap/src/bin/bottlecap/main.rs index 6cb53476a..e000b7a34 100644 --- a/bottlecap/src/bin/bottlecap/main.rs +++ b/bottlecap/src/bin/bottlecap/main.rs @@ -141,6 +141,13 @@ async fn main() -> anyhow::Result<()> { .map_err(|e| anyhow::anyhow!("Failed to create client: {e:?}"))?; cold_start_durations.create_extensions_api_client_ms = elapsed_ms(step_start_time); + // Lambda starts the runtime once every extension has registered, and + // tracers probe `/info` early in their init. Bind the trace agent's port + // before registering, so that probe queues instead of being refused. + let trace_listener = bottlecap::traces::trace_agent::TraceAgent::bind_listener(None) + .inspect_err(|e| error!("Error starting trace agent: {e:?}")) + .ok(); + let cloned_client = client.clone(); let runtime_api = aws_config.runtime_api.clone(); let managed_instance_mode = aws_config.is_managed_instance_mode(); @@ -193,6 +200,7 @@ async fn main() -> anyhow::Result<()> { &r, Arc::clone(&api_key_factory), cold_start_durations, + trace_listener, ) .await { @@ -326,6 +334,7 @@ async fn extension_loop_idle( } #[allow(clippy::too_many_lines)] +#[allow(clippy::too_many_arguments)] async fn extension_loop_active( aws_config: Arc, config: &Arc, @@ -334,6 +343,7 @@ async fn extension_loop_active( r: &RegisterResponse, api_key_factory: Arc, mut cold_start_durations: ColdStartDurations, + trace_listener: Option, ) -> anyhow::Result<()> { let (mut event_bus, event_bus_tx) = EventBus::run(); @@ -462,6 +472,7 @@ async fn extension_loop_active( appsec_processor.clone(), &shared_client, Arc::clone(&proxy_aggregator), + trace_listener, ); let api_runtime_proxy_shutdown_signal = start_api_runtime_proxy( diff --git a/bottlecap/src/startup.rs b/bottlecap/src/startup.rs index cf12b0939..8a19072c5 100644 --- a/bottlecap/src/startup.rs +++ b/bottlecap/src/startup.rs @@ -223,15 +223,19 @@ pub fn build_trace_agent( } /// Builds the trace-processing pipeline with [`build_trace_agent`] and spawns -/// the [`trace_agent::TraceAgent`] HTTP listener onto the current tokio -/// runtime. Convenience entry point for callers that do not need to +/// the [`trace_agent::TraceAgent`] HTTP server on `listener` onto the current +/// tokio runtime. Convenience entry point for callers that do not need to /// further configure the `TraceAgent` before spawning it. /// -/// Errors from `TraceAgent::start` (TCP bind failures, router-extension -/// failures, axum serve errors) are logged and discarded; the surrounding -/// pipeline keeps running with a dead trace channel. Callers that need to -/// react to startup errors should use [`build_trace_agent`] and spawn the -/// agent themselves. +/// `listener` comes from [`trace_agent::TraceAgent::bind_listener`], called +/// before the extension registers. `None` means that bind failed (the caller +/// logged it), and no trace agent is served. +/// +/// Errors from `TraceAgent::start_with_listener` (router-extension failures, +/// axum serve errors) are logged and discarded; the surrounding pipeline keeps +/// running with a dead trace channel. Callers that need to react to startup +/// errors should use [`build_trace_agent`] and spawn the agent themselves. +#[allow(clippy::too_many_arguments)] pub fn start_trace_agent( config: &Arc, api_key_factory: &Arc, @@ -240,6 +244,7 @@ pub fn start_trace_agent( appsec_processor: Option>>, client: &reqwest::Client, proxy_aggregator: Arc>, + listener: Option, ) -> TraceAgentPipeline { let (trace_agent, pipeline) = build_trace_agent( config, @@ -252,11 +257,13 @@ pub fn start_trace_agent( None, ); - tokio::spawn(async move { - if let Err(e) = trace_agent.start().await { - error!("Error starting trace agent: {e:?}"); - } - }); + if let Some(listener) = listener { + tokio::spawn(async move { + if let Err(e) = trace_agent.start_with_listener(listener).await { + error!("Error starting trace agent: {e:?}"); + } + }); + } pipeline } diff --git a/bottlecap/src/traces/trace_agent.rs b/bottlecap/src/traces/trace_agent.rs index bb7bdf6cf..c002ff20f 100644 --- a/bottlecap/src/traces/trace_agent.rs +++ b/bottlecap/src/traces/trace_agent.rs @@ -53,6 +53,8 @@ use crate::traces::stats_concentrator_service::StatsConcentratorHandle; /// Default TCP port the trace agent listener binds to on 127.0.0.1. pub const TRACE_AGENT_PORT: usize = 8126; +/// Listen backlog; the kernel caps it at `net.core.somaxconn`. +const LISTEN_BACKLOG: u32 = 1024; // Agent endpoints const V4_TRACE_ENDPOINT_PATH: &str = "/v0.4/traces"; @@ -115,11 +117,12 @@ pub struct ProxyState { /// deterministic drain endpoint) without adding a dedicated field per /// route to [`TraceAgent`]. /// -/// Returning `Err` propagates out of [`TraceAgent::start`], aborting the -/// HTTP listener task. Note that the production convenience entry point -/// [`crate::startup::start_trace_agent`] spawns `start` and only logs its -/// error; the surrounding pipeline does not observe the failure. Callers -/// that need to react to startup errors must use +/// Returning `Err` propagates out of [`TraceAgent::start_with_listener`] +/// (and [`TraceAgent::start`]), aborting the HTTP listener task. Note that +/// the production convenience entry point +/// [`crate::startup::start_trace_agent`] spawns `start_with_listener` and +/// only logs its error; the surrounding pipeline does not observe the +/// failure. Callers that need to react to startup errors must use /// [`crate::startup::build_trace_agent`] and spawn the agent themselves. /// /// Note that a path collision is not an `Err`: `Router::merge` panics when @@ -367,21 +370,47 @@ impl TraceAgent { self } - #[allow(clippy::cast_possible_truncation)] + /// Binds the trace agent's listening socket on `port`, or + /// [`TRACE_AGENT_PORT`] when `None`. Connections are accepted (and queued) + /// from here on, before the agent serves. + /// + /// The extension calls this before it registers: Lambda starts the runtime + /// once every extension has registered, and tracers probe `/info` early in + /// their init, so the socket must already accept connections by then. + /// + /// Sets `SO_REUSEADDR` like `tokio::net::TcpListener::bind`. Must run + /// inside the tokio runtime. + pub fn bind_listener(port: Option) -> std::io::Result { + let port = + port.unwrap_or(u16::try_from(TRACE_AGENT_PORT).expect("TRACE_AGENT_PORT is too large")); + let socket = tokio::net::TcpSocket::new_v4()?; + socket.set_reuseaddr(true)?; + socket.bind(SocketAddr::from(([127, 0, 0, 1], port)))?; + let listener = socket.listen(LISTEN_BACKLOG)?; + debug!( + "TRACE AGENT | Listening on port {}", + listener.local_addr()?.port() + ); + Ok(listener) + } + pub async fn start(&self) -> Result<(), Box> { + let listener = Self::bind_listener(self.receiver_port)?; + self.start_with_listener(listener).await + } + + /// Serves on a listener from [`TraceAgent::bind_listener`]. + #[allow(clippy::cast_possible_truncation)] + pub async fn start_with_listener( + &self, + listener: tokio::net::TcpListener, + ) -> Result<(), Box> { let now = Instant::now(); let router = self.make_router()?; - let port = self - .receiver_port - .unwrap_or(u16::try_from(TRACE_AGENT_PORT).expect("TRACE_AGENT_PORT is too large")); - let socket = SocketAddr::from(([127, 0, 0, 1], port)); - let listener = tokio::net::TcpListener::bind(&socket).await?; - - debug!("TRACE AGENT | Listening on port {port}"); debug!( - "TRACE AGENT | Time taken to start: {} ms", + "TRACE AGENT | Time taken to build the router: {} ms", now.elapsed().as_millis() ); @@ -1189,6 +1218,15 @@ mod tests { SendDataBuilderInfo::new(builder, size, owned_tags) } + /// Once `bind_listener` returns, clients can connect before the agent serves. + #[tokio::test] + async fn bind_accepts_connections_before_serving() { + let listener = TraceAgent::bind_listener(Some(0)).expect("bind"); + let addr = listener.local_addr().expect("local_addr"); + + std::net::TcpStream::connect(addr).expect("connect before serving"); + } + /// A payload is only handed to the aggregator by a background forwarder /// task, so accepting it does not by itself make it visible to a flush. /// `ingest_barrier` closes that gap: on this single-threaded runtime the diff --git a/bottlecap/tests/common/fake_extensions_api.rs b/bottlecap/tests/common/fake_extensions_api.rs index 1e39d8eb9..4b370f968 100644 --- a/bottlecap/tests/common/fake_extensions_api.rs +++ b/bottlecap/tests/common/fake_extensions_api.rs @@ -74,6 +74,9 @@ struct Shared { telemetry_url: Arc>>, /// `types` of the last accepted subscribe body. telemetry_types: Arc>>, + /// Whether the trace agent port accepted a connection while the first + /// register request was handled, `None` until then. + trace_agent_listening_at_register: Arc>>, } pub struct FakeExtensionsApi { @@ -121,6 +124,7 @@ impl FakeExtensionsApi { subscribe_status: Arc::new(AtomicU16::new(StatusCode::OK.as_u16())), telemetry_url: Arc::new(Mutex::new(None)), telemetry_types: Arc::new(Mutex::new(Vec::new())), + trace_agent_listening_at_register: Arc::new(Mutex::new(None)), }; let router = Router::new() .route("/2020-01-01/extension/register", post(register)) @@ -239,6 +243,16 @@ impl FakeExtensionsApi { self.shared.requests.lock().expect("requests lock").clone() } + /// Whether the trace agent port accepted a connection while the first + /// register request was handled, `None` before any register. + pub fn trace_agent_listening_at_register(&self) -> Option { + *self + .shared + .trace_agent_listening_at_register + .lock() + .expect("trace agent probe lock") + } + // --- Telemetry API sender ---------------------------------------------- /// Port the extension asked Lambda to deliver telemetry to, `None` until @@ -345,6 +359,21 @@ async fn register( body: Bytes, ) -> Result { record(&shared, "POST", "/2020-01-01/extension/register"); + // Lambda starts the runtime, and with it the tracer's early `/info` + // probe, once every extension has registered. Probe the trace agent + // once, before answering, the way that tracer would. + { + let mut at_register = shared + .trace_agent_listening_at_register + .lock() + .expect("trace agent probe lock"); + if at_register.is_none() { + let addr = SocketAddr::from(([127, 0, 0, 1], super::port_guard::TRACE_AGENT_PORT)); + let listening = + std::net::TcpStream::connect_timeout(&addr, Duration::from_millis(200)).is_ok(); + *at_register = Some(listening); + } + } // Lambda answers `403` without a name. Any name but the file under // `/opt/extensions` registers an internal extension, which cannot take SHUTDOWN. if headers diff --git a/bottlecap/tests/lifecycle_contract_test.rs b/bottlecap/tests/lifecycle_contract_test.rs index 1882b3f1c..f955edbf4 100644 --- a/bottlecap/tests/lifecycle_contract_test.rs +++ b/bottlecap/tests/lifecycle_contract_test.rs @@ -89,3 +89,46 @@ async fn invoke_then_shutdown_exits_cleanly() { s.intake.summary() ); } + +/// Lambda starts the runtime, and the tracer's first `/info` call, once every +/// extension has registered, not at the extension's first `/next`. The fake +/// probes the trace agent port while it handles register, so the trace agent +/// must already be listening by then. +async fn assert_trace_agent_listens_before_register(managed_instance: bool) { + let s = Scenario::start(FaultMode::Ok).await; + let builder = s.bottlecap(); + let mut bc = if managed_instance { + builder.managed_instance().spawn() + } else { + builder.spawn() + }; + + assert!( + s.api.wait_for_next_calls(1, bc.startup_deadline()).await, + "no /next poll. {}", + bc.log_tail(40) + ); + assert_eq!( + s.api.trace_agent_listening_at_register(), + Some(true), + "the trace agent was not listening on 8126 when the extension registered. {}", + bc.log_tail(40) + ); + // Only the startup order matters here. The process is killed rather than + // shut down because the On-Demand loop does not yet act on a SHUTDOWN that + // answers the first `/next`. + bc.assert_no_panic(); + bc.kill(); +} + +#[tokio::test(flavor = "multi_thread")] +#[serial_test::serial(bottlecap_ports)] +async fn trace_agent_listens_before_register() { + assert_trace_agent_listens_before_register(false).await; +} + +#[tokio::test(flavor = "multi_thread")] +#[serial_test::serial(bottlecap_ports)] +async fn trace_agent_listens_before_register_managed_instance() { + assert_trace_agent_listens_before_register(true).await; +}