Assign attributes to a transaction from a different async context

[tobobo]

Is your feature request related to a problem? Please describe.

I have a web server with long-running WebSocket connections that make periodic requests to external APIs out of my control. To monitor the health of my server, I am creating background transactions for each request so that I can make sure that request volumes are consistent and monitor for errors in the external APIs.

I had a solution that worked for this, but I'm in the process of switching to undici and I'm struggling to instrument these requests because of New Relic's strictness around async contexts.

Undici uses diagnostics_channel to allow for instrumentation (as you know from its usage in this repo to enable distributed tracing with undici). The undici:request:create channel can be used to start a background transaction, and undici:request:trailers can be used to end it.

However, undici:request:create occurs in a different async context from messages on the other channels for the same request. So, despite the fact that I can manually end the transaction from another async context (the undici:request:trailers handler), I can't add additional properties to it—in my case, response info like status code and headers.

Feature Description

I would like there to be atransaction.addCustomAttributes method to solve this. I am not sure why I can end a transaction from a different async context, but I can't add attributes. It also allows for more explicit control of the transaction attributes, rather than needing to trust the New Relic agent to attach attributes to the appropriate transaction via async context.

Describe Alternatives

I am temporarily extending undici's Client, but this is less flexible and requires more work in my application code to achieve instrumentation throughout my app.

I might be able to instrument this as a segment, but that would mean inconsistent instrumentation when comparing node's http/https modules and undici. Also, the only meaningful transaction encompassing this segment would be the long-running WebSocket, which could be a days-long transaction and presents its own challenges around maintaining async context.

Priority

Please help us better understand this feature request by choosing a priority from the following options:
[Really Want]

[workato-integration]

https://new-relic.atlassian.net/browse/NR-264186

[bizob2828]

@tobobo do you have a reproduction case? I could provide some possible solutions but hard to say from just reading your issue

[tobobo]

After playing around with it a bit more, I was able to accomplish it using Promise.withResolvers to make things a bit cleaner. I'm storing the promise resolvers so that I can return the necessary data back to the original async context so it's properly added to the transaction.

That said, I'm still not clear on why this kind of hack should be necessary—an escape hatch to use the transaction handle to add attributes would be much more explicit and would avoid the workaround of passing around promise resolvers to jump back to the original async context.

Here's a bare-bones example of turning undici requests in to New Relic transactions. In this example, it's not possible to set attributes on the transaction directly from onResponseHeaders or onDone, as these are in a different async context and New Relic can't retrieve the transaction from AsyncLocalStorage. Hence the workaround of using promise resolvers as a reference to the original async context.

const undici = require('undici');
const diagch = require('diagnostics_channel');
const newrelic = require('newrelic');

// Simple polyfill for Promise.withResolvers
// https://tc39.es/proposal-promise-with-resolvers/
Promise.withResolvers || (Promise.withResolvers = function withResolvers() {
  var a, b, c = new this(function (resolve, reject) {
    a = resolve;
    b = reject;
  });
  return {resolve: a, reject: b, promise: c};
});

const requestInstrumentationMap = new WeakMap();

function onRequestCreated({
  request,
}) {
  const hostname = new URL(request.origin).hostname;

  newrelic.startBackgroundTransaction(
    `ApiCall/${hostname}/${request.path}`,
    async () => {
      console.log('requesting', `${hostname}${request.path}`);
      
      const attributes = {};
      
      const { resolve: resolveTransactionPromise, reject: rejectTransactionPromise, promise: donePromise } = Promise.withResolvers();
      
      requestInstrumentationMap.set(request, {
        attributes,
        resolveTransactionPromise,
        rejectTransactionPromise,
      });
      
      try {
        await donePromise;
        console.log('request completed', `${hostname}${request.path}`);
      } catch (e) {
        console.log('request failed', `${hostname}${request.path}`, e);
        newrelic.noticeError(e, attributes);
      } finally {
        newrelic.addCustomAttributes(attributes);
        requestInstrumentationMap.delete(request);
      }
    },
  );
}

function onResponseHeaders({
  request,
  response,
}) {
  const requestInstrumentation = requestInstrumentationMap.get(request);
  if (!requestInstrumentation) return;
  
  Object.assign(requestInstrumentation.attributes, {
    'http.statusCode': response.statusCode,
  });
}

function onDone({ request }) {
  const requestInstrumentation = requestInstrumentationMap.get(request);
  if (!requestInstrumentation) return;
  
  if (requestInstrumentation.attributes['http.statusCode'] >= 400) {
    requestInstrumentation.rejectTransactionPromise(new Error(`HTTP error ${requestInstrumentation.attributes['http.statusCode']}`));
  } else {
    requestInstrumentation.resolveTransactionPromise();
  }
}

diagch.channel('undici:request:create').subscribe(onRequestCreated);
diagch.channel('undici:request:headers').subscribe(onResponseHeaders);
diagch.channel('undici:request:trailers').subscribe(onDone);

undici.request('http://example.com')
  .then(async () => {
    await undici.request('http://example.com/404');
  });

I also have a more fully-featured TypeScript version that I'm using on my server.

import * as diagch from 'diagnostics_channel';
import newrelic from 'newrelic';
import { type DiagnosticsChannel } from 'undici';
import { getGlobalLogger } from '../logger';
import { getCommonNewRelicAttributes } from './customEvents';
import { normalizeHeaders } from './normalizeHeaders';

const IGNORE_HOSTNAMES = ['newrelic.com'];

