Instrumenting an AWS Amplify Web Application Using AWS X-Ray
Alexander WitteIntroduction
Modern mobile and web applications increasingly rely on cloud and serverless technologies to expedite development and improve scaling, but this can add complexity to the provisioning stage. AWS Amplify is a framework that simplifies the provisioning of the backend and frontend services and is designed specifically for web and mobile app development.
"…a complete solution that lets frontend web and mobile developers easily build, ship, and host full-stack applications on AWS, with the flexibility to leverage the breadth of AWS services as use cases evolve. No cloud expertise needed." – Amazon AWS
The simplicity that Amplify provides comes with tradeoffs. As the success and your user base grow, the need for holistic observability becomes vital. AWS offers X-Ray for analyzing distributed application tracing. However, Amplify does not provide the tooling to easily integrate X-Ray into the Amplify framework.
In this post, we’ll share code snippets for instrumenting Amplify components: a Lambda function, an AppSync API, and frontend code to generate traces. The goal behind this is to annotate a trace allowing us to filter our traces by some information that we define. We will implement Amplify and X-Ray and then filter traces by an email address.
A REAL Quick Primer on AWS X-Ray
X-Ray works on a concept of traces. A trace tracks the path of a request through your application. Each trace is broken up into segments which contains information about the work being performed by compute resources in your environment. Segments can also be broken up in subsegments if some more granular measurements need to be made, or for calls to external systems. Alongside those, traces also con contain metadata and annotations. Metadata are key:value pairs you can use to describe components of your trace, but cannot be used for searching. Annotations are also key:value pairs but you can use these for searching within your trace. There is also a service map that allows us to look at a visual representation of our application.
Please review the AWS X-Ray Documentation to learn more about X-Ray in general.
Some AWS Services can be enabled to support X-Ray. These services have varying X-Ray functionality but will generate the trace ID (and header) for us or propagate existing traces to downstream components. Quite commonly we can lean on the API layer of our application as the starting point to generate this X-Ray Trace. In our case however, we’d like to generate the trace header ourselves beforehand and then send that existing trace to the API. This is the case if we’d like to annotate our trace within the front end code. This is referrenced within AWS’ sample Scorekeep X-Ray project. This blog is an adaptation of that idea to work within the Amplify framework.
Enabling X-Ray for AWS Lambda within Amplify
AWS Lambda requires X-Ray to be enabled to allow it to submit segment information for a trace. The Amplify CLI does not give us an option to do this, so we’ll need to get our hands dirty and update the Cloudformation code that Amplify had generated. Head over to your Lambda functions Cloudformation template and make these updates to the Lambda configuration and the IAM Policy.
1"Resources": {
2 "LambdaFunction": {
3 "Type": "AWS::Lambda::Function",
4 "Properties": {
5 "TracingConfig": {
6 "Mode": "Active"
7 },
8...
9 "lambdaexecutionpolicy": {
10 "Type": "AWS::IAM::Policy",
11 "PolicyDocument": {
12 "Statement": [
13 {
14 "Effect": "Allow",
15 "Action": [
16 "xray:PutTraceSegments",
17 "xray:PutTelemetryRecords"
18 ],
19 "Resource": "*"
20 }
21...
amplify/backend/function/functionName-cloudformation-template.json
When you push this change you will notice your Lambda function is now enabled for X-Ray.
Enabling X-Ray for an AWS Appsync API within Amplify
For our Appsync API we will also need to enable X-Ray. We can do this through an Amplify override. To generate the override file we type this command:
amplify override api
This will generate the amplify/backend/api/override.ts file. Let’s edit it:
1import { AmplifyApiGraphQlResourceStackTemplate, AmplifyProjectInfo } from '@aws-amplify/cli-extensibility-helper';
2
3export function override(resources: AmplifyApiGraphQlResourceStackTemplate, amplifyProjectInfo: AmplifyProjectInfo) {
4 resources.api.GraphQLAPI.xrayEnabled = true;
5}
Don’t forget to push your changes when you’re done.
Instrumenting the Front End
This is where the real fun begins. We’ll need to do a couple things here:
- Update our project’s identity pool auth IAM role: Amplify uses AWS Cognito to provide our application with IAM credentials for AWS services. Amplify will stand up a Cognito identity pool for us but we will need to update the IAM policy associated to the role our application will assume for authenticated users. The policy will allow access to X-Ray.
- Create X-Ray trace & segment generation code: We’ll need to create some javascript functions to create an X-Ray trace, segment and annotate the trace using the email address obtained from the Cognito User Pool token.
- Generate the X-Ray trace ID and segment within our API calls: Here we call the X-Ray functions we had created earlier to generate (and close) X-Ray segments for our Appsync/Graphql calls.
Update our project’s identity pool auth IAM role
If you’ve configured Cognito auth within your Amplify project, Amplify should have created a Cognito identity pool for you. This pool should have an IAM Role associated to it for authenticated users. We just need to update the policy for this role to allow X-Ray access. We can do that with an Amplify project override.
amplify override project
This will create the backend/amplify/cloudformation/override.ts file. Edit it with this.
1import { AmplifyProjectInfo, AmplifyRootStackTemplate } from '@aws-amplify/cli-extensibility-helper';
2
3export function override(resources: AmplifyRootStackTemplate, amplifyProjectInfo: AmplifyProjectInfo) {
4 const authRole = resources.authRole;
5
6 const basePolicies = Array.isArray(authRole.policies)
7 ? authRole.policies
8 : [authRole.policies];
9
10 authRole.policies = [
11 ...basePolicies,
12 {
13 policyName: "amplify-permissions-xray",
14 policyDocument: {
15 Version: "2012-10-17",
16 Statement: [
17 {
18 Resource: "*",
19 Action: ["xray:PutTraceSegments","xray:PutTelemetryRecords"],
20 Effect: "Allow",
21 },
22 ],
23 },
24 },
25 ];
26}
This takes the existing policy and adds some new permissions to it. Push the changes.
Create X-Ray trace & segment generation code
First we’ll need to install some libraries. Firstly the Javascript X-Ray SDK and then get-random-values to generate some values for our Trace ID.
npm i @aws-sdk/client-xray get-random-values
I’ve created an X-Ray.js file to contain all the X-Ray related functions. Again, I stole a lot of this from AWS’ demo Scorekeep Application so please go there to learn more.
1import { XRayClient, PutTraceSegmentsCommand } from "@aws-sdk/client-xray";
2import getRandomValues from "get-random-values";
3import { Auth } from 'aws-amplify'
4import awsconfig from "./aws-exports";
5
6const getHexTime = () => {
7 return Math.round(new Date().getTime() / 1000).toString(16);
8};
9
10const getEpochTime = () => {
11 return new Date().getTime() / 1000;
12};
13
14const getHexId = (length) => {
15 const bytes = new Uint8Array(length);
16 getRandomValues(bytes);
17 let hex = "";
18 for (let i = 0; i < bytes.length; i++) {
19 hex += bytes[i].toString(16);
20 }
21 return hex.substring(0, length);
22};
23
24export const getTraceHeader = (segment) => {
25 return "Root=" + segment.trace_id + ";Parent=" + segment.id + ";Sampled=1";
26};
27
28export const beginSegment = (user) => {
29 let segment = {};
30 const traceId = "1-" + getHexTime() + "-" + getHexId(24);
31 const id = getHexId(16);
32 const startTime = getEpochTime();
33
34 segment.trace_id = traceId;
35 segment.id = id;
36 segment.start_time = startTime;
37 segment.end_time = startTime;
38 segment.name = "Web-Application";
39 segment.in_progress = true;
40 segment.annotations = {
41 "email": user.email,
42 };
43 let documents = [];
44 documents[0] = JSON.stringify(segment);
45 putDocuments(documents);
46 return segment;
47};
48
49const putDocuments = async (documents) => {
50 try {
51 let params = {
52 TraceSegmentDocuments: documents,
53 };
54 const credentials = await Auth.currentCredentials();
55 const client = new XRayClient({
56 region: awsconfig.aws_project_region,
57 credentials: Auth.essentialCredentials(credentials),
58 });
59 const command = new PutTraceSegmentsCommand(params);
60 const response = await client.send(command);
61 return response;
62
63 } catch(err) {
64 console.log(err);
65 }
66};
67
68export const endSegment = async (segment) => {
69 const endTime = getEpochTime();
70 segment.end_time = endTime;
71 segment.in_progress = false;
72 let documents = [];
73 documents[0] = JSON.stringify(segment);
74 putDocuments(documents);
75};
Add the Trace Header to our API calls
OK, now we just need to call these functions when making Appsync API calls. Upon sending an API request to Appsync, our application will call our functions exported from X-Ray.js to generate our Trace ID header and a segment. It will place the trace and segment identifiers in the X-Amzn-Trace-Id that Appsync in listening for. After the call is made to Appsync we’ll end the segment (allowing us to time the call).
We’re also annotating the trace with a string that we can then use to filter traces with.
1import { beginSegment, getTraceHeader, endSegment } from "./xray";
2
3...
4
5const App = () => {
6
7 // ... other React code ommitted for simplicity
8
9 let emailAddress = "myemail@email.com" // in the real world we'd assign this dynamically from our Cognito token
10
11 async function someAPICall() {
12 try {
13 let segment = beginSegment({email: emailAddress});
14 let response = await API.graphql(graphqlOperation(myAPICall, {input: {}}), // "myAPICall" is defined in your graphql schema file
15 {"X-Amzn-Trace-Id": getTraceHeader(segment)}
16 )
17 endSegment(segment);
18 }
19 catch (err) {
20 console.log(err)
21 }
22 }
23
24 return (
25 <div>
26 <button onClick={someAPICall}>Button that calls Appsync</button>
27 </div>
28 )
29}
30
31export default withAuthenticator(App)
If you pop into your service map now in the Cloudfront section of the AWS Console you should be able to see your application. You should be able to filter traces on your email address as well.
Conclusion
To conclude, I’ve shown some an example of how we can enable a Lambda function, an Appsync API and some front end code for X-Ray within an Amplify application. This blog was born out frustration trying to figure this out stuff out. I hope this has been helpful!