Aws-xray-sdk-node: XRay has problems recording subsegments when await keyword is used

[wilson-burea]

I have a method with a code similar to the code below were I used await method when calling pool.query() which is a pg pool method. The subsegment that is after the pool.query() never gets recorded but the code execute correctly. It returns that actual rows without any problem.

import {Pool} from 'pg';
import format from "pg-format"

const poolConfig = {
    database: <DATABASE>
host: <HOST>,
port: <PORT>,
user: <USER>,
password: <PASSWORD>
}

pool = new Pool(poolConfig)


export const getDynamicLinkStats = async () => {
    return await AWSXRay.captureAsyncFunc('Query-creation-and-postgres-call', async 
        function(subsegment){
            const textQuery = <QUERY>;

            subsegment.addAnnotation('test', '1');

            try {
                const { rows } = await pool.query(textQuery);
                subsegment.addAnnotation('test2', '2');

                return rows;
            } catch(error) {
                return [];
            } finally {
                subsegment.close();
            }
    });
};

[willarmiros]

Hi @wilson-burea,
Thank you for following up on GitHub following our external conversation. I'll include my current findings for public visibility.

This issue is more with Lambda runtimes than with X-Ray. It is a known issue wherein the X-Ray segment generated by the Lambda worker is sent before the Lambda finishes execution. I believe this is because the Lambda runtime waits for the event loop to be empty before terminating by default, which means it will process all callbacks/unresolved promises before terminating. The X-Ray segment is sent before this, meaning it cannot pick up changes in such callbacks. I've cut a ticket to Lambda to investigate this and will follow up on this ticket if it's flipped back to me.

I am investigating this issue with the Lambda and will update if they come back with anything.

[AdonousTech]

In this situation you may need to use promises on the back-end and async await on the calling code. Have you tried wrapping captureAsyncFunc in a promise? Then you can use basic promise syntax to call the pool.query method (e.g. pool.query().then(...)). The caller of the getDynamicLinkStats method can use async await pattern. There's no need to use await in the back-end/implementation code.

For example, here is a captureAsyncFunc wrapped in a promise in some back-end class:

import AWSXRay = require('aws-xray-sdk-core');
export class TestAsync {
    static mockAsync(): Promise<boolean> {
        return new Promise((resolve, reject) => {
            AWSXRay.captureAsyncFunc('test-async-capture', (subsegment: AWSXRay.Subsegment | undefined) => {
                console.log('ready to capture...');
                simAsync().then((result) => {  
                    console.log('capture complete...');
                    subsegment?.addMetadata('critical-event', {
                        key1: 'value1',
                        key2: 'value2'
                    })
                    subsegment?.close();
                    resolve(true);
                })
            })

            function simAsync(): Promise<boolean> {
                return new Promise((resolve, reject) => {
                    setTimeout(() => {
                      resolve(true);  
                    }, 1500);
                })
            }
        })
    }
}

The implementation above executes the captureAsyncFunc properly.

I only care about the benefits of async await in the main handler code. There it is convenient to use a synchronous pattern and return the response when complete.

import AWSXRay = require('aws-xray-sdk-core');
import { TestAsync } from './lib/test-async';
const AWS = AWSXRay.captureAWS(require('aws-sdk'));

export const handler = async (event: any) => {
    console.log('received event :: ', event);
    const response = {
        statusCode: 200,
        body: JSON.stringify('Hello from Lambda!')
    };

   // call the back-end method which will execute and complete the
  // AWSXRay.captureFuncAsync call before returning
    await TestAsync.mockAsync();

    return response;
}

[willarmiros]

Thank you for the workaround @AdonousTech, I tested it in Lambda and verified the trace contains the expected data. @wilson-burea Have you considered adding a promise wrapper similar to the one Matt proposed around your third party async calls?

[willarmiros]

Hi @wilson-burea,

Are you still experiencing this issue? I tried this sample code, although I was getting ECONNREFUSED errors since I didn't have a postgres database set up, so I used a simulated async query function instead. I'm using node 12.x and aws-xray-sdk 3.0.0 in my example.

var AWSXRay = require('aws-xray-sdk-core');
const pg = require('pg');

var pool = new pg.Pool();

exports.handler = async (event, context) => {
    return getDynamicLinkStats();
};



const getDynamicLinkStats = async () => {
    return await AWSXRay.captureAsyncFunc('Query-creation-and-postgres-call',  
        async function(subsegment) {
            const textQuery = '';

            subsegment.addAnnotation('test', '1');

            try {
                // const { rows } = await pool.query(textQuery);  // giving me errors
                const rows = await fakeHandleQuery(textQuery);
                subsegment.addAnnotation('test2', '2');
                subsegment.close();
                return rows;
            } catch(error) {
                console.log('got error: ', error);
                subsegment.close(error);
                return [];
            }
    });
};

async function fakeHandleQuery(event) {
  let promise = new Promise((resolve, reject) => {
    setTimeout(() => resolve("done!"), 1000);
  });

  let result = await promise;
  return result;
}

It's possible that using cls-hooked helped resolve this issue with async/await. Please let me know if you're still encountering this issue.