feat(lambda): add SnapStart support for arm64 functions (#30898)

### Issue # (if applicable)

Closes #30899.

### Reason for this change

Lambda released SnapStart support for arm64 functions: https://aws.amazon.com/about-aws/whats-new/2024/07/aws-lambda-snapstart-java-functions-arm64-architecture/



### Description of changes

Allow configuring SnapStart on arm64 functions.



### Description of how you validated changes



```yarn test```

### Checklist
- [x] My code adheres to the [CONTRIBUTING GUIDE](https://github.com/aws/aws-cdk/blob/main/CONTRIBUTING.md) and [DESIGN GUIDELINES](https://github.com/aws/aws-cdk/blob/main/docs/DESIGN_GUIDELINES.md)

----

*By submitting this pull request, I confirm that my contribution is made under the terms of the Apache-2.0 license*

packages/@aws-cdk-testing/framework-integ/test/aws-lambda/test/integ.lambda.js.snapshot/aws-cdk-lambda-1.template.json

@@ -201,6 +201,66 @@
    "DependsOn": [
     "MySnapStartLambdaServiceRoleE0F04324"
    ]
+  },
+  "MySnapStartLambdaArmServiceRole731A8472": {
+   "Type": "AWS::IAM::Role",
+   "Properties": {
+    "AssumeRolePolicyDocument": {
+     "Statement": [
+      {
+       "Action": "sts:AssumeRole",
+       "Effect": "Allow",
+       "Principal": {
+        "Service": "lambda.amazonaws.com"
+       }
+      }
+     ],
+     "Version": "2012-10-17"
+    },
+    "ManagedPolicyArns": [
+     {
+      "Fn::Join": [
+       "",
+       [
+        "arn:",
+        {
+         "Ref": "AWS::Partition"
+        },
+        ":iam::aws:policy/service-role/AWSLambdaBasicExecutionRole"
+       ]
+      ]
+     }
+    ]
+   }
+  },
+  "MySnapStartLambdaArmEDD3F420": {
+   "Type": "AWS::Lambda::Function",
+   "Properties": {
+    "Architectures": [
+     "arm64"
+    ],
+    "Code": {
+     "S3Bucket": {
+      "Fn::Sub": "cdk-hnb659fds-assets-${AWS::AccountId}-${AWS::Region}"
+     },
+     "S3Key": "a37d3ef54c18e7738fe5dc008504591bd3b1f14c6a09ee91eac6d55f7ca5ba5f.zip"
+    },
+    "Description": "version-hash:522e22f485651f1644be7e317c07836f",
+    "Handler": "example.Handler::handleRequest",
+    "Role": {
+     "Fn::GetAtt": [
+      "MySnapStartLambdaArmServiceRole731A8472",
+      "Arn"
+     ]
+    },
+    "Runtime": "java21",
+    "SnapStart": {
+     "ApplyOn": "PublishedVersions"
+    }
+   },
+   "DependsOn": [
+    "MySnapStartLambdaArmServiceRole731A8472"
+   ]
   }
  },
  "Parameters": {

packages/@aws-cdk-testing/framework-integ/test/aws-lambda/test/integ.lambda.ts

@@ -45,4 +45,12 @@ new lambda.Function(stack, 'MySnapStartLambda', {
   snapStart: lambda.SnapStartConf.ON_PUBLISHED_VERSIONS,
 });
 
+new lambda.Function(stack, 'MySnapStartLambdaArm', {
+  code: lambda.Code.fromAsset(path.join(__dirname, 'handler-snapstart.zip')),
+  handler: 'example.Handler::handleRequest',
+  runtime: lambda.Runtime.JAVA_21,
+  architecture: lambda.Architecture.ARM_64,
+  snapStart: lambda.SnapStartConf.ON_PUBLISHED_VERSIONS,
+});
+
 app.synth();

packages/aws-cdk-lib/aws-lambda/README.md

@@ -967,7 +967,7 @@ managing concurrency.
 
 ## Lambda with SnapStart
 
-SnapStart is currently supported only on Java 11/Java 17 runtime. SnapStart does not support provisioned concurrency, the arm64 architecture, Amazon Elastic File System (Amazon EFS), or ephemeral storage greater than 512 MB. After you enable Lambda SnapStart for a particular Lambda function, publishing a new version of the function will trigger an optimization process.
+SnapStart is currently supported only on Java 11 and later [Java managed runtimes](https://docs.aws.amazon.com/lambda/latest/dg/lambda-runtimes.html). SnapStart does not support provisioned concurrency, Amazon Elastic File System (Amazon EFS), or ephemeral storage greater than 512 MB. After you enable Lambda SnapStart for a particular Lambda function, publishing a new version of the function will trigger an optimization process.
 
 See [the AWS documentation](https://docs.aws.amazon.com/lambda/latest/dg/snapstart.html) to learn more about AWS Lambda SnapStart
 

packages/aws-cdk-lib/aws-lambda/lib/function.ts

@@ -1554,7 +1554,7 @@ Environment variables can be marked for removal when used in Lambda@Edge by sett
       return undefined;
     }
 
-    // SnapStart does not support arm64 architecture, Amazon Elastic File System (Amazon EFS), or ephemeral storage greater than 512 MB.
+    // SnapStart does not support Amazon Elastic File System (Amazon EFS), or ephemeral storage greater than 512 MB.
     // SnapStart doesn't support provisioned concurrency either, but that's configured at the version level,
     // so it can't be checked at function set up time
     // SnapStart supports the Java 11 and Java 17 (java11 and java17) managed runtimes.
@@ -1565,10 +1565,6 @@ Environment variables can be marked for removal when used in Lambda@Edge by sett
       throw new Error(`SnapStart currently not supported by runtime ${props.runtime.name}`);
     }
 
-    if (props.architecture == Architecture.ARM_64) {
-      throw new Error('SnapStart is currently not supported on Arm_64');
-    }
-
     if (props.filesystem) {
       throw new Error('SnapStart is currently not supported using EFS');
     }

packages/aws-cdk-lib/aws-lambda/test/function.test.ts

@@ -3429,16 +3429,29 @@ describe('function', () => {
       })).toThrowError('SnapStart currently not supported by runtime nodejs18.x');
     });
 