function headersToObject(headers: Buffer[] | string[]): Record<string, string> {
  const headersObject: Record<string, string> = {};

  for (let i = 0; i < headers.length; i += 2) {
    headersObject[headers[i].toString()] = headers[i + 1].toString();
  }

  return headersObject;
}

function instrumentUndici(): void {
  const requestInstrumentationMap = new WeakMap<
    any,
    {
      attributes: {
        ['request.type']: 'undici';
        ['request.uri']: string;
        ['request.method']: string;
        ['request.headers.host']: string;
        ['http.statusCode']?: number;
        ['http.statusText']?: string;
        [key: `request.headers.${string}`]: string;
        [key: `response.headers.${string}`]: string;
      };
      resolveTransactionPromise: () => void;
      rejectTransactionPromise: (e: Error) => void;
    }
  >();

  const log = getGlobalLogger().verbose;

  function onRequestCreated({
    request,
  }: DiagnosticsChannel.RequestCreateMessage): void {
    const hostname = request.origin
      ? new URL(request.origin).hostname
      : 'unknown';

    const isIgnoredHostname = IGNORE_HOSTNAMES.some((ignoredHostname) =>
      hostname.endsWith(ignoredHostname),
    );
    if (isIgnoredHostname) return;

    const uri = request.path;

    const attributes = {
      'request.type': 'undici' as const,
      'request.uri': uri,
      'request.method': request.method ?? 'unknown',
      'request.headers.host': hostname,
    };

    void newrelic.startBackgroundTransaction(
      `ApiCall/${hostname}/${request.path.split('?')[0]}`,
      async () => {
        newrelic.addCustomAttributes(getCommonNewRelicAttributes());
        newrelic.addCustomAttributes(attributes);

        const {
          resolve: resolveTransactionPromise,
          reject: rejectTransactionPromise,
          promise: donePromise,
        } = Promise.withResolvers();

        requestInstrumentationMap.set(request, {
          attributes,
          resolveTransactionPromise,
          rejectTransactionPromise,
        });

        log(`Requesting ${hostname}${request.path}`, attributes);

        try {
          await donePromise;
        } catch (err) {
          newrelic.noticeError(err as Error, attributes);
        } finally {
          newrelic.addCustomAttributes(attributes);
          requestInstrumentationMap.delete(request);
        }
      },
    );
  }

  function onResponseHeaders({
    request,
    response,
  }: DiagnosticsChannel.RequestHeadersMessage): void {
    const record = requestInstrumentationMap.get(request);
    if (!record) return;

    const { attributes } = record;

    const additionalAttributes = {
      'http.statusCode': response.statusCode,
      'http.statusText': response.statusText,
      ...(response.statusCode >= 400 &&
        normalizeHeaders(
          headersToObject(response.headers),
          `response.headers.`,
        )),
      ...(response.statusCode >= 400 &&
        normalizeHeaders(headersToObject(request.headers), 'request.headers.')),
    };

    Object.assign(attributes, additionalAttributes);
  }

  function onDone({
    request,
  }: DiagnosticsChannel.RequestTrailersMessage): void {
    const record = requestInstrumentationMap.get(request);
    if (!record) return;

    const { resolveTransactionPromise, rejectTransactionPromise, attributes } =
      record;

    if (attributes['http.statusCode'] && attributes['http.statusCode'] < 400) {
      resolveTransactionPromise();
      log(
        `Received success response from ${attributes['request.headers.host']}${attributes['request.uri']}`,
        attributes,
      );
    } else {
      log(
        `Received error response from ${attributes['request.headers.host']}${attributes['request.uri']}`,
        attributes,
      );
      rejectTransactionPromise(
        new Error(
          `Received error response from ${attributes['request.headers.host']}${attributes['request.uri']}`,
        ),
      );
    }
  }

  function subscribeToChannel(
    diagnosticChannel: string,
    onMessage: (message: any) => void,
  ): void {
    const channel = diagch.channel(diagnosticChannel);

    channel.subscribe(onMessage); // Bind the listener function directly to handle messages.
  }

  subscribeToChannel('undici:request:create', onRequestCreated);
  subscribeToChannel('undici:request:headers', onResponseHeaders);
  subscribeToChannel('undici:request:trailers', onDone);
}

instrumentUndici();

[bizob2828]

@tobobo typically you'd be calling undici within some webframework handle or at the very least if it's a background job wrapping those calls in newrelic.startBackgroundTransaction. So taking this code and wrapping it

newrelic.startBackgroundTransaction('transactionName', () => {
  undici.request('http://example.com')
    .then(async () => {
      await undici.request('http://example.com/404');
    });
})

Our async context manager is built around spans/segments and they definitely do not cross different transaction contexts. If this code is working then great but I'm having a hard time understanding the use case. It reads interesting to see you starting transactions within the undici hooks. In your undici hooks you could call newrelic.addCustomAttributes and they'll attach to the active transaction that was creating outside the hooks.

This example works just fine

const express = require('express')
const app = express()
const { PORT = 3000 } = process.env
const undici = require('undici')
const diagch = require('diagnostics_channel');

function onRequestCreated({
  request,
}) {
  debugger

  const attributes = {foo: 'bar', baz: 'bat'};
  newrelic.addCustomAttributes(attributes);
}

diagch.channel('undici:request:create').subscribe(onRequestCreated);


app.get('/undici', (req, res) => {
  undici.request('http://example.com')
  .then(async () => {
    await undici.request('http://example.com/404');
  }).then(() => {
    res.send(200)
  });
})