Is there a way to achieve correct segment parent-child relationship with batch utility and aws sdk clients captured with the tracer utility?

[artkhod]

Hi, as title states I have found myself stuck on an issue of producing AWS X-ray segments per each record processed with a batch utility.

As documentation correctly suggests Batch processing | Tracing with AWS X-Ray, when working inside of a recordHandler, it makes sense to work with the created subsegment directly since they provide an object with which we can work independently for every invocation of the recordHandler.

On the other hand, other utilities like logger would require us to spawn a child logger to eliminate a possibility of appended keys' override by handlers which were invoked concurrently.

This brings us to the issue: what if my recordHandler uses aws sdk client patched by tracer?

import {
  BatchProcessor,
  EventType,
  processPartialResponse,
} from '@aws-lambda-powertools/batch';
import { Tracer } from '@aws-lambda-powertools/tracer';
import { captureLambdaHandler } from '@aws-lambda-powertools/tracer/middleware';
import middy from '@middy/core';
import type { SQSEvent, SQSHandler, SQSRecord } from 'aws-lambda';
import { S3Client, PutObjectCommand } from "@aws-sdk/client-s3";

const processor = new BatchProcessor(EventType.SQS);
const tracer = new Tracer({ serviceName: 'serverlessAirline' });
const s3Client = tracer.captureAWSv3Client(
  new S3Client({ region: 'eu-west-1' })
);

const recordHandler = async (record: SQSRecord): Promise<void> => {
  const subsegment = tracer.getSegment()?.addNewSubsegment('### recordHandler');
  const command = new PutObjectCommand({
    Body: record.body,
    Bucket: "examplebucket",
    Key: "objectkey"
  }); 
  await s3Client.send(command)
  subsegment?.close(); 
};

export const handler: SQSHandler = middy(async (event: SQSEvent, context) =>
  processPartialResponse(event, recordHandler, processor, {
    context,
  })
).use(captureLambdaHandler(tracer));

Ideally, we would want to be able to produce the following structure, however there is no way for tracer to know to which segment it should attach a new subsegment to correctly show the connections.

## handler
├── ### recordHandler
│   ├── S3  PutObject
├── ### recordHandler
│   ├── S3  PutObject
├── ### recordHandler
│   ├── S3  PutObject
├── ### recordHandler
│   ├── S3  PutObject

The use case for such x-ray trace would be that it represents the reality in which recordHandlers run concurrently + if recordHandler segment would put an annotation with some sort of message id, then it would allow for queries like: Which segments have failed while processing a message with id "123".

However, as you might imagine, depending on how long each recorHandler takes to execute the x-ray trace may take a shape of an arbitrary structure which would make it hard to interpret parent/child relationships of trace segment.

So, in summary, do you have an idea of how such tracing may be accomplished and whether it's even possible?

[artkhod]

HI, I have come up with an idea how I would be able to solve it, yet the method used from aws-xray-sdk-core Powertools tracer provider was restricted to 1 parameter instead of its original 2.

I wouldn't mind if it was on the tracer itself, but I thought that the provider object would allow me to use the provider (in this case aws-xray-sdk) the way it was written x-ray-sdk.

[dreamorosi]

Hey, sorry for the delay - I spent some time thinking about your question and unfortunately I didn't come up with anything very useful.

The use case you're describing is valid however you're hitting a limitation of the X-Ray SDK for Node.js when working with promises. The X-Ray SDK uses on an older version of context tracking based on this module, which in turns uses Node.js async_hooks instead of the more modern AsyncLocalStorage.

Because of this, the whole context tracking for segments breaks down when running parallel async code - which is also why we added the workaround you linked to the docs.

When it comes to patched AWS SDK clients however the parent segment is resolved automatically based on the currently active one, which breaks the workaround.

With this in mind, I am unsure how to get around this beyond suggesting to not use tracer.captureAWSv3Client() and creating your own subsegment for these operations manually - which is less than ideal.