AWS X-Ray and SQS with AWSTraceHeader

[toleabivol]

Hello,

recently it was announced that:

AWS X-Ray integrates with Amazon Simple Queue Service (Amazon SQS) to trace messages that are passed through an Amazon SQS queue.

What I fail to get is a nice map as shown here : https://docs.aws.amazon.com/xray/latest/devguide/xray-services-sqs.html

Specifically the arrow between SQS and custom sqs poller.

I have : API GW->LAMBDA->SQS->Custom python script with instrumented aws xray sdk that polls sqs

In the Lambda I have the following code:

...
        trace_entity = xray_recorder.get_trace_entity()
        trace_header = TraceHeader(root=trace_entity.trace_id, parent=trace_entity.id,
                                   sampled=trace_entity.sampled)
        message['MessageSystemAttributes'] = {
            'AWSTraceHeader': {
                'StringValue': trace_header.to_header_str(),
                'DataType': 'String'
            }
        }

        sqs_client.send_message(**message)

and in the sqs poller :

def _yielded_receive(self) -> List[SQSMessage]:
        return self.sqs.receive_messages(MessageAttributeNames=['All'],
                                             AttributeNames=['All'],
                                             VisibilityTimeout=self.VisibilityTimeout,
                                             WaitTimeSeconds=self.WaitTimeSeconds,
                                             MaxNumberOfMessages=self.MaxNumberOfMessages)

for raw_sqs_message in self._yielded_receive():
    trace_header = 
    TraceHeader.from_header_str(raw_sqs_message.attributes.get('AWSTraceHeader'))
    trace_id = trace_header.root
    sampled = trace_header.sampled

    segment = xray_recorder.begin_segment('SQS Receive',
                                                              traceid=trace_id,
                                                              sampling=sampled,
                                                              parent_id=trace_header.parent)

And with this I get in the x-ray map : API GW -> Lambda -> SQS Poller ; and SQS tied to Lambda separately.

[chanchiem]

Hey,

It looks like the expectation here is to have the SQS Receive poller node have a bidirectional edge between the SQS node and itself.

I think the issue is that the parent ID of the trace is set to the Lambda segment's ID. Specifically, in the trace header generation in the Lambda function here:

        trace_entity = xray_recorder.get_trace_entity()
        trace_header = TraceHeader(root=trace_entity.trace_id, parent=trace_entity.id,
                                   sampled=trace_entity.sampled)
...

You set parent of trace_header to be trace_entity.id, telling the poller that the parent entity of this outgoing trace stored in the message should be the Lambda segment.

I think what you want to do is acquire the subsegment that's generated as a result of the (presumably instrumented) SQS send message call. sqs_client.send_message(**message); however, since the subsegment is generated and ended during the send_message call, this wouldn't be possible.

By default, the AWS X-Ray SDK automatically propagates the entity trace header in the instrumented AWS SDK calls. You can read more about it in the Default HTTP Header section of the documentation.

My suggestion is to rely on the instrumented SQS client to do it for you using the default HTTP Header instead of trying to propagate the trace header yourself. Your lambda function would therefore look like this (assuming sqs_client is patched):

...
        sqs_client.send_message(**message)

When your receive side poller acquires the message, all you have to do is acquire the trace header string from the HTTP header and set the trace header information in the segment.

If you need more information and guidance on how to do this, please let me know.

Thanks,
Chan Chiem Jeffery Saeteurn

[lexxns]

Hi chanchiem, would you be able to provide me any guidance on how I would set the trace header information in a segment?

Thank you

[jconway45]

@chanchiem also interested in this.. I have the same example, when I receive a message from SQS it has the 'AWSTraceHeader' attribute available.. but the parent id is from the message sender not the subsegment generated by SQS. Would the X-Amzn-Trace-Id header contain the SQS subsegment id?

@lexxns let me know if you work anything out for this, I think we are trying to do the same thing..

[lexxns]

@jconway45 currently I don't believe it is possible in my situation, we are using lambdas triggered by SQS & the segment is created automatically before the lambda starts. This means I am unable to set trace information as the segment is immutable. So despite having access to the trace information I am unable to make use of it.

[jconway45]

@lexxns seems like it's impossible at the moment.

[chanchiem]

@lexxns Today it's impossible if the receive side is Lambda since Lambda creates the segment for you, and this generated segment is immutable in the Lambda Function.

@jconway45 The parent ID stored in the Trace Header should be the SQS subsegment ID. Are you making sure to use the patched SQS client to send the message to the queue?

[jconway45]

@chanchiem we're using the vanilla SDK client from AWS, nuget package (running in .net core)

By patched client, do you mean the AWS SDK after the xray integration?

We're using nuget version 3.3.102.13

[GrahamLea]

@chanchiem Can I please ask you to clarify your response to @lexxns: I think you are saying that xray currently cannot be used to trace from SQS to a SQS-triggered Lambda? It's just not supported? Do you know if it's on a roadmap somewhere? Thanks.

[chanchiem]

Sorry about the lack of correspondence. Since the issue was closed (and I had forgotten to reopen it), I was not been getting notifications from this issue.

@jconway45 Yes, that's what I mean. If you're running the AWS SDK client in .net, the .net SDK provides instrumentors that you can use to instrument your client. Link

@GrahamLea Today, you cant get the trace from SQS to SQS-triggered Lambda. This is a feature that has been asked from many customers and it is something that we have in our backlog. Please stay tuned

[gsogol]

Is there an update regarding SQS triggered lambdas?

[willarmiros]

This feature continues to be under development and we will update this thread when it's released. Thank you for your patience!

[victor-kironde]

Hi @willarmiros, Is there any update on this?

[mrgambal]

Bumping the topic: it's still an open and problematic issue.

[Sam-Martin]

I spent a little time working around this for lambda functions triggered by SQS queues and created a helper shim to create an entirely new segment inside the lambda function and add the trace id and request information from the incoming lambda function.
You can find it here: https://pypi.org/project/aws-xray-lambda-segment-shim/

This works for me, but be warned that:

1. This is a workaround and we're setting all sorts of undocumented (as far as I can see) attributes and monkey patching the origin on to the Segment object.
2. You will have duplicate traces: the one lambda creates and the one we create, the former won't be shown as triggered by SQS but the latter will
3. This is totally unsupported by AWS, if it works for you great. If it doesn't, raise an issue here

[stale]

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs in next 7 days. Thank you for your contributions.

[polamayster]

Is it possible to at least make it easier to patch on a receiving end similar to how it is documented for Java SDK: https://docs.aws.amazon.com/xray/latest/devguide/xray-services-sqs.html#xray-services-sqs-retrieving?
Being open for more than 3 years looks like this issue has no priority at all.

[polamayster]

Just in case someone missed it: https://aws.amazon.com/about-aws/whats-new/2022/11/aws-x-ray-trace-linking-event-driven-applications-amazon-sqs-lambda/
solves the issue and no code changes are needed

[srprash]

Hi all!
As pointed out by @polamayster, SQS tracing to AWS Lambda is now supported out of the box!
Thanks all of you for being patient. Please do give the new feature a try and share your feedback. I am closing out this issue now.