-    test('arm64 validation for snapStart', () => {
+    test('arm64 function using snapStart', () => {
       const stack = new cdk.Stack();
-
-      expect(() => new lambda.Function(stack, 'MyLambda', {
+      //WHEN
+      new lambda.Function(stack, 'MyLambda', {
         code: lambda.Code.fromAsset(path.join(__dirname, 'handler.zip')),
         handler: 'example.Handler::handleRequest',
         runtime: lambda.Runtime.JAVA_11,
         architecture: lambda.Architecture.ARM_64,
         snapStart: lambda.SnapStartConf.ON_PUBLISHED_VERSIONS,
-      })).toThrowError('SnapStart is currently not supported on Arm_64');
+      });
+
+      //THEN
+      Template.fromStack(stack).hasResource('AWS::Lambda::Function', {
+        Properties:
+            {
+              Handler: 'example.Handler::handleRequest',
+              Runtime: 'java11',
+              Architectures: ['arm64'],
+              SnapStart: {
+                ApplyOn: 'PublishedVersions',
+              },
+            },
+      });
     });
 
     test('EFS validation for snapStart', () => {
@@ -3463,7 +3476,7 @@ describe('function', () => {
       })).toThrowError('SnapStart is currently not supported using EFS');
     });
 
-    test('arm64 validation for snapStart', () => {
+    test('ephemeral storage limit validation for snapStart', () => {
       const stack = new cdk.Stack();
 
       expect(() => new lambda.Function(stack, 'MyLambda', {