[SVLS-10070] fix(lifecycle): remove finished tasks from the listener's JoinSet

bottlecap/src/lifecycle/listener.rs

@@ -140,6 +140,7 @@ impl Listener {
         }
 
         let mut join_set = tasks.lock().await;
+        Self::reap_finished(&mut join_set);
         join_set.spawn(async move {
             Self::universal_instrumentation_start(
                 headers,
@@ -164,6 +165,7 @@ impl Listener {
         let (parts, body) = extract_request_body_or_empty(request).await;
 
         let mut join_set = tasks.lock().await;
+        Self::reap_finished(&mut join_set);
         join_set.spawn(async move {
             Self::universal_instrumentation_end(&parts.headers, body, invocation_processor_handle)
                 .await;
@@ -172,6 +174,16 @@ impl Listener {
         (StatusCode::OK, json!({}).to_string()).into_response()
     }
 
+    /// A `JoinSet` keeps each finished task until it is joined, so remove them before adding
+    /// one. Without this, the set grows by two tasks per invocation for the sandbox's lifetime.
+    fn reap_finished(join_set: &mut JoinSet<()>) {
+        while let Some(result) = join_set.try_join_next() {
+            if let Err(e) = result {
+                error!("Lifecycle API | Task error: {e}");
+            }
+        }
+    }
+
     #[allow(clippy::unused_async)]
     async fn handle_hello() -> Response {
         warn!("[DEPRECATED] Please upgrade your tracing library, the /hello route is deprecated");
@@ -342,6 +354,34 @@ mod tests {
         );
     }
 
+    #[tokio::test]
+    async fn test_finished_tasks_are_removed_from_join_set() {
+        let propagator = Arc::new(DatadogCompositePropagator::new(Arc::new(
+            crate::config::Config::default(),
+        )));
+        let tasks = Arc::new(Mutex::new(JoinSet::new()));
+        let state: ListenerState = (InvocationProcessorHandle::noop(), propagator, tasks.clone());
+        let router = Listener::make_router(state);
+
+        for i in 0..100 {
+            for path in [START_INVOCATION_PATH, END_INVOCATION_PATH] {
+                let req = Request::builder()
+                    .method("POST")
+                    .uri(path)
+                    .header("lambda-runtime-aws-request-id", format!("req-{i}"))
+                    .body(Body::from("{}"))
+                    .expect("failed to build request");
+                let response = router.clone().oneshot(req).await.expect("request failed");
+                assert_eq!(response.status(), StatusCode::OK);
+            }
+            // Let the spawned tasks finish before the next request.
+            tokio::task::yield_now().await;
+        }
+
+        let len = tasks.lock().await.len();
+        assert!(len <= 2, "join set holds {len} tasks after 200 requests");
+    }
+
     #[test]
     fn test_extract_request_id_from_header() {
         let headers = HashMap::from([(