diff --git a/package-lock.json b/package-lock.json index c564484c37..caa70c2401 100644 --- a/package-lock.json +++ b/package-lock.json @@ -8422,13 +8422,13 @@ "license": "MIT-0", "dependencies": { "@aws-lambda-powertools/commons": "2.35.0", - "@aws/lambda-invoke-store": "0.3.0", "@standard-schema/spec": "^1.1.0" }, "devDependencies": { "@aws-lambda-powertools/parser": "2.35.0", "@aws-lambda-powertools/testing-utils": "file:../testing", "@aws-sdk/client-sqs": "^3.1121.0", + "@aws/lambda-invoke-store": "0.3.0", "zod": "^4.5.4" } }, @@ -8593,12 +8593,12 @@ "version": "2.35.0", "license": "MIT-0", "dependencies": { - "@aws-lambda-powertools/commons": "2.35.0", - "@aws/lambda-invoke-store": "0.3.0" + "@aws-lambda-powertools/commons": "2.35.0" }, "devDependencies": { "@aws-lambda-powertools/testing-utils": "file:../testing", "@aws-sdk/client-lambda": "^3.1121.0", + "@aws/lambda-invoke-store": "0.3.0", "@types/promise-retry": "^1.1.3", "promise-retry": "^2.0.1" }, @@ -8620,12 +8620,12 @@ "version": "2.35.0", "license": "MIT-0", "dependencies": { - "@aws-lambda-powertools/commons": "2.35.0", - "@aws/lambda-invoke-store": "0.3.0" + "@aws-lambda-powertools/commons": "2.35.0" }, "devDependencies": { "@aws-lambda-powertools/testing-utils": "file:../testing", "@aws-sdk/client-cloudwatch": "^3.1121.0", + "@aws/lambda-invoke-store": "0.3.0", "@types/promise-retry": "^1.1.3", "promise-retry": "^2.0.1" }, diff --git a/packages/batch/package.json b/packages/batch/package.json index cd49dd2cee..65adb055e5 100644 --- a/packages/batch/package.json +++ b/packages/batch/package.json @@ -83,7 +83,6 @@ "nodejs" ], "dependencies": { - "@aws/lambda-invoke-store": "0.3.0", "@aws-lambda-powertools/commons": "2.35.0", "@standard-schema/spec": "^1.1.0" }, @@ -91,6 +90,7 @@ "@aws-lambda-powertools/parser": "2.35.0", "@aws-lambda-powertools/testing-utils": "file:../testing", "@aws-sdk/client-sqs": "^3.1121.0", + "@aws/lambda-invoke-store": "0.3.0", "zod": "^4.5.4" } } diff --git a/packages/batch/src/BatchProcessingStore.ts b/packages/batch/src/BatchProcessingStore.ts index 6081661f54..32f0de19d0 100644 --- a/packages/batch/src/BatchProcessingStore.ts +++ b/packages/batch/src/BatchProcessingStore.ts @@ -1,5 +1,4 @@ -import '@aws/lambda-invoke-store'; -import { shouldUseInvokeStore } from '@aws-lambda-powertools/commons/utils/env'; +import { InvocationScoped } from '@aws-lambda-powertools/commons/utils/invocation-scoped'; import type { BaseRecord, BatchProcessingOptions, @@ -10,178 +9,96 @@ import type { /** * Manages storage of batch processing state with automatic context detection. * - * This class abstracts the storage mechanism for batch processing state, - * automatically choosing between InvokeStore (when in Lambda context) and - * fallback instance variables (when outside Lambda context). The decision is - * made at runtime on every method call to support Lambda's concurrent execution - * isolation. + * This class abstracts the storage mechanism for state that must be scoped to + * a single invocation, automatically choosing between the InvokeStore (when + * invocations run concurrently in the same execution environment) and an + * instance-level fallback shared across sequential invocations. The decision + * is made at runtime on every access to support Lambda's transition to async + * contexts. */ class BatchProcessingStore { - readonly #recordsKey = Symbol('powertools.batch.records'); - readonly #handlerKey = Symbol('powertools.batch.handler'); - readonly #optionsKey = Symbol('powertools.batch.options'); - readonly #failureMessagesKey = Symbol('powertools.batch.failureMessages'); - readonly #successMessagesKey = Symbol('powertools.batch.successMessages'); - readonly #batchResponseKey = Symbol('powertools.batch.batchResponse'); - readonly #errorsKey = Symbol('powertools.batch.errors'); - - #fallbackRecords: BaseRecord[] = []; - #fallbackHandler: CallableFunction = () => {}; - #fallbackOptions?: BatchProcessingOptions; - #fallbackFailureMessages: EventSourceDataClassTypes[] = []; - #fallbackSuccessMessages: EventSourceDataClassTypes[] = []; - #fallbackBatchResponse: PartialItemFailureResponse = { - batchItemFailures: [], - }; - #fallbackErrors: Error[] = []; + readonly #records = new InvocationScoped( + 'powertools.batch.records', + { fresh: () => [] } + ); + readonly #handler = new InvocationScoped( + 'powertools.batch.handler', + { fresh: () => () => {} } + ); + readonly #options = new InvocationScoped( + 'powertools.batch.options', + { initial: undefined } + ); + readonly #failureMessages = new InvocationScoped( + 'powertools.batch.failureMessages', + { fresh: () => [] } + ); + readonly #successMessages = new InvocationScoped( + 'powertools.batch.successMessages', + { fresh: () => [] } + ); + readonly #batchResponse = new InvocationScoped( + 'powertools.batch.batchResponse', + { fresh: () => ({ batchItemFailures: [] }) } + ); + readonly #errors = new InvocationScoped('powertools.batch.errors', { + fresh: () => [], + }); public getRecords(): BaseRecord[] { - if (!shouldUseInvokeStore()) { - return this.#fallbackRecords; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - return (store.get(this.#recordsKey) as BaseRecord[]) ?? []; + return this.#records.get(); } public setRecords(records: BaseRecord[]): void { - if (!shouldUseInvokeStore()) { - this.#fallbackRecords = records; - return; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - store.set(this.#recordsKey, records); + this.#records.set(records); } public getHandler(): CallableFunction { - if (!shouldUseInvokeStore()) { - return this.#fallbackHandler; - } - - return ( - (globalThis.awslambda?.InvokeStore?.get( - this.#handlerKey - ) as CallableFunction) ?? (() => {}) - ); + return this.#handler.get(); } public setHandler(handler: CallableFunction): void { - if (!shouldUseInvokeStore()) { - this.#fallbackHandler = handler; - return; - } - - globalThis.awslambda?.InvokeStore?.set(this.#handlerKey, handler); + this.#handler.set(handler); } public getOptions(): BatchProcessingOptions | undefined { - if (!shouldUseInvokeStore()) { - return this.#fallbackOptions; - } - - return globalThis.awslambda?.InvokeStore?.get(this.#optionsKey) as - | BatchProcessingOptions - | undefined; + return this.#options.get(); } public setOptions(options: BatchProcessingOptions | undefined): void { - if (!shouldUseInvokeStore()) { - this.#fallbackOptions = options; - return; - } - - globalThis.awslambda?.InvokeStore?.set(this.#optionsKey, options); + this.#options.set(options); } public getFailureMessages(): EventSourceDataClassTypes[] { - if (!shouldUseInvokeStore()) { - return this.#fallbackFailureMessages; - } - - return ( - (globalThis.awslambda?.InvokeStore?.get( - this.#failureMessagesKey - ) as EventSourceDataClassTypes[]) ?? [] - ); + return this.#failureMessages.get(); } public setFailureMessages(messages: EventSourceDataClassTypes[]): void { - if (!shouldUseInvokeStore()) { - this.#fallbackFailureMessages = messages; - return; - } - - globalThis.awslambda?.InvokeStore?.set(this.#failureMessagesKey, messages); + this.#failureMessages.set(messages); } public getSuccessMessages(): EventSourceDataClassTypes[] { - if (!shouldUseInvokeStore()) { - return this.#fallbackSuccessMessages; - } - - return ( - (globalThis.awslambda?.InvokeStore?.get( - this.#successMessagesKey - ) as EventSourceDataClassTypes[]) ?? [] - ); + return this.#successMessages.get(); } public setSuccessMessages(messages: EventSourceDataClassTypes[]): void { - if (!shouldUseInvokeStore()) { - this.#fallbackSuccessMessages = messages; - return; - } - - globalThis.awslambda?.InvokeStore?.set(this.#successMessagesKey, messages); + this.#successMessages.set(messages); } public getBatchResponse(): PartialItemFailureResponse { - if (!shouldUseInvokeStore()) { - return this.#fallbackBatchResponse; - } - - return ( - (globalThis.awslambda?.InvokeStore?.get( - this.#batchResponseKey - ) as PartialItemFailureResponse) ?? { batchItemFailures: [] } - ); + return this.#batchResponse.get(); } public setBatchResponse(response: PartialItemFailureResponse): void { - if (!shouldUseInvokeStore()) { - this.#fallbackBatchResponse = response; - return; - } - - globalThis.awslambda?.InvokeStore?.set(this.#batchResponseKey, response); + this.#batchResponse.set(response); } public getErrors(): Error[] { - if (!shouldUseInvokeStore()) { - return this.#fallbackErrors; - } - - return ( - (globalThis.awslambda?.InvokeStore?.get(this.#errorsKey) as Error[]) ?? [] - ); + return this.#errors.get(); } public setErrors(errors: Error[]): void { - if (!shouldUseInvokeStore()) { - this.#fallbackErrors = errors; - return; - } - - globalThis.awslambda?.InvokeStore?.set(this.#errorsKey, errors); + this.#errors.set(errors); } } diff --git a/packages/batch/src/SqsFifoProcessorStore.ts b/packages/batch/src/SqsFifoProcessorStore.ts index c56e42c5ff..1218d3b337 100644 --- a/packages/batch/src/SqsFifoProcessorStore.ts +++ b/packages/batch/src/SqsFifoProcessorStore.ts @@ -1,51 +1,31 @@ -import '@aws/lambda-invoke-store'; -import { shouldUseInvokeStore } from '@aws-lambda-powertools/commons/utils/env'; +import { InvocationScoped } from '@aws-lambda-powertools/commons/utils/invocation-scoped'; /** * Manages storage of SQS FIFO processor state with automatic context detection. * - * This class abstracts the storage mechanism for SQS FIFO processing state, - * automatically choosing between InvokeStore (when in Lambda context) and - * fallback instance variables (when outside Lambda context). The decision is - * made at runtime on every method call to support Lambda's concurrent execution - * isolation. + * This class abstracts the storage mechanism for state that must be scoped to + * a single invocation, automatically choosing between the InvokeStore (when + * invocations run concurrently in the same execution environment) and an + * instance-level fallback shared across sequential invocations. The decision + * is made at runtime on every access to support Lambda's transition to async + * contexts. */ class SqsFifoProcessorStore { - readonly #currentGroupIdKey = Symbol( - 'powertools.batch.sqsFifo.currentGroupId' + readonly #currentGroupId = new InvocationScoped( + 'powertools.batch.sqsFifo.currentGroupId', + { initial: undefined } ); - readonly #failedGroupIdsKey = Symbol( - 'powertools.batch.sqsFifo.failedGroupIds' + readonly #failedGroupIds = new InvocationScoped>( + 'powertools.batch.sqsFifo.failedGroupIds', + { fresh: () => new Set() } ); - #fallbackCurrentGroupId?: string; - #fallbackFailedGroupIds = new Set(); - public getCurrentGroupId(): string | undefined { - if (!shouldUseInvokeStore()) { - return this.#fallbackCurrentGroupId; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - return store.get(this.#currentGroupIdKey) as string | undefined; + return this.#currentGroupId.get(); } public setCurrentGroupId(groupId: string | undefined): void { - if (!shouldUseInvokeStore()) { - this.#fallbackCurrentGroupId = groupId; - return; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - store.set(this.#currentGroupIdKey, groupId); + this.#currentGroupId.set(groupId); } public addFailedGroupId(groupId: string): void { @@ -57,38 +37,11 @@ class SqsFifoProcessorStore { } public getFailedGroupIds(): Set { - if (!shouldUseInvokeStore()) { - return this.#fallbackFailedGroupIds; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - let failedGroupIds = store.get(this.#failedGroupIdsKey) as - | Set - | undefined; - if (failedGroupIds == null) { - failedGroupIds = new Set(); - store.set(this.#failedGroupIdsKey, failedGroupIds); - } - - return failedGroupIds; + return this.#failedGroupIds.get(); } public clearFailedGroupIds(): void { - if (!shouldUseInvokeStore()) { - this.#fallbackFailedGroupIds = new Set(); - return; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - store.set(this.#failedGroupIdsKey, new Set()); + this.#failedGroupIds.reset(); } } diff --git a/packages/batch/tests/unit/concurrency/BatchProcessingStore.test.ts b/packages/batch/tests/unit/concurrency/BatchProcessingStore.test.ts index 331db6d8be..6c080b768a 100644 --- a/packages/batch/tests/unit/concurrency/BatchProcessingStore.test.ts +++ b/packages/batch/tests/unit/concurrency/BatchProcessingStore.test.ts @@ -236,6 +236,26 @@ describe('BatchProcessingStore concurrent invocation isolation', () => { } ); + it('shadows shared options with undefined in an invocation context', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const invokeStore = await InvokeStore.getInstanceAsync(); + const store = new BatchProcessingStore(); + const sharedOptions = { throwOnFullBatchFailure: false }; + store.setOptions(sharedOptions); + + // Act + const scopedOptions = invokeStore.run({}, () => { + store.setOptions(undefined); + + return store.getOptions(); + }); + + // Assess + expect(scopedOptions).toBeUndefined(); + expect(store.getOptions()).toBe(sharedOptions); + }); + describe('InvokeStore error handling', () => { beforeEach(() => { vi.stubGlobal('awslambda', undefined); @@ -245,26 +265,56 @@ describe('BatchProcessingStore concurrent invocation isolation', () => { vi.unstubAllGlobals(); }); - it('throws error when getRecords is called with AWS_LAMBDA_MAX_CONCURRENCY set but InvokeStore is not available', () => { - // Prepare - vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); - const store = new BatchProcessingStore(); - - // Act & Assess - expect(() => { - store.getRecords(); - }).toThrow('InvokeStore is not available'); - }); - - it('throws error when setRecords is called with AWS_LAMBDA_MAX_CONCURRENCY set but InvokeStore is not available', () => { - // Prepare - vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); - const store = new BatchProcessingStore(); + it.each([ + ['getRecords', (store: BatchProcessingStore) => store.getRecords()], + ['setRecords', (store: BatchProcessingStore) => store.setRecords([])], + ['getHandler', (store: BatchProcessingStore) => store.getHandler()], + [ + 'setHandler', + (store: BatchProcessingStore) => store.setHandler(() => {}), + ], + ['getOptions', (store: BatchProcessingStore) => store.getOptions()], + [ + 'setOptions', + (store: BatchProcessingStore) => store.setOptions(undefined), + ], + [ + 'getFailureMessages', + (store: BatchProcessingStore) => store.getFailureMessages(), + ], + [ + 'setFailureMessages', + (store: BatchProcessingStore) => store.setFailureMessages([]), + ], + [ + 'getSuccessMessages', + (store: BatchProcessingStore) => store.getSuccessMessages(), + ], + [ + 'setSuccessMessages', + (store: BatchProcessingStore) => store.setSuccessMessages([]), + ], + [ + 'getBatchResponse', + (store: BatchProcessingStore) => store.getBatchResponse(), + ], + [ + 'setBatchResponse', + (store: BatchProcessingStore) => + store.setBatchResponse({ batchItemFailures: [] }), + ], + ['getErrors', (store: BatchProcessingStore) => store.getErrors()], + ['setErrors', (store: BatchProcessingStore) => store.setErrors([])], + ])( + 'throws error when %s is called with InvokeStore unavailable', + (_method, action) => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = new BatchProcessingStore(); - // Act & Assess - expect(() => { - store.setRecords([]); - }).toThrow('InvokeStore is not available'); - }); + // Act & Assess + expect(() => action(store)).toThrow('InvokeStore is not available'); + } + ); }); }); diff --git a/packages/batch/tests/unit/concurrency/BatchProcessor.test.ts b/packages/batch/tests/unit/concurrency/BatchProcessor.test.ts index f9d10a3ec4..d2568a282f 100644 --- a/packages/batch/tests/unit/concurrency/BatchProcessor.test.ts +++ b/packages/batch/tests/unit/concurrency/BatchProcessor.test.ts @@ -32,6 +32,7 @@ describe('BatchProcessor concurrent invocation isolation', () => { // Prepare if (useInvokeStore) { vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); } const processor = new BatchProcessor(EventType.SQS); const recordsA = [sqsRecordFactory('record-A')]; @@ -97,6 +98,7 @@ describe('BatchProcessor concurrent invocation isolation', () => { // Prepare if (useInvokeStore) { vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); } const processor = new BatchProcessor(EventType.SQS); const recordsA = [sqsRecordFactory('record-A')]; @@ -152,6 +154,7 @@ describe('BatchProcessor concurrent invocation isolation', () => { // Prepare if (useInvokeStore) { vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); } const processor = new BatchProcessor(EventType.SQS); const recordsA = [sqsRecordFactory('fail')]; @@ -225,6 +228,7 @@ describe('BatchProcessor concurrent invocation isolation', () => { // Prepare if (useInvokeStore) { vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); } const processor = new BatchProcessor(EventType.SQS); const recordsA = [sqsRecordFactory('fail-1'), sqsRecordFactory('fail-2')]; diff --git a/packages/batch/tests/unit/concurrency/SqsFifoPartialProcessor.test.ts b/packages/batch/tests/unit/concurrency/SqsFifoPartialProcessor.test.ts index 08e3d9e844..52a40ebc47 100644 --- a/packages/batch/tests/unit/concurrency/SqsFifoPartialProcessor.test.ts +++ b/packages/batch/tests/unit/concurrency/SqsFifoPartialProcessor.test.ts @@ -74,6 +74,7 @@ describe('SQS FIFO Processors concurrent invocation isolation', () => { // Prepare if (useInvokeStore) { vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); } const processor = new processorClass(); const recordsA = [sqsRecordFactory('record-A', '1')]; @@ -150,6 +151,7 @@ describe('SQS FIFO Processors concurrent invocation isolation', () => { // Prepare if (useInvokeStore) { vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); } const processor = new processorClass(); const recordsA = [sqsRecordFactory('record-A', '1')]; @@ -227,6 +229,7 @@ describe('SQS FIFO Processors concurrent invocation isolation', () => { // Prepare if (useInvokeStore) { vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); } const processor = new processorClass(); const recordsA = [sqsRecordFactory('body-A-2', '1')]; @@ -308,6 +311,7 @@ describe('SQS FIFO Processors concurrent invocation isolation', () => { // Prepare if (useInvokeStore) { vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); } const processor = new processorClass(); const recordsA = [ diff --git a/packages/commons/package.json b/packages/commons/package.json index 5f9f755533..ec9a19db7e 100644 --- a/packages/commons/package.json +++ b/packages/commons/package.json @@ -77,6 +77,10 @@ "import": "./lib/esm/LRUCache.js", "require": "./lib/cjs/LRUCache.js" }, + "./utils/invocation-scoped": { + "import": "./lib/esm/InvocationScoped.js", + "require": "./lib/cjs/InvocationScoped.js" + }, "./utils/unmarshallDynamoDB": { "import": "./lib/esm/unmarshallDynamoDB.js", "require": "./lib/cjs/unmarshallDynamoDB.js" @@ -112,6 +116,10 @@ "lib/cjs/LRUCache.d.ts", "lib/esm/LRUCache.d.ts" ], + "utils/invocation-scoped": [ + "lib/cjs/InvocationScoped.d.ts", + "lib/esm/InvocationScoped.d.ts" + ], "utils/unmarshallDynamoDB": [ "lib/cjs/unmarshallDynamoDB.d.ts", "lib/esm/unmarshallDynamoDB.d.ts" diff --git a/packages/commons/src/InvocationScoped.ts b/packages/commons/src/InvocationScoped.ts new file mode 100644 index 0000000000..d140626bef --- /dev/null +++ b/packages/commons/src/InvocationScoped.ts @@ -0,0 +1,207 @@ +import type { InvokeStoreBase } from '@aws/lambda-invoke-store'; +import '@aws/lambda-invoke-store'; +import { shouldUseInvokeStore } from './envUtils.js'; +import type { InvocationScopedOptions } from './types/InvocationScoped.js'; + +/** + * A value wrapped in a single-element array, so that a slot holding no value at + * all is told apart from one holding a value that is itself `undefined`. + * + * Neither the InvokeStore nor a class field can express "empty" other than as + * `undefined`, which is itself a legitimate value to store: a Lambda context + * that has never been set has to read differently from one deliberately + * cleared. + * + * @typeParam T - The type of the value being wrapped + */ +type Box = [T]; + +/** + * Holds a value scoped to a single invocation when invocations may run + * concurrently in the same execution environment, and shared across + * invocations otherwise. + * + * The storage is resolved on every access rather than once at construction, + * so a value created during initialization keeps working once an invocation + * context exists. Reads and writes are scoped to the invocation only when the + * `AWS_LAMBDA_MAX_CONCURRENCY` environment variable is set, the RIC exposes an + * `InvokeStore`, and an invocation context is active; in every other case they + * fall through to the shared value, which is also what + * {@link InvocationScoped.getShared | `getShared()`} and + * {@link InvocationScoped.setShared | `setShared()`} always operate on. + * + * Values are handed out by reference: a caller that retains the value can + * mutate it, so containers meant to stay per-invocation should not be passed + * outside the invocation that owns them. + * + * @example + * ```ts + * import { InvocationScoped } from '@aws-lambda-powertools/commons/utils/invocation-scoped'; + * + * // A value with a shared default that invocations can override + * const logLevel = new InvocationScoped('powertools.logger.logLevel', { + * initial: 12, + * }); + * logLevel.set(8); // only this invocation logs at DEBUG + * + * // A container created anew for each invocation + * const attributes = new InvocationScoped>( + * 'powertools.logger.temporaryAttributes', + * { fresh: () => ({}) } + * ); + * attributes.get().foo = 'bar'; + * ``` + * + * @typeParam T - The type of the value being held + * + * @internal + */ +class InvocationScoped { + /** + * Key the boxed value is stored under in the InvokeStore, unique to this + * instance. + */ + readonly #key: symbol; + /** + * Creates the value, either by calling the `fresh` factory or by handing back + * the `initial` value. + */ + readonly #createValue: () => T; + /** + * Whether each invocation gets its own value created for it (`fresh`) instead + * of falling back to the shared one (`initial`). + */ + readonly #createsPerInvocation: boolean; + /** + * The value shared across invocations, boxed like the invocation-scoped one. + */ + #sharedValue: Box | undefined; + + /** + * @param name - Description of the value, used for the storage key, e.g. `powertools.logger.logLevel` + * @param {Object} options - Options to create the value + * @param options.initial - Value shared across invocations until one sets its own + * @param options.fresh - Creates the value, called once per invocation on first access, and once for the shared value + */ + public constructor(name: string, options: InvocationScopedOptions) { + this.#key = Symbol(name); + if ('fresh' in options) { + this.#createValue = options.fresh; + this.#createsPerInvocation = true; + } else { + const { initial } = options; + this.#createValue = () => initial; + this.#createsPerInvocation = false; + } + } + + /** + * Whether reads and writes are currently scoped to a single invocation. + */ + public get isScoped(): boolean { + return this.#getScopedStore() !== undefined; + } + + /** + * Gets the invocation-scoped value when there is one, otherwise the shared + * value. + * + * A `fresh` value is created on first access, and the same one is handed back + * for the rest of the invocation. + */ + public get(): T { + const store = this.#getScopedStore(); + if (store === undefined) { + return this.getShared(); + } + + const scopedValue = store.get>(this.#key); + if (scopedValue !== undefined) { + return scopedValue[0]; + } + if (!this.#createsPerInvocation) { + return this.getShared(); + } + + const freshValue = this.#createValue(); + store.set(this.#key, [freshValue]); + + return freshValue; + } + + /** + * Sets the value for the current invocation, or the shared value when no + * invocation context is active, i.e. during initialization. + * + * `undefined` is a value like any other: it shadows the shared value until + * {@link InvocationScoped.reset | `reset()`} discards it. + * + * @param value - The value to store + */ + public set(value: T): void { + const store = this.#getScopedStore(); + if (store === undefined) { + this.setShared(value); + return; + } + store.set(this.#key, [value]); + } + + /** + * Discards the value, so that a `fresh` one is created on the next access and + * an `initial` one goes back to the shared value. + * + * Outside an invocation context the shared value is discarded instead, which + * puts an `initial` value back to the one given at construction. + */ + public reset(): void { + const store = this.#getScopedStore(); + if (store === undefined) { + this.#sharedValue = undefined; + return; + } + // The InvokeStore has no way to remove a key, so an empty slot is + // represented by the absence of a box, which reads fall back from. + store.set(this.#key, undefined); + } + + /** + * Gets the value shared across invocations, ignoring any invocation context. + * + * It is created on first access, whether from here or from + * {@link InvocationScoped.get | `get()`}. + */ + public getShared(): T { + this.#sharedValue ??= [this.#createValue()]; + + return this.#sharedValue[0]; + } + + /** + * Sets the value shared across invocations, ignoring any invocation context. + * + * @param value - The value to store + */ + public setShared(value: T): void { + this.#sharedValue = [value]; + } + + /** + * Gets the InvokeStore when reads and writes should be scoped to the current + * invocation, or `undefined` when they should go to the shared value. + */ + #getScopedStore(): InvokeStoreBase | undefined { + if (!shouldUseInvokeStore()) { + return undefined; + } + + const store = globalThis.awslambda?.InvokeStore; + if (store === undefined) { + throw new Error('InvokeStore is not available'); + } + + return store.hasContext() ? store : undefined; + } +} + +export { InvocationScoped }; diff --git a/packages/commons/src/envUtils.ts b/packages/commons/src/envUtils.ts index fa9ac3baa1..46f4b8f0bf 100644 --- a/packages/commons/src/envUtils.ts +++ b/packages/commons/src/envUtils.ts @@ -324,6 +324,15 @@ const isRequestXRaySampled = (): boolean => { return xRayTraceData?.Sampled === '1'; }; +/** + * Determine whether invocations may run concurrently in the same execution + * environment, and state that must not leak between them should be held in the + * Lambda `InvokeStore` rather than on an instance. + * + * Concurrency is opted into by setting the `AWS_LAMBDA_MAX_CONCURRENCY` + * environment variable, which the runtime also uses to decide whether to back + * the `InvokeStore` with an `AsyncLocalStorage`. + */ const shouldUseInvokeStore = (): boolean => { const res = getStringFromEnv({ key: AWS_LAMBDA_MAX_CONCURRENCY, diff --git a/packages/commons/src/types/InvocationScoped.ts b/packages/commons/src/types/InvocationScoped.ts new file mode 100644 index 0000000000..4ed8447191 --- /dev/null +++ b/packages/commons/src/types/InvocationScoped.ts @@ -0,0 +1,27 @@ +/** + * Options to create an {@link InvocationScoped | `InvocationScoped`} value. + * + * Exactly one of the two shapes must be provided: `initial` for a value that + * falls back to a shared default until an invocation sets its own, `fresh` for + * a container that must be created anew for each invocation. + * + * @typeParam T - The type of the value being held + * + * @internal + */ +type InvocationScopedOptions = + | { + /** + * Value shared across invocations until one sets its own. + */ + initial: T; + } + | { + /** + * Creates the value, called once per invocation on first access, and once + * for the shared value. + */ + fresh: () => T; + }; + +export type { InvocationScopedOptions }; diff --git a/packages/commons/tests/unit/InvocationScoped.test.ts b/packages/commons/tests/unit/InvocationScoped.test.ts new file mode 100644 index 0000000000..69f106c6e9 --- /dev/null +++ b/packages/commons/tests/unit/InvocationScoped.test.ts @@ -0,0 +1,463 @@ +import { InvokeStore } from '@aws/lambda-invoke-store'; +import { sequence } from '@aws-lambda-powertools/testing-utils'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { InvocationScoped } from '../../src/InvocationScoped.js'; + +describe('Class: InvocationScoped', () => { + beforeEach(() => { + InvokeStore._testing?.reset(); + }); + + afterEach(() => { + vi.unstubAllEnvs(); + vi.unstubAllGlobals(); + }); + + describe('without concurrency', () => { + it('returns the initial value until one is set', () => { + // Prepare + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act & Assess + expect(cell.get()).toBe(12); + expect(cell.isScoped).toBe(false); + }); + + it('creates the value only once when the cell is fresh', () => { + // Prepare + const fresh = vi.fn(() => ({ count: 0 })); + const cell = new InvocationScoped('powertools.test.attributes', { + fresh, + }); + + // Act + const first = cell.get(); + const second = cell.get(); + + // Assess + expect(fresh).toHaveBeenCalledTimes(1); + expect(second).toBe(first); + }); + + it('sets the value shared across invocations', () => { + // Prepare + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act + cell.set(8); + + // Assess + expect(cell.get()).toBe(8); + expect(cell.getShared()).toBe(8); + }); + + it('restores the initial value when reset', () => { + // Prepare + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + cell.set(8); + + // Act + cell.reset(); + + // Assess + expect(cell.get()).toBe(12); + }); + + it('creates a new value when a fresh cell is reset', () => { + // Prepare + const cell = new InvocationScoped('powertools.test.attributes', { + fresh: () => ({ count: 0 }), + }); + const first = cell.get(); + + // Act + cell.reset(); + + // Assess + expect(cell.get()).not.toBe(first); + }); + + it('holds a value that is itself undefined', () => { + // Prepare + const cell = new InvocationScoped( + 'powertools.test.level', + { initial: 12 } + ); + + // Act + cell.set(undefined); + + // Assess + expect(cell.get()).toBeUndefined(); + }); + + it('sets the shared value directly', () => { + // Prepare + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act + cell.setShared(8); + + // Assess + expect(cell.getShared()).toBe(8); + expect(cell.get()).toBe(8); + }); + }); + + describe('with an unavailable InvokeStore', () => { + beforeEach(() => { + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + vi.stubGlobal('awslambda', undefined); + }); + + it.each([ + { + method: 'get', + action: (cell: InvocationScoped) => cell.get(), + }, + { + method: 'set', + action: (cell: InvocationScoped) => cell.set(8), + }, + { + method: 'reset', + action: (cell: InvocationScoped) => cell.reset(), + }, + { + method: 'isScoped', + action: (cell: InvocationScoped) => cell.isScoped, + }, + ])('throws when calling $method', ({ action }) => { + // Prepare + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act & Assess + expect(() => action(cell)).toThrow('InvokeStore is not available'); + }); + + it('does not consult the InvokeStore for the shared value', () => { + // Prepare + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act + cell.setShared(8); + + // Assess + expect(cell.getShared()).toBe(8); + }); + }); + + describe('without an active invocation context', () => { + it('reads and writes the value shared across invocations', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act + cell.set(8); + + // Assess + expect(cell.get()).toBe(8); + expect(cell.getShared()).toBe(8); + expect(cell.isScoped).toBe(false); + }); + + it('restores the initial value when reset', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + cell.set(8); + + // Act + cell.reset(); + + // Assess + expect(cell.getShared()).toBe(12); + }); + }); + + describe('within an invocation context', () => { + it('scopes the value to the invocation', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act + const scopedValue = store.run({}, () => { + cell.set(8); + + return cell.get(); + }); + + // Assess + expect(scopedValue).toBe(8); + expect(cell.isScoped).toBe(false); + expect(cell.getShared()).toBe(12); + }); + + it('falls back to the value shared across invocations', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + cell.setShared(8); + + // Act + const scopedValue = store.run({}, () => cell.get()); + + // Assess + expect(scopedValue).toBe(8); + }); + + it('creates a fresh value for each invocation', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.attributes', { + fresh: () => ({ count: 0 }), + }); + + // Act + const [firstValue, firstValueAgain] = store.run({}, () => [ + cell.get(), + cell.get(), + ]); + const secondValue = store.run({}, () => cell.get()); + + // Assess + expect(firstValueAgain).toBe(firstValue); + expect(secondValue).not.toBe(firstValue); + }); + + it('bypasses the invocation context for the shared value', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act + const scopedValue = store.run({}, () => { + cell.setShared(8); + + return cell.get(); + }); + + // Assess + expect(scopedValue).toBe(8); + expect(cell.getShared()).toBe(8); + }); + + it('reports that reads and writes are scoped', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act + const isScoped = store.run({}, () => cell.isScoped); + + // Assess + expect(isScoped).toBe(true); + }); + + it('falls back to the shared value when reset', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act + const scopedValue = store.run({}, () => { + cell.set(8); + cell.reset(); + + return cell.get(); + }); + + // Assess + expect(scopedValue).toBe(12); + }); + + it('creates a new value when a fresh cell is reset', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.attributes', { + fresh: () => ({ count: 0 }), + }); + + // Act + const [firstValue, secondValue] = store.run({}, () => { + const first = cell.get(); + cell.reset(); + + return [first, cell.get()]; + }); + + // Assess + expect(secondValue).not.toBe(firstValue); + }); + + it('scopes a value created before any invocation context exists', async () => { + // Prepare + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + + // Act + const scopedValue = store.run({}, () => { + cell.set(8); + + return cell.get(); + }); + + // Assess + expect(scopedValue).toBe(8); + expect(cell.getShared()).toBe(12); + }); + + it('shadows the shared value with an undefined value', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped( + 'powertools.test.level', + { initial: 12 } + ); + cell.setShared(8); + + // Act + const scopedValue = store.run({}, () => { + cell.set(undefined); + + return cell.get(); + }); + + // Assess + expect(scopedValue).toBeUndefined(); + expect(cell.getShared()).toBe(8); + }); + + it('falls back to the shared value when reset after an undefined value', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped( + 'powertools.test.level', + { initial: 12 } + ); + cell.setShared(8); + + // Act + const scopedValue = store.run({}, () => { + cell.set(undefined); + cell.reset(); + + return cell.get(); + }); + + // Assess + expect(scopedValue).toBe(8); + }); + + it('shadows the created value with an undefined value', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped<{ count: number } | undefined>( + 'powertools.test.attributes', + { fresh: () => ({ count: 0 }) } + ); + + // Act + const scopedValue = store.run({}, () => { + cell.get(); + cell.set(undefined); + + return cell.get(); + }); + + // Assess + expect(scopedValue).toBeUndefined(); + }); + + it('creates a new value when reset after an undefined value', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const store = await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped<{ count: number } | undefined>( + 'powertools.test.attributes', + { fresh: () => ({ count: 0 }) } + ); + + // Act + const scopedValue = store.run({}, () => { + cell.set(undefined); + cell.reset(); + + return cell.get(); + }); + + // Assess + expect(scopedValue).toEqual({ count: 0 }); + }); + + it('isolates concurrent invocations from each other', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); + const cell = new InvocationScoped('powertools.test.level', { + initial: 12, + }); + + // Act + const [firstValue, secondValue] = await sequence( + { + sideEffects: [() => cell.set(8), () => {}], + return: () => cell.get(), + }, + { + sideEffects: [() => {}, () => cell.set(20)], + return: () => cell.get(), + }, + { useInvokeStore: true } + ); + + // Assess + expect(firstValue).toBe(8); + expect(secondValue).toBe(20); + expect(cell.getShared()).toBe(12); + }); + }); +}); diff --git a/packages/commons/tests/unit/envUtils.test.ts b/packages/commons/tests/unit/envUtils.test.ts index 35e538bad8..53cae74417 100644 --- a/packages/commons/tests/unit/envUtils.test.ts +++ b/packages/commons/tests/unit/envUtils.test.ts @@ -428,7 +428,7 @@ describe('Functions: envUtils', () => { }); describe('Function: shouldUseInvokeStore', () => { - it('returns true when AWS_LAMBDA_MAX_CONCURRENCY is not set', () => { + it('returns false when AWS_LAMBDA_MAX_CONCURRENCY is not set', () => { // Act const result = shouldUseInvokeStore(); @@ -436,7 +436,7 @@ describe('Functions: envUtils', () => { expect(result).toBe(false); }); - it('returns false when AWS_LAMBDA_MAX_CONCURRENCY is set', () => { + it('returns true when AWS_LAMBDA_MAX_CONCURRENCY is set', () => { // Prepare vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); diff --git a/packages/logger/package.json b/packages/logger/package.json index ead4b99c3b..89f376a21b 100644 --- a/packages/logger/package.json +++ b/packages/logger/package.json @@ -76,6 +76,7 @@ "devDependencies": { "@aws-lambda-powertools/testing-utils": "file:../testing", "@aws-sdk/client-lambda": "^3.1121.0", + "@aws/lambda-invoke-store": "0.3.0", "@types/promise-retry": "^1.1.3", "promise-retry": "^2.0.1" }, @@ -102,7 +103,6 @@ "url": "https://github.com/aws-powertools/powertools-lambda-typescript/issues" }, "dependencies": { - "@aws/lambda-invoke-store": "0.3.0", "@aws-lambda-powertools/commons": "2.35.0" }, "keywords": [ diff --git a/packages/logger/src/LogAttributesStore.ts b/packages/logger/src/LogAttributesStore.ts index f8d517bfc2..9e6f0fb521 100644 --- a/packages/logger/src/LogAttributesStore.ts +++ b/packages/logger/src/LogAttributesStore.ts @@ -1,6 +1,5 @@ -import '@aws/lambda-invoke-store'; import { deepMerge } from '@aws-lambda-powertools/commons'; -import { shouldUseInvokeStore } from '@aws-lambda-powertools/commons/utils/env'; +import { InvocationScoped } from '@aws-lambda-powertools/commons/utils/invocation-scoped'; import type { LambdaFunctionContext, LogAttributes } from './types/logKeys.js'; /** @@ -12,55 +11,26 @@ import type { LambdaFunctionContext, LogAttributes } from './types/logKeys.js'; * every method call to support Lambda's transition to async contexts. */ class LogAttributesStore { - readonly #temporaryAttributesKey = Symbol( - 'powertools.logger.temporaryAttributes' + readonly #temporaryAttributes = new InvocationScoped( + 'powertools.logger.temporaryAttributes', + { fresh: () => ({}) } ); - readonly #keysKey = Symbol('powertools.logger.keys'); - readonly #lambdaContextKey = Symbol('powertools.logger.lambdaContext'); + readonly #keys = new InvocationScoped>( + 'powertools.logger.keys', + { fresh: () => new Map() } + ); + readonly #lambdaContext = new InvocationScoped< + LambdaFunctionContext | undefined + >('powertools.logger.lambdaContext', { initial: undefined }); - #fallbackTemporaryAttributes: LogAttributes = {}; - readonly #fallbackKeys: Map = new Map(); #persistentAttributes: LogAttributes = {}; - #fallbackLambdaContext: LambdaFunctionContext | undefined; #getTemporaryAttributes(): LogAttributes { - if (!shouldUseInvokeStore()) { - return this.#fallbackTemporaryAttributes; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - let stored = store.get(this.#temporaryAttributesKey) as - | LogAttributes - | undefined; - if (stored == null) { - stored = {}; - store.set(this.#temporaryAttributesKey, stored); - } - return stored; + return this.#temporaryAttributes.get(); } #getKeys(): Map { - if (!shouldUseInvokeStore()) { - return this.#fallbackKeys; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - let stored = store.get(this.#keysKey) as - | Map - | undefined; - if (stored == null) { - stored = new Map(); - store.set(this.#keysKey, stored); - } - return stored; + return this.#keys.get(); } public appendTemporaryKeys(attributes: LogAttributes): void { @@ -104,33 +74,15 @@ class LogAttributesStore { } } - if (!shouldUseInvokeStore()) { - this.#fallbackTemporaryAttributes = {}; - return; - } - - globalThis.awslambda.InvokeStore?.set(this.#temporaryAttributesKey, {}); + this.#temporaryAttributes.reset(); } public setLambdaContext(context: LambdaFunctionContext): void { - if (!shouldUseInvokeStore()) { - this.#fallbackLambdaContext = context; - return; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - globalThis.awslambda.InvokeStore.set(this.#lambdaContextKey, context); + this.#lambdaContext.set(context); } public getLambdaContext(): LambdaFunctionContext | undefined { - if (!shouldUseInvokeStore()) { - return this.#fallbackLambdaContext; - } - - return globalThis.awslambda?.InvokeStore?.get(this.#lambdaContextKey); + return this.#lambdaContext.get(); } public setPersistentAttributes(attributes: LogAttributes): void { diff --git a/packages/logger/src/LogInvocationStore.ts b/packages/logger/src/LogInvocationStore.ts index c8c5b9a2fb..2ad77701c6 100644 --- a/packages/logger/src/LogInvocationStore.ts +++ b/packages/logger/src/LogInvocationStore.ts @@ -1,5 +1,4 @@ -import '@aws/lambda-invoke-store'; -import { shouldUseInvokeStore } from '@aws-lambda-powertools/commons/utils/env'; +import { InvocationScoped } from '@aws-lambda-powertools/commons/utils/invocation-scoped'; import { CircularMap, type SizedSet } from './logBuffer.js'; /** @@ -20,14 +19,17 @@ import { CircularMap, type SizedSet } from './logBuffer.js'; * `maxBytes` (only after {@link LogInvocationStore.configureBuffer | `configureBuffer()`}). */ class LogInvocationStore { - readonly #logLevelKey = Symbol('powertools.logger.logLevel'); - readonly #bufferKey = Symbol('powertools.logger.buffer'); - #baseLogLevel: number; + readonly #logLevel: InvocationScoped; + readonly #buffer = new InvocationScoped>( + 'powertools.logger.buffer', + { fresh: () => new CircularMap({ maxBytesSize: this.#maxBytes }) } + ); #maxBytes!: number; - #fallbackBuffer!: CircularMap; public constructor(logLevel: number) { - this.#baseLogLevel = logLevel; + this.#logLevel = new InvocationScoped('powertools.logger.logLevel', { + initial: logLevel, + }); } /** @@ -35,19 +37,7 @@ class LogInvocationStore { * set, otherwise the base level. */ public getLogLevel(): number { - if (!shouldUseInvokeStore()) { - return this.#baseLogLevel; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - return ( - (globalThis.awslambda.InvokeStore.get(this.#logLevelKey) as - | number - | undefined) ?? this.#baseLogLevel - ); + return this.#logLevel.get(); } /** @@ -55,7 +45,7 @@ class LogInvocationStore { * InvokeStore. */ public getBaseLogLevel(): number { - return this.#baseLogLevel; + return this.#logLevel.getShared(); } /** @@ -64,18 +54,7 @@ class LogInvocationStore { * context (e.g. during initialization) the base level is set instead. */ public setLogLevel(logLevel: number): void { - if (shouldUseInvokeStore()) { - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - if (store.hasContext()) { - store.set(this.#logLevelKey, logLevel); - return; - } - } - this.#baseLogLevel = logLevel; + this.#logLevel.set(logLevel); } /** @@ -84,7 +63,7 @@ class LogInvocationStore { * Controls, which apply environment-wide. */ public setBaseLogLevel(logLevel: number): void { - this.#baseLogLevel = logLevel; + this.#logLevel.setShared(logLevel); } /** @@ -97,25 +76,11 @@ class LogInvocationStore { */ public configureBuffer(maxBytes: number): void { this.#maxBytes = maxBytes; - this.#fallbackBuffer = new CircularMap({ maxBytesSize: maxBytes }); + this.#buffer.setShared(new CircularMap({ maxBytesSize: maxBytes })); } #getBuffer(): CircularMap { - if (!shouldUseInvokeStore()) { - return this.#fallbackBuffer; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - let buffer = store.get(this.#bufferKey) as CircularMap | undefined; - if (buffer == null) { - buffer = new CircularMap({ maxBytesSize: this.#maxBytes }); - store.set(this.#bufferKey, buffer); - } - return buffer; + return this.#buffer.get(); } /** @@ -132,7 +97,7 @@ class LogInvocationStore { // retaining stale logs. Under concurrency the buffer is scoped to the // invocation via the InvokeStore, so no cleanup is needed and clearing // would wipe other in-flight invocations' logs. - if (!shouldUseInvokeStore() && buffer.has(traceId) === false) { + if (!this.#buffer.isScoped && buffer.has(traceId) === false) { buffer.clear(); } buffer.setItem(traceId, value, logLevel); diff --git a/packages/logger/tests/unit/concurrency/logAttributesStore.test.ts b/packages/logger/tests/unit/concurrency/logAttributesStore.test.ts index 3edf629ab8..452d597113 100644 --- a/packages/logger/tests/unit/concurrency/logAttributesStore.test.ts +++ b/packages/logger/tests/unit/concurrency/logAttributesStore.test.ts @@ -1,6 +1,6 @@ import { InvokeStore } from '@aws/lambda-invoke-store'; import { sequence } from '@aws-lambda-powertools/testing-utils'; -import { beforeEach, describe, expect, it, vi } from 'vitest'; +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { LogAttributesStore } from '../../../src/LogAttributesStore.js'; describe('LogAttributesStore concurrent invocation isolation', () => { @@ -10,6 +10,98 @@ describe('LogAttributesStore concurrent invocation isolation', () => { vi.unstubAllEnvs(); }); + afterEach(() => { + vi.unstubAllGlobals(); + }); + + it('stores shared temporary attributes outside an active invocation context', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); + const store = new LogAttributesStore(); + + // Act + store.appendTemporaryKeys({ shared: true }); + + // Assess + expect(store.getTemporaryAttributes()).toEqual({ shared: true }); + }); + + it('clears shared temporary attributes outside an active invocation context', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); + const store = new LogAttributesStore(); + store.appendTemporaryKeys({ shared: true }); + + // Act + store.clearTemporaryAttributes(); + + // Assess + expect(store.getTemporaryAttributes()).toEqual({}); + }); + + it('shares lambda context with invocations when set outside an active context', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const invokeStore = await InvokeStore.getInstanceAsync(); + const store = new LogAttributesStore(); + const lambdaContext = { + functionName: 'test-function', + memoryLimitInMB: '128', + functionVersion: '$LATEST', + invokedFunctionArn: 'arn:aws:lambda:us-east-1:123:function:test', + awsRequestId: 'request-id', + tenantId: 'tenant-id', + coldStart: true, + }; + + // Act + store.setLambdaContext(lambdaContext); + const scopedContext = invokeStore.run({}, () => store.getLambdaContext()); + + // Assess + expect(scopedContext).toBe(lambdaContext); + }); + + it('isolates fresh temporary attributes from the shared value', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const invokeStore = await InvokeStore.getInstanceAsync(); + const store = new LogAttributesStore(); + store.appendTemporaryKeys({ shared: true }); + + // Act + const scopedAttributes = invokeStore.run({}, () => + store.getTemporaryAttributes() + ); + + // Assess + expect(scopedAttributes).toEqual({}); + }); + + it.each([ + { + method: 'clearTemporaryAttributes', + action: (store: LogAttributesStore) => store.clearTemporaryAttributes(), + }, + { + method: 'getLambdaContext', + action: (store: LogAttributesStore) => store.getLambdaContext(), + }, + ])( + 'throws when calling $method with InvokeStore unavailable', + ({ action }) => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + vi.stubGlobal('awslambda', undefined); + const store = new LogAttributesStore(); + + // Act & Assess + expect(() => action(store)).toThrow('InvokeStore is not available'); + } + ); + it.each([ { description: 'without InvokeStore', diff --git a/packages/metrics/package.json b/packages/metrics/package.json index 905895cfbd..7222aefe8e 100644 --- a/packages/metrics/package.json +++ b/packages/metrics/package.json @@ -67,6 +67,7 @@ "devDependencies": { "@aws-lambda-powertools/testing-utils": "file:../testing", "@aws-sdk/client-cloudwatch": "^3.1121.0", + "@aws/lambda-invoke-store": "0.3.0", "@types/promise-retry": "^1.1.3", "promise-retry": "^2.0.1" }, @@ -89,7 +90,6 @@ "url": "https://github.com/aws-powertools/powertools-lambda-typescript/issues" }, "dependencies": { - "@aws/lambda-invoke-store": "0.3.0", "@aws-lambda-powertools/commons": "2.35.0" }, "keywords": [ diff --git a/packages/metrics/src/DimensionsStore.ts b/packages/metrics/src/DimensionsStore.ts index 888d3b5d6a..1e4566bb2d 100644 --- a/packages/metrics/src/DimensionsStore.ts +++ b/packages/metrics/src/DimensionsStore.ts @@ -1,5 +1,4 @@ -import '@aws/lambda-invoke-store'; -import { shouldUseInvokeStore } from '@aws-lambda-powertools/commons/utils/env'; +import { InvocationScoped } from '@aws-lambda-powertools/commons/utils/invocation-scoped'; import type { Dimensions } from './types/Metrics.js'; /** @@ -11,47 +10,22 @@ import type { Dimensions } from './types/Metrics.js'; * every method call to support Lambda's transition to async contexts. */ class DimensionsStore { - readonly #dimensionsKey = Symbol('powertools.metrics.dimensions'); - readonly #dimensionSetsKey = Symbol('powertools.metrics.dimensionSets'); - - #fallbackDimensions: Dimensions = {}; - #fallbackDimensionSets: Dimensions[] = []; + readonly #dimensions = new InvocationScoped( + 'powertools.metrics.dimensions', + { fresh: () => ({}) } + ); + readonly #dimensionSets = new InvocationScoped( + 'powertools.metrics.dimensionSets', + { fresh: () => [] } + ); #defaultDimensions: Dimensions = {}; #getDimensions(): Dimensions { - if (!shouldUseInvokeStore()) { - return this.#fallbackDimensions; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - let stored = store.get(this.#dimensionsKey) as Dimensions | undefined; - if (stored == null) { - stored = {}; - store.set(this.#dimensionsKey, stored); - } - return stored; + return this.#dimensions.get(); } #getDimensionSets(): Dimensions[] { - if (!shouldUseInvokeStore()) { - return this.#fallbackDimensionSets; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - let stored = store.get(this.#dimensionSetsKey) as Dimensions[] | undefined; - if (stored == null) { - stored = []; - store.set(this.#dimensionSetsKey, stored); - } - return stored; + return this.#dimensionSets.get(); } public addDimension(name: string, value: string): string { @@ -73,19 +47,8 @@ class DimensionsStore { } public clearRequestDimensions(): void { - if (!shouldUseInvokeStore()) { - this.#fallbackDimensions = {}; - this.#fallbackDimensionSets = []; - return; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - store.set(this.#dimensionsKey, {}); - store.set(this.#dimensionSetsKey, []); + this.#dimensions.reset(); + this.#dimensionSets.reset(); } public clearDefaultDimensions(): void { diff --git a/packages/metrics/src/MetadataStore.ts b/packages/metrics/src/MetadataStore.ts index 9f74a51331..cb11707414 100644 --- a/packages/metrics/src/MetadataStore.ts +++ b/packages/metrics/src/MetadataStore.ts @@ -1,5 +1,4 @@ -import '@aws/lambda-invoke-store'; -import { shouldUseInvokeStore } from '@aws-lambda-powertools/commons/utils/env'; +import { InvocationScoped } from '@aws-lambda-powertools/commons/utils/invocation-scoped'; /** * Manages storage of metrics #metadata with automatic context detection. @@ -10,28 +9,13 @@ import { shouldUseInvokeStore } from '@aws-lambda-powertools/commons/utils/env'; * every method call to support Lambda's transition to async contexts. */ class MetadataStore { - readonly #metadataKey = Symbol('powertools.metrics.metadata'); - - #fallbackStorage: Record = {}; + readonly #metadata = new InvocationScoped>( + 'powertools.metrics.metadata', + { fresh: () => ({}) } + ); #getStorage(): Record { - if (!shouldUseInvokeStore()) { - return this.#fallbackStorage; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - let stored = store.get(this.#metadataKey) as - | Record - | undefined; - if (stored == null) { - stored = {}; - store.set(this.#metadataKey, stored); - } - return stored; + return this.#metadata.get(); } public set(key: string, value: string): string { @@ -44,17 +28,7 @@ class MetadataStore { } public clear(): void { - if (!shouldUseInvokeStore()) { - this.#fallbackStorage = {}; - return; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - store.set(this.#metadataKey, {}); + this.#metadata.reset(); } } diff --git a/packages/metrics/src/MetricsStore.ts b/packages/metrics/src/MetricsStore.ts index 5446465ffd..675c35412f 100644 --- a/packages/metrics/src/MetricsStore.ts +++ b/packages/metrics/src/MetricsStore.ts @@ -1,6 +1,5 @@ -import '@aws/lambda-invoke-store'; import { isIntegerNumber } from '@aws-lambda-powertools/commons/typeutils'; -import { shouldUseInvokeStore } from '@aws-lambda-powertools/commons/utils/env'; +import { InvocationScoped } from '@aws-lambda-powertools/commons/utils/invocation-scoped'; import { MetricResolution as MetricResolutions } from './constants.js'; import type { MetricResolution, @@ -18,28 +17,17 @@ import type { * every method call to support Lambda's transition to async contexts. */ class MetricsStore { - readonly #storedMetricsKey = Symbol('powertools.metrics.storedMetrics'); - readonly #timestampKey = Symbol('powertools.metrics.timestamp'); - - #fallbackStorage: StoredMetrics = {}; - #fallbackTimestamp?: number; + readonly #storedMetrics = new InvocationScoped( + 'powertools.metrics.storedMetrics', + { fresh: () => ({}) } + ); + readonly #timestamp = new InvocationScoped( + 'powertools.metrics.timestamp', + { initial: undefined } + ); #getStorage(): StoredMetrics { - if (!shouldUseInvokeStore()) { - return this.#fallbackStorage; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - let stored = store.get(this.#storedMetricsKey) as StoredMetrics | undefined; - if (stored == null) { - stored = {}; - store.set(this.#storedMetricsKey, stored); - } - return stored; + return this.#storedMetrics.get(); } public getMetric(name: string): StoredMetric | undefined { @@ -109,19 +97,8 @@ class MetricsStore { } public clearMetrics(): void { - if (!shouldUseInvokeStore()) { - this.#fallbackStorage = {}; - this.#fallbackTimestamp = undefined; - return; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - store.set(this.#storedMetricsKey, {}); - store.set(this.#timestampKey, undefined); + this.#storedMetrics.reset(); + this.#timestamp.set(undefined); } public hasMetrics(): boolean { @@ -133,32 +110,13 @@ class MetricsStore { } public getTimestamp(): number | undefined { - if (!shouldUseInvokeStore()) { - return this.#fallbackTimestamp; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - return store.get(this.#timestampKey) as number | undefined; + return this.#timestamp.get(); } public setTimestamp(timestamp: number | Date): number { const timestampMs = this.#convertTimestampToEmfFormat(timestamp); - if (!shouldUseInvokeStore()) { - this.#fallbackTimestamp = timestampMs; - return timestampMs; - } - - if (globalThis.awslambda?.InvokeStore === undefined) { - throw new Error('InvokeStore is not available'); - } - - const store = globalThis.awslambda.InvokeStore; - store.set(this.#timestampKey, timestampMs); + this.#timestamp.set(timestampMs); return timestampMs; } diff --git a/packages/metrics/tests/unit/concurrency/dimensionsStore.test.ts b/packages/metrics/tests/unit/concurrency/dimensionsStore.test.ts index 15f2521cd7..42ec3379ba 100644 --- a/packages/metrics/tests/unit/concurrency/dimensionsStore.test.ts +++ b/packages/metrics/tests/unit/concurrency/dimensionsStore.test.ts @@ -70,6 +70,21 @@ describe('DimensionsStore concurrent invocation isolation', () => { }); }); + it('uses shared dimensions when no invocation context is active', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); + const store = new DimensionsStore(); + + // Act + store.addDimension('env', 'prod'); + store.addDimensionSet({ service: 'api' }); + + // Assess + expect(store.getDimensions()).toEqual({ env: 'prod' }); + expect(store.getDimensionSets()).toEqual([{ service: 'api' }]); + }); + it.each([ { description: 'without InvokeStore', diff --git a/packages/metrics/tests/unit/concurrency/metadataStore.test.ts b/packages/metrics/tests/unit/concurrency/metadataStore.test.ts index 8dcf6dc22f..16d435dac1 100644 --- a/packages/metrics/tests/unit/concurrency/metadataStore.test.ts +++ b/packages/metrics/tests/unit/concurrency/metadataStore.test.ts @@ -45,6 +45,19 @@ describe('MetadataStore concurrent invocation isolation', () => { }); }); + it('uses shared metadata when no invocation context is active', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); + const store = new MetadataStore(); + + // Act + store.set('env', 'prod'); + + // Assess + expect(store.getAll()).toEqual({ env: 'prod' }); + }); + it.each([ { description: 'without InvokeStore', diff --git a/packages/metrics/tests/unit/concurrency/metricsStore.test.ts b/packages/metrics/tests/unit/concurrency/metricsStore.test.ts index bbdbd67b28..6e76903e34 100644 --- a/packages/metrics/tests/unit/concurrency/metricsStore.test.ts +++ b/packages/metrics/tests/unit/concurrency/metricsStore.test.ts @@ -69,6 +69,43 @@ describe('MetricsStore concurrent invocation isolation', () => { }); }); + it('uses shared metrics when no invocation context is active', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + await InvokeStore.getInstanceAsync(); + const store = new MetricsStore(); + + // Act + store.setMetric('count', MetricUnit.Count, 1, 60); + store.setTimestamp(1000); + + // Assess + expect(store.getMetric('count')).toEqual({ + name: 'count', + unit: MetricUnit.Count, + value: 1, + resolution: 60, + }); + expect(store.getTimestamp()).toBe(1000); + }); + + it('clears a scoped timestamp when a shared timestamp exists', async () => { + // Prepare + vi.stubEnv('AWS_LAMBDA_MAX_CONCURRENCY', '10'); + const invokeStore = await InvokeStore.getInstanceAsync(); + const store = new MetricsStore(); + store.setTimestamp(1000); + + // Act + const timestamp = invokeStore.run({}, () => { + store.clearMetrics(); + return store.getTimestamp(); + }); + + // Assess + expect(timestamp).toBeUndefined(); + }); + it.each([ { description: 'without InvokeStore',