diff --git a/src/logging/log-patch.test.ts b/src/logging/log-patch.test.ts index dfc9c73..7f06bdd 100644 --- a/src/logging/log-patch.test.ts +++ b/src/logging/log-patch.test.ts @@ -10,6 +10,7 @@ describe("LogPatch", async () => { const originalConsole = { ...console }; const mockTimestamp = "2024-01-01T00:00:00.000Z"; const mockRequestId = "test-request-id"; + const mockTenantId = "test-tenant-id"; const invokeStore = await InvokeStore.getInstanceAsync(); @@ -25,6 +26,7 @@ describe("LogPatch", async () => { // Mock InvokeStore vi.spyOn(invokeStore, "getRequestId").mockReturnValue(mockRequestId); + vi.spyOn(invokeStore, "getTenantId").mockReturnValue(mockTenantId); // Mock Date vi.spyOn(Date.prototype, "toISOString").mockReturnValue(mockTimestamp); @@ -230,6 +232,9 @@ describe("LogPatch", async () => { expect(process.stdout.write).toHaveBeenCalledWith( expect.stringContaining(`"requestId":"${mockRequestId}"`), ); + expect(process.stdout.write).toHaveBeenCalledWith( + expect.stringContaining(`"tenantId":"${mockTenantId}"`), + ); expect(process.stdout.write).toHaveBeenCalledWith( expect.stringContaining(`"level":"INFO"`), ); diff --git a/src/runtime/runtime.test.ts b/src/runtime/runtime.test.ts index 781cb22..3637fcd 100644 --- a/src/runtime/runtime.test.ts +++ b/src/runtime/runtime.test.ts @@ -353,6 +353,30 @@ describe("Runtime", () => { ]); }); + it("should set tenantId in InvokeStore for buffered invocations", async () => { + // GIVEN + const mockRapidClient = createMockRapidClient(-1, false, true); + const storeOrder: (string | undefined)[] = []; + + const BufferedHandler: BufferedHandler = async () => { + storeOrder.push(invokeStore.getTenantId()); + await new Promise((resolve) => setTimeout(resolve, 50)); // Simulate work + return "done"; + }; + + const runtime = Runtime.create({ + rapidClient: mockRapidClient, + handler: BufferedHandler, + isMultiConcurrent: false, + }); + + // WHEN + await startAndTickRuntime(runtime, 60); + + // THEN - tenantId is pinned to the sandbox, constant across invocations + expect(storeOrder).toEqual(["test-tenant", "test-tenant"]); + }); + it("should set requestId and traceId in InvokeStore for streaming invocations", async () => { // GIVEN const mockRapidClient = createMockRapidClient(3, true); @@ -556,6 +580,45 @@ describe("Runtime", () => { expect(mockRapidClient.postInvocationResponse).toHaveBeenCalledTimes(2); }); + it("should set tenantId in InvokeStore for concurrent invocations", async () => { + // GIVEN + const mockRapidClient = createMockRapidClient(3, false, true); + const storeOrder: (string | undefined)[] = []; + const completionOrder: (string | undefined)[] = []; + + const BufferedHandler: BufferedHandler = async () => { + const invocationNumber = storeOrder.length + 1; + + storeOrder.push(invokeStore.getTenantId()); + + await new Promise((resolve) => + setTimeout(resolve, invocationNumber === 1 ? 10 : 5), + ); + + completionOrder.push(invokeStore.getTenantId()); + + return "done"; + }; + + const runtime = Runtime.create({ + rapidClient: mockRapidClient, + handler: BufferedHandler, + isMultiConcurrent: true, + }); + + // WHEN + runtime.start(); + await new Promise((resolve) => setTimeout(resolve, 20)); + + // THEN - tenantId is pinned to the sandbox, so every concurrent + // invocation reads the same value, both at start and after async gaps. + expect(storeOrder).toEqual(["test-tenant", "test-tenant"]); + expect(completionOrder).toEqual(["test-tenant", "test-tenant"]); + + expect(mockRapidClient.nextInvocation).toHaveBeenCalledTimes(3); + expect(mockRapidClient.postInvocationResponse).toHaveBeenCalledTimes(2); + }); + it("should maintain context when mixing different timer types in RIC", async () => { // GIVEN const mockRapidClient = createMockRapidClient(3, true); @@ -897,6 +960,7 @@ async function startAndTickRuntime( function createMockRapidClient( throttleAt: number = -1, passXRayTraceId: boolean = false, + passTenantId: boolean = false, ): RAPIDClient { let invocationCount = 0; return { @@ -915,6 +979,14 @@ function createMockRapidClient( "lambda-runtime-trace-id": `trace-id-${invocationCount}`, } : {}), + ...(passTenantId + ? { + // tenantId is pinned to the sandbox, so it's constant across + // every invocation in the execution environment's lifecycle + // (unlike requestId/traceId, which change per invocation). + "lambda-runtime-aws-tenant-id": "test-tenant", + } + : {}), }, }); }), diff --git a/src/runtime/runtime.ts b/src/runtime/runtime.ts index 379ee8e..419258c 100644 --- a/src/runtime/runtime.ts +++ b/src/runtime/runtime.ts @@ -67,6 +67,7 @@ export class Runtime { await this.runWithInvokeContext( context.awsRequestId, context.xRayTraceId, + context.tenantId, () => processor.processInvoke(context, event), ); } @@ -82,6 +83,7 @@ export class Runtime { await this.runWithInvokeContext( context.awsRequestId, context.xRayTraceId, + context.tenantId, () => processor.processInvoke(context, event), ); }); @@ -91,6 +93,7 @@ export class Runtime { private async runWithInvokeContext( requestId: string, xRayTraceId: string | undefined, + tenantId: string | undefined, fn: () => Promise, ): Promise { const invokeStore = await InvokeStore.getInstanceAsync(); @@ -98,6 +101,7 @@ export class Runtime { { [InvokeStoreBase.PROTECTED_KEYS.REQUEST_ID]: requestId, [InvokeStoreBase.PROTECTED_KEYS.X_RAY_TRACE_ID]: xRayTraceId, + [InvokeStoreBase.PROTECTED_KEYS.TENANT_ID]: tenantId, }, fn, );