feat(stepfunctions): support X-Ray tracing (#10371)

packages/@aws-cdk/aws-stepfunctions/README.md

@@ -568,6 +568,22 @@ new stepfunctions.StateMachine(stack, 'MyStateMachine', {
 });
 ```
 
+## X-Ray tracing
+
+Enable X-Ray tracing for StateMachine:
+
+```ts
+const logGroup = new logs.LogGroup(stack, 'MyLogGroup');
+
+new stepfunctions.StateMachine(stack, 'MyStateMachine', {
+    definition: stepfunctions.Chain.start(new stepfunctions.Pass(stack, 'Pass')),
+    tracingEnabled: true
+});
+```
+
+See [the AWS documentation](https://docs.aws.amazon.com/step-functions/latest/dg/concepts-xray-tracing.html)
+to learn more about AWS Step Functions's X-Ray support.
+
 ## State Machine Permission Grants
 
 IAM roles, users, or groups which need to be able to work with a State Machine should be granted IAM permissions.

packages/@aws-cdk/aws-stepfunctions/lib/state-machine.ts

@@ -118,6 +118,13 @@ export interface StateMachineProps {
    * @default No logging
    */
   readonly logs?: LogOptions;
+
+  /**
+   * Specifies whether Amazon X-Ray tracing is enabled for this state machine.
+   *
+   * @default false
+   */
+  readonly tracingEnabled?: boolean;
 }
 
 /**
@@ -272,37 +279,13 @@ export class StateMachine extends StateMachineBase {
 
     this.stateMachineType = props.stateMachineType ? props.stateMachineType : StateMachineType.STANDARD;
 
-    let loggingConfiguration: CfnStateMachine.LoggingConfigurationProperty | undefined;
-    if (props.logs) {
-      const conf = props.logs;
-      loggingConfiguration = {
-        destinations: [{ cloudWatchLogsLogGroup: { logGroupArn: conf.destination.logGroupArn } }],
-        includeExecutionData: conf.includeExecutionData,
-        level: conf.level || 'ERROR',
-      };
-      // https://docs.aws.amazon.com/step-functions/latest/dg/cw-logs.html#cloudwatch-iam-policy
-      this.addToRolePolicy(new iam.PolicyStatement({
-        effect: iam.Effect.ALLOW,
-        actions: [
-          'logs:CreateLogDelivery',
-          'logs:GetLogDelivery',
-          'logs:UpdateLogDelivery',
-          'logs:DeleteLogDelivery',
-          'logs:ListLogDeliveries',
-          'logs:PutResourcePolicy',
-          'logs:DescribeResourcePolicies',
-          'logs:DescribeLogGroups',
-        ],
-        resources: ['*'],
-      }));
-    }
-
     const resource = new CfnStateMachine(this, 'Resource', {
       stateMachineName: this.physicalName,
       stateMachineType: props.stateMachineType ? props.stateMachineType : undefined,
       roleArn: this.role.roleArn,
       definitionString: Stack.of(this).toJsonString(graph.toGraphJson()),
-      loggingConfiguration,
+      loggingConfiguration: this.buildLoggingConfiguration(props),
+      tracingConfiguration: this.buildTracingConfiguration(props),
     });
 
     resource.node.addDependency(this.role);
@@ -324,7 +307,7 @@ export class StateMachine extends StateMachineBase {
    * Add the given statement to the role's policy
    */
   public addToRolePolicy(statement: iam.PolicyStatement) {
-    this.role.addToPolicy(statement);
+    this.role.addToPrincipalPolicy(statement);
   }
 
   /**
@@ -404,6 +387,49 @@ export class StateMachine extends StateMachineBase {
   public metricTime(props?: cloudwatch.MetricOptions): cloudwatch.Metric {
     return this.metric('ExecutionTime', props);
   }
+
+  private buildLoggingConfiguration(props: StateMachineProps) {
+    if (props.logs === undefined) {
+      return undefined;
+    }
+
+    // https://docs.aws.amazon.com/step-functions/latest/dg/cw-logs.html#cloudwatch-iam-policy
+    this.addToRolePolicy(new iam.PolicyStatement({
+      effect: iam.Effect.ALLOW,
+      actions: [
+        'logs:CreateLogDelivery',
+        'logs:GetLogDelivery',
+        'logs:UpdateLogDelivery',
+        'logs:DeleteLogDelivery',
+        'logs:ListLogDeliveries',
+        'logs:PutResourcePolicy',
+        'logs:DescribeResourcePolicies',
+        'logs:DescribeLogGroups',
+      ],
+      resources: ['*'],
+    }));
+
+    return {
+      destinations: [{ cloudWatchLogsLogGroup: { logGroupArn: props.logs.destination.logGroupArn } }],
+      includeExecutionData: props.logs.includeExecutionData,
+      level: props.logs.level || 'ERROR',
+    };
+  }
+
+  private buildTracingConfiguration(props: StateMachineProps) {
+    if (!props.tracingEnabled) {
+      return undefined;
+    }
+
+    this.addToRolePolicy(new iam.PolicyStatement({
+      actions: ['xray:PutTraceSegments', 'xray:PutTelemetryRecords'],
+      resources: ['*'],
+    }));
+
+    return {
+      enabled: true,
+    };
+  }
 }
 
 /**

packages/@aws-cdk/aws-stepfunctions/test/state-machine.test.ts

@@ -120,4 +120,42 @@ describe('State Machine', () => {
     });
   });
 
-});
+  test('tracing configuration', () => {
+    // GIVEN
+    const stack = new cdk.Stack();
+
+    // WHEN
+    new stepfunctions.StateMachine(stack, 'MyStateMachine', {
+      definition: stepfunctions.Chain.start(new stepfunctions.Pass(stack, 'Pass')),
+      tracingEnabled: true,
+    });
+
+    // THEN
+    expect(stack).toHaveResource('AWS::StepFunctions::StateMachine', {
+      DefinitionString: '{"StartAt":"Pass","States":{"Pass":{"Type":"Pass","End":true}}}',
+      TracingConfiguration: {
+        Enabled: true,
+      },
+    });
+
+    expect(stack).toHaveResource('AWS::IAM::Policy', {
+      PolicyDocument: {
+        Statement: [{
+          Action: [
+            'xray:PutTraceSegments',
+            'xray:PutTelemetryRecords',
+          ],
+          Effect: 'Allow',
+          Resource: '*',
+        }],
+        Version: '2012-10-17',
+      },
+      PolicyName: 'MyStateMachineRoleDefaultPolicyE468EB18',
+      Roles: [
+        {
+          Ref: 'MyStateMachineRoleD59FFEBC',
+        },
+      ],
+    });
+  });
+});