[SVLS-10110] fix(traces): bind the trace agent listener before registering

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<AwsConfig>,
     config: &Arc<Config>,
@@ -334,6 +343,7 @@ async fn extension_loop_active(
     r: &RegisterResponse,
     api_key_factory: Arc<ApiKeyFactory>,
     mut cold_start_durations: ColdStartDurations,
+    trace_listener: Option<tokio::net::TcpListener>,
 ) -> 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(

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<Config>,
     api_key_factory: &Arc<ApiKeyFactory>,
@@ -240,6 +244,7 @@ pub fn start_trace_agent(
     appsec_processor: Option<Arc<TokioMutex<AppSecProcessor>>>,
     client: &reqwest::Client,
     proxy_aggregator: Arc<TokioMutex<proxy_aggregator::Aggregator>>,
+    listener: Option<tokio::net::TcpListener>,
 ) -> 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
 }

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<u16>) -> std::io::Result<tokio::net::TcpListener> {
+        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<dyn std::error::Error + Send + Sync>> {
+        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<dyn std::error::Error + Send + Sync>> {
         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

bottlecap/tests/common/fake_extensions_api.rs

@@ -74,6 +74,9 @@ struct Shared {
     telemetry_url: Arc<Mutex<Option<reqwest::Url>>>,
     /// `types` of the last accepted subscribe body.
     telemetry_types: Arc<Mutex<Vec<String>>>,
+    /// Whether the trace agent port accepted a connection while the first
+    /// register request was handled, `None` until then.
+    trace_agent_listening_at_register: Arc<Mutex<Option<bool>>>,
 }
 
 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<bool> {
+        *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<impl IntoResponse, StatusCode> {
     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

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;
+}