Support ARM layers

integration_tests/correct_extension_apigateway_snapshot.json

@@ -815,7 +815,10 @@
           }
         },
         "MethodResponses": []
-      }
+      },
+      "DependsOn": [
+        "PythonHello27LambdaPermissionApiGateway"
+      ]
     },
     "ApiGatewayMethodUsersCreatePost": {
       "Type": "AWS::ApiGateway::Method",
@@ -858,7 +861,10 @@
           }
         },
         "MethodResponses": []
-      }
+      },
+      "DependsOn": [
+        "JavascriptHello10DashxLambdaPermissionApiGateway"
+      ]
     },
     "ApiGatewayDeploymentxxxx": {
       "Type": "AWS::ApiGateway::Deployment",

integration_tests/correct_forwarder_snapshot.json

@@ -997,7 +997,10 @@
           }
         },
         "MethodResponses": []
-      }
+      },
+      "DependsOn": [
+        "PythonHello27LambdaPermissionApiGateway"
+      ]
     },
     "ApiGatewayMethodUsersCreatePost": {
       "Type": "AWS::ApiGateway::Method",
@@ -1040,7 +1043,10 @@
           }
         },
         "MethodResponses": []
-      }
+      },
+      "DependsOn": [
+        "JavascriptHello10DashxLambdaPermissionApiGateway"
+      ]
     },
     "ApiGatewayDeploymentxxxx": {
       "Type": "AWS::ApiGateway::Deployment",

package.json

@@ -26,7 +26,7 @@
     "mock-fs": "4.13.0",
     "nock": "^13.0.5",
     "prettier": "^2.2.1",
-    "serverless": "2.38.0",
+    "serverless": "2.63.0",
     "ts-jest": "^26.4.4",
     "tslint": "^6.1.3",
     "typescript": "^4.1.3"

src/index.ts

@@ -12,7 +12,14 @@ import * as govLayers from "./layers-gov.json";
 import { version } from "../package.json";
 
 import { getConfig, setEnvConfiguration, forceExcludeDepsFromWebpack, hasWebpackPlugin, Configuration } from "./env";
-import { applyExtensionLayer, applyLambdaLibraryLayers, findHandlers, FunctionInfo, RuntimeType } from "./layer";
+import {
+  applyExtensionLayer,
+  applyLambdaLibraryLayers,
+  DEFAULT_ARCHITECTURE,
+  findHandlers,
+  FunctionInfo,
+  RuntimeType,
+} from "./layer";
 import { TracingMode, enableTracing } from "./tracing";
 import { redirectHandlers } from "./wrapper";
 import { addCloudWatchForwarderSubscriptions, addExecutionLogGroupsAndSubscriptions } from "./forwarder";
@@ -22,12 +29,6 @@ import { setMonitors } from "./monitors";
 import { getCloudFormationStackId } from "./monitor-api-requests";
 import { stringify } from "querystring";
 
-// Separate interface since DefinitelyTyped currently doesn't include tags or env
-export interface ExtendedFunctionDefinition extends FunctionDefinition {
-  tags?: { [key: string]: string };
-  environment?: { [key: string]: string };
-}
-
 enum TagKeys {
   Service = "service",
   Env = "env",

src/layer.spec.ts

@@ -24,11 +24,12 @@ type FunctionDefinitionAll = FunctionDefinitionHandler | FunctionDefinitionImage
 function createMockService(
   region: string,
   funcs: { [funcName: string]: Partial<FunctionDefinitionAll> },
+  architecture?: string,
   plugins?: string[],
   layers?: string[],
 ): Service {
   const service: Partial<Service> & { functions: any; plugins: any } = {
-    provider: { region, layers } as any,
+    provider: { region, layers, architecture } as any,
     getAllFunctionsNames: () => Object.keys(funcs),
     getFunction: (name) => funcs[name] as FunctionDefinitionAll,
     functions: funcs as any,
@@ -164,6 +165,7 @@ describe("applyLambdaLibraryLayers", () => {
       {
         "node-function": { handler: "myfile.handler", runtime: "nodejs10.x" },
       },
+      "x86_64",
       [],
       ["my-layer-1", "my-layer-2"],
     );
@@ -321,6 +323,161 @@ describe("applyLambdaLibraryLayers", () => {
       layers: ["node:2", "extension:5"],
     });
   });
+
+  it("adds correct lambda layer given architecture in function level", () => {
+    const handler = {
+      handler: { runtime: "python3.9", architecture: "arm64" },
+      type: RuntimeType.PYTHON,
+      runtime: "python3.9",
+    } as FunctionInfo;
+    const layers: LayerJSON = {
+      regions: {
+        "us-east-1": {
+          "python3.9": "python:3.9",
+          "python3.9-arm": "python-arm:3.9",
+          extension: "extension:11",
+          "extension-arm": "extension-arm:11",
+        },
+      },
+    };
+    const mockService = createMockService(
+      "us-east-1",
+      {
+        "python-function": { handler: "myfile.handler", runtime: "python3.9" },
+      },
+      "arm64",
+    );
+    applyLambdaLibraryLayers(mockService, [handler], layers);
+    applyExtensionLayer(mockService, [handler], layers);
+    expect(handler.handler).toEqual({
+      architecture: "arm64",
+      runtime: "python3.9",
+      layers: ["python-arm:3.9", "extension-arm:11"],
+    });
+  });
+
+  it("adds correct lambda layer given architecture in provider level", () => {
+    const handler = {
+      handler: { runtime: "python3.9" },
+      type: RuntimeType.PYTHON,
+      runtime: "python3.9",
+    } as FunctionInfo;
+    const layers: LayerJSON = {
+      regions: {
+        "us-east-1": {
+          "python3.9": "python:3.9",
+          "python3.9-arm": "python-arm:3.9",
+          extension: "extension:11",
+          "extension-arm": "extension-arm:11",
+        },
+      },
+    };
+    const mockService = createMockService(
+      "us-east-1",
+      {
+        "python-function": { handler: "myfile.handler", runtime: "python3.9" },
+      },
+      "arm64",
+    );
+    applyLambdaLibraryLayers(mockService, [handler], layers);
+    applyExtensionLayer(mockService, [handler], layers);
+    expect(handler.handler).toEqual({
+      runtime: "python3.9",
+      layers: ["python-arm:3.9", "extension-arm:11"],
+    });
+  });
+
+  it("uses default runtime layer if architecture not available for specified runtime", () => {
+    const handler = {
+      handler: { runtime: "python3.7", architecture: "arm64" },
+      type: RuntimeType.PYTHON,
+      runtime: "python3.7",
+    } as FunctionInfo;
+    const layers: LayerJSON = {
+      regions: {
+        "us-east-1": {
+          "python3.7": "python:3.7",
+          extension: "extension:11",
+          "extension-arm": "extension-arm:11",
+        },
+      },
+    };
+    const mockService = createMockService(
+      "us-east-1",
+      {
+        "python-function": { handler: "myfile.handler", runtime: "python3.7" },
+      },
+      "arm64",
+    );
+    applyLambdaLibraryLayers(mockService, [handler], layers);
+    applyExtensionLayer(mockService, [handler], layers);
+    expect(handler.handler).toEqual({
+      architecture: "arm64",
+      runtime: "python3.7",
+      layers: ["python:3.7", "extension:11"],
+    });
+  });
+
+  it("swaps previous layer when specifying arm architecture in functions", () => {
+    let handler = {
+      handler: { runtime: "python3.9", architecture: "arm64" },
+      type: RuntimeType.PYTHON,
+      runtime: "python3.9",
+    } as FunctionInfo;
+    (handler.handler as any).layers = ["python:3.9", "extension:11"];
+    const layers: LayerJSON = {
+      regions: {
+        "us-east-1": {
+          "python3.9": "python:3.9",
+          "python3.9-arm": "python-arm:3.9",
+          extension: "extension:11",
+          "extension-arm": "extension-arm:11",
+        },
+      },
+    };
+    const mockService = createMockService("us-east-1", {
+      "python-function": { handler: "myfile.handler", runtime: "python3.9" },
+    });
+    applyLambdaLibraryLayers(mockService, [handler], layers);
+    applyExtensionLayer(mockService, [handler], layers);
+    expect(handler.handler).toEqual({
+      architecture: "arm64",
+      runtime: "python3.9",
+      layers: ["python-arm:3.9", "extension-arm:11"],
+    });
+  });
+
+  it("swaps previous layer when specifying arm architecture in provider level", () => {
+    let handler = {
+      handler: { runtime: "python3.9" },
+      type: RuntimeType.PYTHON,
+      runtime: "python3.9",
+    } as FunctionInfo;
+    (handler.handler as any).layers = ["python:3.9", "extension:11"];
+    const layers: LayerJSON = {
+      regions: {
+        "us-east-1": {
+          "python3.9": "python:3.9",
+          "python3.9-arm": "python-arm:3.9",
+          extension: "extension:11",
+          "extension-arm": "extension-arm:11",
+        },
+      },
+    };
+    const mockService = createMockService(
+      "us-east-1",
+      {
+        "python-function": { handler: "myfile.handler", runtime: "python3.9" },
+      },
+      "arm64",
+    );
+    applyLambdaLibraryLayers(mockService, [handler], layers);
+    applyExtensionLayer(mockService, [handler], layers);
+    expect(handler.handler).toEqual({
+      runtime: "python3.9",
+      layers: ["python-arm:3.9", "extension-arm:11"],
+    });
+  });
 });
 
 describe("pushLayerARN", () => {

src/layer.ts

@@ -7,7 +7,6 @@
  */
 import { FunctionDefinition, FunctionDefinitionHandler } from "serverless";
 import Service from "serverless/classes/Service";
-const extensionLayerKey: string = "extension";
 export enum RuntimeType {
   NODE,
   PYTHON,
@@ -17,10 +16,19 @@ export enum RuntimeType {
 export interface FunctionInfo {
   name: string;
   type: RuntimeType;
-  handler: FunctionDefinition;
+  handler: ExtendedFunctionDefinition;
   runtime?: string;
 }
 
+export const X86_64_ARCHITECTURE = "x86_64";
+export const ARM64_ARCHITECTURE = "arm64";
+export const DEFAULT_ARCHITECTURE = X86_64_ARCHITECTURE;
+
+// Separate interface since DefinitelyTyped currently doesn't include tags or env
+export interface ExtendedFunctionDefinition extends FunctionDefinition {
+  architecture?: string;
+}
+
 export interface LayerJSON {
   regions: {
     [region: string]:
@@ -42,6 +50,12 @@ export const runtimeLookup: { [key: string]: RuntimeType } = {
   "python3.9": RuntimeType.PYTHON,
 };
 
+export const armKeys: { [key: string]: string } = {
+  "python3.8": "python3.8-arm",
+  "python3.9": "python3.9-arm",
+  extension: "extension-arm",
+};
+
 export function findHandlers(service: Service, exclude: string[], defaultRuntime?: string): FunctionInfo[] {
   return Object.entries(service.functions)
     .map(([name, handler]) => {
@@ -62,7 +76,6 @@ export function findHandlers(service: Service, exclude: string[], defaultRuntime
 
 export function applyLambdaLibraryLayers(service: Service, handlers: FunctionInfo[], layers: LayerJSON) {
   const { region } = service.provider;
-
   const regionRuntimes = layers.regions[region];
   if (regionRuntimes === undefined) {
     return;
@@ -74,9 +87,16 @@ export function applyLambdaLibraryLayers(service: Service, handlers: FunctionInf
     }
 
     const { runtime } = handler;
-    const lambdaLayerARN = runtime !== undefined ? regionRuntimes[runtime] : undefined;
-    let currentLayers = getLayers(service, handler);
+    const architecture =
+      (handler.handler as any).architecture ?? (service.provider as any).architecture ?? DEFAULT_ARCHITECTURE;
+    let runtimeKey: string | undefined = runtime;
+    if (architecture === ARM64_ARCHITECTURE && runtime && runtime in armKeys) {
+      runtimeKey = armKeys[runtime];
+      removePreviousLayer(service, handler, regionRuntimes[runtime]);
+    }
 
+    const lambdaLayerARN = runtimeKey !== undefined ? regionRuntimes[runtimeKey] : undefined;
+    let currentLayers = getLayers(service, handler);
     if (lambdaLayerARN) {
       currentLayers = pushLayerARN([lambdaLayerARN], currentLayers);
       setLayers(handler, currentLayers);
@@ -95,7 +115,17 @@ export function applyExtensionLayer(service: Service, handlers: FunctionInfo[],
     if (handler.type === RuntimeType.UNSUPPORTED) {
       continue;
     }
+    const { runtime } = handler;
+    const architecture =
+      (handler.handler as any).architecture ?? (service.provider as any).architecture ?? DEFAULT_ARCHITECTURE;
     let extensionLayerARN: string | undefined;
+    let extensionLayerKey: string = "extension";
+
+    if (architecture === ARM64_ARCHITECTURE && runtime && runtime in armKeys) {
+      removePreviousLayer(service, handler, regionRuntimes[extensionLayerKey]);
+      extensionLayerKey = armKeys[extensionLayerKey];
+    }
+
     extensionLayerARN = regionRuntimes[extensionLayerKey];
     let currentLayers = getLayers(service, handler);
     if (extensionLayerARN) {
@@ -126,6 +156,14 @@ function getLayers(service: Service, handler: FunctionInfo): string[] {
   return layerList;
 }
 
+function removePreviousLayer(service: Service, handler: FunctionInfo, previousLayer: string | undefined) {
+  let layersList = getLayers(service, handler);
+  if (new Set(layersList).has(previousLayer!)) {
+    layersList = layersList?.filter((layer) => layer !== previousLayer);
+  }
+  setLayers(handler, layersList);
+}
+
 function setLayers(handler: FunctionInfo, layers: string[]) {
   (handler.handler as any).layers = layers;
 }