From 5445630843c6b399d834d39842923cb23bf29f12 Mon Sep 17 00:00:00 2001
From: julienboulay <julienboulay@users.noreply.github.com>
Date: Tue, 11 Mar 2025 20:54:36 +0100
Subject: [PATCH] Fix lambda container is dropped by ryuk (#848)

---
 .../src/localstack-container.test.ts          | 14 +++++++++--
 .../localstack/src/localstack-container.ts    | 23 ++++++++++++++++++-
 2 files changed, 34 insertions(+), 3 deletions(-)

diff --git a/packages/modules/localstack/src/localstack-container.test.ts b/packages/modules/localstack/src/localstack-container.test.ts
index c172d3c0f..18d3d26ea 100644
--- a/packages/modules/localstack/src/localstack-container.test.ts
+++ b/packages/modules/localstack/src/localstack-container.test.ts
@@ -1,6 +1,7 @@
-import { LOCALSTACK_PORT, LocalstackContainer } from "./localstack-container";
-import { HeadBucketCommand, S3Client, CreateBucketCommand } from "@aws-sdk/client-s3";
+import { CreateBucketCommand, HeadBucketCommand, S3Client } from "@aws-sdk/client-s3";
 import { GenericContainer, log, Network, StartedTestContainer } from "testcontainers";
+import { LABEL_TESTCONTAINERS_SESSION_ID } from "testcontainers/src/utils/labels";
+import { LocalstackContainer, LOCALSTACK_PORT } from "./localstack-container";
 
 const runAwsCliAgainstDockerNetworkContainer = async (
   command: string,
@@ -103,4 +104,13 @@ describe("LocalStackContainer", () => {
 
     await container.stop();
   });
+
+  it("should add LAMBDA_DOCKER_FLAGS with sessionId label", async () => {
+    const container = await new LocalstackContainer().start();
+    const sessionId = container.getLabels()[LABEL_TESTCONTAINERS_SESSION_ID];
+
+    const { output, exitCode } = await container.exec(["printenv", "LAMBDA_DOCKER_FLAGS"]);
+    expect(exitCode).toBe(0);
+    expect(output).toContain(`${LABEL_TESTCONTAINERS_SESSION_ID}=${sessionId}`);
+  });
 });
diff --git a/packages/modules/localstack/src/localstack-container.ts b/packages/modules/localstack/src/localstack-container.ts
index 0b56a143d..6f73a7a7f 100644
--- a/packages/modules/localstack/src/localstack-container.ts
+++ b/packages/modules/localstack/src/localstack-container.ts
@@ -1,10 +1,20 @@
-import { AbstractStartedContainer, GenericContainer, log, StartedTestContainer, Wait } from "testcontainers";
+import {
+  AbstractStartedContainer,
+  GenericContainer,
+  getContainerRuntimeClient,
+  log,
+  StartedTestContainer,
+  Wait,
+} from "testcontainers";
+import { getReaper } from "testcontainers/src/reaper/reaper";
+import { LABEL_TESTCONTAINERS_SESSION_ID } from "testcontainers/src/utils/labels";
 
 export const LOCALSTACK_PORT = 4566;
 
 export class LocalstackContainer extends GenericContainer {
   constructor(image = "localstack/localstack:2.2.0") {
     super(image);
+
     this.withExposedPorts(LOCALSTACK_PORT).withWaitStrategy(Wait.forLogMessage("Ready", 1)).withStartupTimeout(120_000);
   }
 
@@ -25,8 +35,19 @@ export class LocalstackContainer extends GenericContainer {
     log.info(`${envVar} environment variable set to "${this.environment[envVar]}" (${hostnameExternalReason})`);
   }
 
+  /**
+   * Configure the LocalStack container to add sessionId when starting lambda container.
+   * This allows the lambda container to be identified by the {@link Reaper} and cleaned up on exit.
+   */
+  private async flagLambdaSessionId(): Promise<void> {
+    const client = await getContainerRuntimeClient();
+    const reaper = await getReaper(client);
+    this.withEnvironment({ LAMBDA_DOCKER_FLAGS: `-l ${LABEL_TESTCONTAINERS_SESSION_ID}=${reaper.sessionId}` });
+  }
+
   protected override async beforeContainerCreated(): Promise<void> {
     this.resolveHostname();
+    await this.flagLambdaSessionId();
   }
 
   public override async start(): Promise<StartedLocalStackContainer> {