feat(awscdk): add AWS Lambda Extension development support (projen#1647)

Adds AWS Lambda Extension development support to AWS CDK projects. Fixes projen#1646 --- By submitting this pull request, I confirm that my contribution is made under the terms of the Apache 2.0 license.
FitX · Mar 13, 2022 · 81614ca · 81614ca
1 parent e0e8ed5
commit 81614ca
Show file tree

Hide file tree

Showing 13 changed files with 885 additions and 17 deletions.
diff --git a/docs/api/API.md b/docs/api/API.md
diff --git a/docs/awscdk.md b/docs/awscdk.md
@@ -101,6 +101,142 @@ new awscdk.LambdaFunction(p, {
 });
 ```
 
+## AWS Lambda Extensions
+
+An AWS [Lambda Extension][lambda-extensions-blog] is a way to integrate your
+preferred development, monitoring, observability, and governance tools with
+AWS Lambda. 
+
+Functionally, AWS [Lambda Extensions][lambda-extensions-blog] are long-running
+executable files that reside in the `extensions` subdirectory of your code
+asset. AWS Lambda executes all extensions before starting your handler's main
+process. These AWS Lambda Extensions interact with your function's main process
+and the [Lambda Extension API][lambda-extensions-api] to integrate with tools
+outside the Lambda environment. Projen helps with bundling and preparing your
+code as reusable Lambda Layers.
+
+[lambda-extensions-blog]: https://aws.amazon.com/blogs/aws/getting-started-with-using-your-favorite-operational-tools-on-aws-lambda-extensions-are-now-generally-available/
+[lambda-extensions-api]: https://docs.aws.amazon.com/lambda/latest/dg/runtimes-extensions-api.html
+
+To create an AWS Lambda Extension with Projen:
+
+- Create a file in your project's source tree called
+  `my-extension.lambda-extension.ts`
+- Run `npx projen`
+- Projen will automatically discover this file, generating an AWS Lambda Layer
+  Version named `MyExtensionLayerVersion` in a file named
+  `my-extension-layer-version.ts`.
+- Now you can instantiate `MyExtensionLayerVersion` and add it to your Lambda
+  functions.
+
+Offical AWS extension examples are available in the [AWS Samples][ext-samples]
+repository.
+
+[ext-samples]: https://github.com/aws-samples/aws-lambda-extensions
+
+**Example of an extension:**
+
+A skeleton for a Lambda extension follows below. Comments with `TODO` describe
+locations where you can provide your custom functionality.
+
+```ts
+#!/usr/bin/env node
+// ^ Don't forget this shebang - Lambda executes the bundled version of this
+// file directly and doesn't otherwise know it's a node script.
+
+import { basename } from 'path';
+
+// This example uses the `got` HTTP client and assumes that you have included
+// `got` in your `devDependencies`. But, you can use any HTTP client you like.
+import got from 'got';
+
+/**
+ * Your Lambda Extension's main loop
+ */
+async function main() {
+  const extensionInfo = await registerExtension([
+    ExtensionEventType.SHUTDOWN,
+    ExtensionEventType.INVOKE,
+  ]);
+
+  // TODO: Put your initialization code here. You can do things like
+  // testing a connection to your external tooling here.
+
+  while (true) {
+    const event = await getNextEvent(extensionInfo.extensionId);
+
+    switch (event.eventType) {
+      case ExtensionEventType.SHUTDOWN:
+        // TODO: Do something when the lambda extension is being
+        // shut down. You might do things here like de-registering
+        // your extension from your external tooling.
+        return 0;
+
+      case ExtensionEventType.INVOKE:
+        // TODO: Do something every time your function is invoked,
+        // such as re-establishing a connection with your external
+        // tooling after the Lambda has thawed from a period of
+        // freezing due to inactivity.
+        break;
+
+      default:
+        console.log(`Unhandled event type ${event.eventType}`);
+    }
+  }
+}
+
+const EXTENSION_API_BASE_URL = `http://${process.env.AWS_LAMBDA_RUNTIME_API}/2020-01-01/extension`;
+
+enum ExtensionEventType {
+  INVOKE = 'INVOKE',
+  SHUTDOWN = 'SHUTDOWN',
+}
+
+interface ExtensionEvent {
+  readonly eventType: ExtensionEventType;
+  // For complete event structures, see:
+  // https://docs.aws.amazon.com/lambda/latest/dg/runtimes-extensions-api.html
+}
+
+async function registerExtension(events: ExtensionEventType[]) {
+  // Do not set a timeout on the GET call, as the extension can be suspended
+  // for a period of time until there is an event to return.
+  const res = await got.post(`${EXTENSION_API_BASE_URL}/register`, {
+    json: { events },
+    headers: {
+      'Lambda-Extension-Name': basename(__filename),
+    },
+  });
+
+  const header = res.headers['lambda-extension-identifier'];
+  const extensionId = Array.isArray(header) ? header[0] : header;
+  const json = JSON.parse(res.body);
+
+  return {
+    extensionId,
+    functionName: json.functionName as string,
+    functionVersion: json.functionVersion as string,
+  };
+}
+
+function getNextEvent(extensionId: string): Promise<ExtensionEvent> {
+  return got(`${EXTENSION_API_BASE_URL}/event/next`, {
+    headers: {
+      'Lambda-Extension-Identifier': extensionId,
+    },
+  }).json();
+}
+
+main()
+  .then(statusCode => {
+    process.exit(statusCode);
+  })
+  .catch(e => {
+    console.error(e);
+    process.exit(1);
+  });
+```
+
 ## Integration Snapshot Tests
 
 Files in the `test/` tree with the `.integ.ts` suffix are recognized as

diff --git a/src/awscdk/auto-discover.ts b/src/awscdk/auto-discover.ts
@@ -10,7 +10,14 @@ import {
   IntegrationTest,
   IntegrationTestCommonOptions,
 } from "./integration-test";
-import { TYPESCRIPT_LAMBDA_EXT } from "./internal";
+import {
+  TYPESCRIPT_LAMBDA_EXT,
+  TYPESCRIPT_LAMBDA_EXTENSION_EXT,
+} from "./internal";
+import {
+  LambdaExtension,
+  LambdaExtensionCommonOptions,
+} from "./lambda-extension";
 import { LambdaFunction, LambdaFunctionCommonOptions } from "./lambda-function";
 
 /**
@@ -68,7 +75,7 @@ export interface LambdaAutoDiscoverOptions extends AutoDiscoverCommonOptions {
   readonly srcdir: string;
 
   /**
-   * Options for auto-discovery of AWS Lambda functions.
+   * Options for AWS Lambda functions.
    */
   readonly lambdaOptions?: LambdaFunctionCommonOptions;
 }
@@ -93,11 +100,49 @@ export class LambdaAutoDiscover extends AutoDiscoverBase {
   }
 }
 
+/**
+ * Options for `LambdaExtensionAutoDiscover`
+ */
+export interface LambdaExtensionAutoDiscoverOptions
+  extends AutoDiscoverCommonOptions {
+  /**
+   * Project source tree (relative to project output directory).
+   */
+  readonly srcdir: string;
+
+  /**
+   * Options for lambda extensions.
+   */
+  readonly lambdaExtensionOptions?: LambdaExtensionCommonOptions;
+}
+
+/**
+ * Creates Lambda Extensions from entrypoints discovered in the project's
+ * source tree.
+ */
+export class LambdaExtensionAutoDiscover extends AutoDiscoverBase {
+  constructor(project: Project, options: LambdaExtensionAutoDiscoverOptions) {
+    super(project, {
+      projectdir: options.srcdir,
+      extension: TYPESCRIPT_LAMBDA_EXTENSION_EXT,
+    });
+
+    for (const entrypoint of this.entrypoints) {
+      new LambdaExtension(this.project, {
+        entrypoint,
+        cdkDeps: options.cdkDeps,
+        ...options.lambdaExtensionOptions,
+      });
+    }
+  }
+}
+
 /**
  * Options for `AutoDiscover`
  */
 export interface AutoDiscoverOptions
   extends LambdaAutoDiscoverOptions,
+    LambdaExtensionAutoDiscoverOptions,
     IntegrationTestAutoDiscoverOptions {
   /**
    * Auto-discover lambda functions.
@@ -106,6 +151,13 @@ export interface AutoDiscoverOptions
    */
   readonly lambdaAutoDiscover?: boolean;
 
+  /**
+   * Auto-discover lambda extensions.
+   *
+   * @default true
+   */
+  readonly lambdaExtensionAutoDiscover?: boolean;
+
   /**
    * Auto-discover integration tests.
    *
@@ -131,6 +183,15 @@ export class AutoDiscover extends Component {
       });
     }
 
+    if (options.lambdaExtensionAutoDiscover ?? true) {
+      new LambdaExtensionAutoDiscover(this.project, {
+        cdkDeps: options.cdkDeps,
+        tsconfigPath: options.tsconfigPath,
+        srcdir: options.srcdir,
+        lambdaExtensionOptions: options.lambdaExtensionOptions,
+      });
+    }
+
     if (options.integrationTestAutoDiscover ?? true) {
       new IntegrationTestAutoDiscover(this.project, {
         cdkDeps: options.cdkDeps,

diff --git a/src/awscdk/awscdk-app-ts.ts b/src/awscdk/awscdk-app-ts.ts
@@ -31,6 +31,15 @@ export interface AwsCdkTypeScriptAppOptions
    */
   readonly lambdaAutoDiscover?: boolean;
 
+  /**
+   * Automatically adds an `awscdk.LambdaExtension` for each `.lambda-extension.ts`
+   * entrypoint in your source tree. If this is disabled, you can manually add an
+   * `awscdk.AutoDiscover` component to your project
+   *
+   * @default true
+   */
+  readonly lambdaExtensionAutoDiscover?: boolean;
+
   /**
    * Automatically discovers and creates integration tests for each `.integ.ts`
    * file in under your test directory.
@@ -154,6 +163,7 @@ export class AwsCdkTypeScriptApp extends TypeScriptAppProject {
       tsconfigPath: this.tsconfigDev.fileName,
       cdkDeps: this.cdkDeps,
       lambdaAutoDiscover: options.lambdaAutoDiscover ?? true,
+      lambdaExtensionAutoDiscover: options.lambdaExtensionAutoDiscover ?? true,
       integrationTestAutoDiscover: options.integrationTestAutoDiscover ?? true,
     });
   }

diff --git a/src/awscdk/awscdk-construct.ts b/src/awscdk/awscdk-construct.ts
@@ -22,6 +22,15 @@ export interface AwsCdkConstructLibraryOptions
    */
   readonly lambdaAutoDiscover?: boolean;
 
+  /**
+   * Automatically adds an `awscdk.LambdaExtension` for each `.lambda-extension.ts`
+   * entrypoint in your source tree. If this is disabled, you can manually add an
+   * `awscdk.AutoDiscover` component to your project
+   *
+   * @default true
+   */
+  readonly lambdaExtensionAutoDiscover?: boolean;
+
   /**
    * Automatically discovers and creates integration tests for each `.integ.ts`
    * file in under your test directory.
@@ -80,6 +89,7 @@ export class AwsCdkConstructLibrary extends ConstructLibrary {
       tsconfigPath: this.tsconfigDev.fileName,
       cdkDeps: this.cdkDeps,
       lambdaAutoDiscover: options.lambdaAutoDiscover ?? true,
+      lambdaExtensionAutoDiscover: options.lambdaExtensionAutoDiscover ?? true,
       integrationTestAutoDiscover: options.integrationTestAutoDiscover ?? true,
     });
   }

diff --git a/src/awscdk/index.ts b/src/awscdk/index.ts
@@ -10,3 +10,4 @@ export * from "./cdk-config";
 export * from "./cdk-tasks";
 export * from "./integration-test";
 export * from "./lambda-function";
+export * from "./lambda-extension";
diff --git a/src/awscdk/internal.ts b/src/awscdk/internal.ts
@@ -1,3 +1,5 @@
+import { sep, posix } from "path";
+
 /**
  * Feature flags as of v1.130.0
  */
@@ -21,3 +23,15 @@ export const FEATURE_FLAGS = [
  * Suffix for AWS Lambda handlers.
  */
 export const TYPESCRIPT_LAMBDA_EXT = ".lambda.ts";
+
+/**
+ * Suffix for AWS Lambda Extensions.
+ */
+export const TYPESCRIPT_LAMBDA_EXTENSION_EXT = ".lambda-extension.ts";
+
+/**
+ * Converts the given path string to posix if it wasn't already.
+ */
+export function convertToPosixPath(p: string) {
+  return p.split(sep).join(posix.sep);
+}