feat(aws-lambda): add support for XRay Tracing (#675)

aws · Sep 13, 2018 · b4435cc · b4435cc
1 parent c6c09bf
commit b4435cc
Show file tree

Hide file tree

Showing 3 changed files with 229 additions and 1 deletion.
diff --git a/packages/@aws-cdk/aws-lambda/README.md b/packages/@aws-cdk/aws-lambda/README.md
@@ -75,3 +75,18 @@ const fn = new lambda.Function(this, 'MyFunction', {
 ```
 See [the AWS documentation](https://docs.aws.amazon.com/lambda/latest/dg/dlq.html)
 to learn more about AWS Lambdas and DLQs.
+
+### Lambda with X-Ray Tracing 
+
+```ts
+import lambda = require('@aws-cdk/aws-lambda');
+
+const fn = new lambda.Function(this, 'MyFunction', {
+    runtime: lambda.Runtime.NodeJS810,
+    handler: 'index.handler'
+    code: lambda.Code.inline('exports.handler = function(event, ctx, cb) { return cb(null, "hi"); }'),
+    tracing: lambda.Tracing.Active
+});
+```
+See [the AWS documentation](https://docs.aws.amazon.com/lambda/latest/dg/lambda-x-ray.html)
+to learn more about AWS Lambda's X-Ray support.
diff --git a/packages/@aws-cdk/aws-lambda/lib/lambda.ts b/packages/@aws-cdk/aws-lambda/lib/lambda.ts
@@ -8,6 +8,26 @@ import { FunctionVersion } from './lambda-version';
 import { cloudformation, FunctionArn, FunctionName } from './lambda.generated';
 import { Runtime } from './runtime';
 
+/**
+ * X-Ray Tracing Modes (https://docs.aws.amazon.com/lambda/latest/dg/API_TracingConfig.html)
+ */
+export enum Tracing {
+    /**
+     * Lambda will respect any tracing header it receives from an upstream service.
+     * If no tracing header is received, Lambda will call X-Ray for a tracing decision.
+     */
+    Active,
+    /**
+     * Lambda will only trace the request from an upstream service
+     * if it contains a tracing header with "sampled=1"
+     */
+    PassThrough,
+    /**
+     * Lambda will not trace any request.
+     */
+    Disabled
+}
+
 export interface FunctionProps {
     /**
      * The source code of your Lambda function. You can point to a file in an
@@ -134,6 +154,13 @@ export interface FunctionProps {
      * @default SQS queue with 14 day retention period if `deadLetterQueueEnabled` is `true`
      */
     deadLetterQueue?: sqs.QueueRef;
+
+    /**
+     * Enable AWS X-Ray Tracing for Lambda Function.
+     *
+     * @default undefined X-Ray tracing disabled
+     */
+    tracing?: Tracing;
 }
 
 /**
@@ -216,6 +243,7 @@ export class Function extends FunctionRef {
             memorySize: props.memorySize,
             vpcConfig: this.configureVpc(props),
             deadLetterConfig: this.buildDeadLetterConfig(props),
+            tracingConfig: this.buildTracingConfig(props)
         });
 
         resource.addDependency(this.role);
@@ -332,4 +360,19 @@ export class Function extends FunctionRef {
             targetArn: deadLetterQueue.queueArn
         };
     }
+
+    private buildTracingConfig(props: FunctionProps) {
+        if (props.tracing === undefined || props.tracing === Tracing.Disabled) {
+            return undefined;
+        }
+
+        this.addToRolePolicy(new cdk.PolicyStatement()
+            .addActions('xray:PutTraceSegments', 'xray:PutTelemetryRecords')
+            .addAllResources());
+
+        return {
+            mode: Tracing[props.tracing]
+        };
+    }
+
 }
diff --git a/packages/@aws-cdk/aws-lambda/test/test.lambda.ts b/packages/@aws-cdk/aws-lambda/test/test.lambda.ts
@@ -1,4 +1,4 @@
-import { countResources, expect, haveResource } from '@aws-cdk/assert';
+import { countResources, expect, haveResource, ResourcePart } from '@aws-cdk/assert';
 import events = require('@aws-cdk/aws-events');
 import iam = require('@aws-cdk/aws-iam');
 import sqs = require('@aws-cdk/aws-sqs');
@@ -911,6 +911,176 @@ export = {
       test.done();
     },
 
+    'default function with Active tracing'(test: Test) {
+      const stack = new cdk.Stack();
+
+      new lambda.Function(stack, 'MyLambda', {
+          code: new lambda.InlineCode('foo'),
+          handler: 'index.handler',
+          runtime: lambda.Runtime.NodeJS610,
+          tracing: lambda.Tracing.Active
+      });
+
+      expect(stack).to(haveResource('AWS::IAM::Policy', {
+        "PolicyDocument": {
+          "Statement": [
+            {
+              "Action": [
+                "xray:PutTraceSegments",
+                "xray:PutTelemetryRecords"
+              ],
+              "Effect": "Allow",
+              "Resource": "*"
+            }
+          ],
+          "Version": "2012-10-17"
+        },
+        "PolicyName": "MyLambdaServiceRoleDefaultPolicy5BBC6F68",
+        "Roles": [
+          {
+            "Ref": "MyLambdaServiceRole4539ECB6"
+          }
+        ]
+      }));
+
+      expect(stack).to(haveResource('AWS::Lambda::Function', {
+        "Properties": {
+          "Code": {
+            "ZipFile": "foo"
+          },
+          "Handler": "index.handler",
+          "Role": {
+            "Fn::GetAtt": [
+              "MyLambdaServiceRole4539ECB6",
+              "Arn"
+            ]
+          },
+          "Runtime": "nodejs6.10",
+          "TracingConfig": {
+            "Mode": "Active"
+          }
+        },
+        "DependsOn": [
+          "MyLambdaServiceRole4539ECB6",
+          "MyLambdaServiceRoleDefaultPolicy5BBC6F68"
+        ]
+      }, ResourcePart.CompleteDefinition));
+
+      test.done();
+    },
+
+    'default function with PassThrough tracing'(test: Test) {
+      const stack = new cdk.Stack();
+
+      new lambda.Function(stack, 'MyLambda', {
+          code: new lambda.InlineCode('foo'),
+          handler: 'index.handler',
+          runtime: lambda.Runtime.NodeJS610,
+          tracing: lambda.Tracing.PassThrough
+      });
+
+      expect(stack).to(haveResource('AWS::IAM::Policy', {
+        "PolicyDocument": {
+          "Statement": [
+            {
+              "Action": [
+                "xray:PutTraceSegments",
+                "xray:PutTelemetryRecords"
+              ],
+              "Effect": "Allow",
+              "Resource": "*"
+            }
+          ],
+          "Version": "2012-10-17"
+        },
+        "PolicyName": "MyLambdaServiceRoleDefaultPolicy5BBC6F68",
+        "Roles": [
+          {
+            "Ref": "MyLambdaServiceRole4539ECB6"
+          }
+        ]
+      }));
+
+      expect(stack).to(haveResource('AWS::Lambda::Function', {
+        "Properties": {
+          "Code": {
+            "ZipFile": "foo"
+          },
+          "Handler": "index.handler",
+          "Role": {
+            "Fn::GetAtt": [
+              "MyLambdaServiceRole4539ECB6",
+              "Arn"
+            ]
+          },
+          "Runtime": "nodejs6.10",
+          "TracingConfig": {
+            "Mode": "PassThrough"
+          }
+        },
+        "DependsOn": [
+          "MyLambdaServiceRole4539ECB6",
+          "MyLambdaServiceRoleDefaultPolicy5BBC6F68"
+        ]
+      }, ResourcePart.CompleteDefinition));
+
+      test.done();
+    },
+
+    'default function with Disabled tracing'(test: Test) {
+      const stack = new cdk.Stack();
+
+      new lambda.Function(stack, 'MyLambda', {
+          code: new lambda.InlineCode('foo'),
+          handler: 'index.handler',
+          runtime: lambda.Runtime.NodeJS610,
+          tracing: lambda.Tracing.Disabled
+      });
+
+      expect(stack).notTo(haveResource('AWS::IAM::Policy', {
+        "PolicyDocument": {
+          "Statement": [
+            {
+              "Action": [
+                "xray:PutTraceSegments",
+                "xray:PutTelemetryRecords"
+              ],
+              "Effect": "Allow",
+              "Resource": "*"
+            }
+          ],
+          "Version": "2012-10-17"
+        },
+        "PolicyName": "MyLambdaServiceRoleDefaultPolicy5BBC6F68",
+        "Roles": [
+          {
+            "Ref": "MyLambdaServiceRole4539ECB6"
+          }
+        ]
+      }));
+
+      expect(stack).to(haveResource('AWS::Lambda::Function', {
+        "Properties": {
+          "Code": {
+            "ZipFile": "foo"
+          },
+          "Handler": "index.handler",
+          "Role": {
+            "Fn::GetAtt": [
+              "MyLambdaServiceRole4539ECB6",
+              "Arn"
+            ]
+          },
+          "Runtime": "nodejs6.10"
+        },
+        "DependsOn": [
+          "MyLambdaServiceRole4539ECB6"
+        ]
+      }, ResourcePart.CompleteDefinition));
+
+      test.done();
+    },
+
 };
 
 function newTestLambda(parent: cdk.Construct) {