feat(commons): add cleanupPowertools function (#1473)

* feat: add cleanupPowertools hook

* tests: fixed tests

* chore: rename cleanup function

Author: dreamorosi

Diff for: packages/commons/src/middleware/cleanupMiddlewares.ts

@@ -0,0 +1,77 @@
+import {
+  TRACER_KEY,
+  METRICS_KEY,
+  LOGGER_KEY,
+  IDEMPOTENCY_KEY,
+} from './constants';
+import type { MiddyLikeRequest, CleanupFunction } from '../types/middy';
+
+// Typeguard to assert that an object is of Function type
+const isFunction = (obj: unknown): obj is CleanupFunction => {
+  return typeof obj === 'function';
+};
+
+/**
+ * Function used to cleanup Powertools for AWS resources when a Middy
+ * middleware [returns early](https://middy.js.org/docs/intro/early-interrupt)
+ * and terminates the middleware chain.
+ *
+ * When a middleware returns early, all the middleware lifecycle functions
+ * that come after it are not executed. This means that if a middleware
+ * was relying on certain logic to be run during the `after` or `onError`
+ * lifecycle functions, that logic will not be executed.
+ *
+ * This is the case for the middlewares that are part of Powertools for AWS
+ * which rely on these lifecycle functions to perform cleanup operations
+ * like closing the current segment in the tracer or flushing any stored
+ * metrics.
+ *
+ * When authoring a middleware that might return early, you can use this
+ * function to cleanup Powertools resources. This function will check if
+ * any cleanup function is present in the `request.internal` object and
+ * execute it.
+ *
+ * @example
+ * ```typescript
+ * import middy from '@middy/core';
+ * import { cleanupMiddlewares } from '@aws-lambda-powertools/commons/lib/middleware';
+ *
+ * // Example middleware that returns early
+ * const myCustomMiddleware = (): middy.MiddlewareObj => {
+ *   const before = async (request: middy.Request): Promise<undefined | string> => {
+ *     // If the request is a GET, return early (as an example)
+ *     if (request.event.httpMethod === 'GET') {
+ *       // Cleanup Powertools resources
+ *       await cleanupMiddlewares(request);
+ *       // Then return early
+ *       return 'GET method not supported';
+ *     }
+ *   };
+ *
+ *   return {
+ *     before,
+ *   };
+ * };
+ * ```
+ *
+ * @param request - The Middy request object
+ * @param options - An optional object that can be used to pass options to the function
+ */
+const cleanupMiddlewares = async (request: MiddyLikeRequest): Promise<void> => {
+  const cleanupFunctionNames = [
+    TRACER_KEY,
+    METRICS_KEY,
+    LOGGER_KEY,
+    IDEMPOTENCY_KEY,
+  ];
+  for (const functionName of cleanupFunctionNames) {
+    if (Object(request.internal).hasOwnProperty(functionName)) {
+      const functionReference = request.internal[functionName];
+      if (isFunction(functionReference)) {
+        await functionReference(request);
+      }
+    }
+  }
+};
+
+export { cleanupMiddlewares };

Diff for: packages/commons/src/middleware/constants.ts

@@ -0,0 +1,12 @@
+/**
+ * These constants are used to store cleanup functions in Middy's `request.internal` object.
+ * They are used by the `cleanupPowertools` function to check if any cleanup function
+ * is present and execute it.
+ */
+const PREFIX = 'powertools-for-aws';
+const TRACER_KEY = `${PREFIX}.tracer`;
+const METRICS_KEY = `${PREFIX}.metrics`;
+const LOGGER_KEY = `${PREFIX}.logger`;
+const IDEMPOTENCY_KEY = `${PREFIX}.idempotency`;
+
+export { TRACER_KEY, METRICS_KEY, LOGGER_KEY, IDEMPOTENCY_KEY };

Diff for: packages/commons/src/middleware/index.ts

@@ -0,0 +1,2 @@
+export * from './cleanupMiddlewares';
+export * from './constants';

Diff for: packages/commons/src/types/middy.ts

@@ -1,4 +1,4 @@
-import { Context } from 'aws-lambda';
+import type { Context } from 'aws-lambda';
 
 /**
  * We need to define these types and interfaces here because we can't import them from @middy/core.
@@ -22,14 +22,14 @@ type Request<
   };
 };
 
-declare type MiddlewareFn<
+type MiddlewareFn<
   TEvent = unknown,
   TResult = unknown,
   TErr = Error,
   TContext extends Context = Context
 > = (request: Request<TEvent, TResult, TErr, TContext>) => unknown;
 
-export type MiddlewareLikeObj<
+type MiddlewareLikeObj<
   TEvent = unknown,
   TResult = unknown,
   TErr = Error,
@@ -40,9 +40,21 @@ export type MiddlewareLikeObj<
   onError?: MiddlewareFn<TEvent, TResult, TErr, TContext>;
 };
 
-export type MiddyLikeRequest = {
+type MiddyLikeRequest = {
   event: unknown;
   context: Context;
   response: unknown | null;
   error: Error | null;
+  internal: {
+    [key: string]: unknown;
+  };
 };
+
+/**
+ * Cleanup function that is used to cleanup resources when a middleware returns early.
+ * Each Powertools for AWS middleware that needs to perform cleanup operations will
+ * store a cleanup function with this signature in the `request.internal` object.
+ */
+type CleanupFunction = (request: MiddyLikeRequest) => Promise<void>;
+
+export { MiddlewareLikeObj, MiddyLikeRequest, CleanupFunction };

Diff for: packages/commons/tests/unit/cleanupMiddlewares.test.ts

@@ -0,0 +1,66 @@
+/**
+ * Test Middy cleanupMiddlewares function
+ *
+ * @group unit/commons/cleanupMiddlewares
+ */
+import {
+  cleanupMiddlewares,
+  TRACER_KEY,
+  METRICS_KEY,
+} from '../../src/middleware';
+import { helloworldContext as context } from '../../src/samples/resources/contexts/hello-world';
+
+describe('Function: cleanupMiddlewares', () => {
+  it('calls the cleanup function that are present', async () => {
+    // Prepare
+    const mockCleanupFunction1 = jest.fn();
+    const mockCleanupFunction2 = jest.fn();
+    const mockRequest = {
+      event: {},
+      context: context,
+      response: null,
+      error: null,
+      internal: {
+        [TRACER_KEY]: mockCleanupFunction1,
+        [METRICS_KEY]: mockCleanupFunction2,
+      },
+    };
+
+    // Act
+    await cleanupMiddlewares(mockRequest);
+
+    // Assess
+    expect(mockCleanupFunction1).toHaveBeenCalledTimes(1);
+    expect(mockCleanupFunction1).toHaveBeenCalledWith(mockRequest);
+    expect(mockCleanupFunction2).toHaveBeenCalledTimes(1);
+    expect(mockCleanupFunction2).toHaveBeenCalledWith(mockRequest);
+  });
+  it('resolves successfully if no cleanup function is present', async () => {
+    // Prepare
+    const mockRequest = {
+      event: {},
+      context: context,
+      response: null,
+      error: null,
+      internal: {},
+    };
+
+    // Act & Assess
+    await expect(cleanupMiddlewares(mockRequest)).resolves.toBeUndefined();
+  });
+  it('resolves successfully if cleanup function is not a function', async () => {
+    // Prepare
+    const mockRequest = {
+      event: {},
+      context: context,
+      response: null,
+      error: null,
+      internal: {
+        [TRACER_KEY]: 'not a function',
+      },
+    };
+
+    // Act & Assess
+    await expect(cleanupMiddlewares(mockRequest)).resolves.toBeUndefined();
+  });
+});