Understanding undici/async_hooks/New Relic

[tobobo]

I'm working on a tough issue—I'm switching to undici for a web server I'm running, but I need to be able to track undici reuqests as New Relic transactions so I can look at error rates when calling an external API.

I've attempted to use diagnostics_channel to tap in to undici's events, but there seems to be an issue—New Relic relies on async_hooks to get the current transaction, and it seems like the undici:request:create returns a different async ID than the other methods (undici:request:headers and undici:request:trailers have the same async ID, for example).

Is it a bug that the async ID is different for undici:request:create vs the other methods? Am I approaching this all wrong?

Here's an example of my instrumentation code below. It records the initial values and duration, but any attributes added after request creation are lost because New Relic can't find the transaction in the original execution context.

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

const IGNORE_HOSTNAMES = ['newrelic.com'];

function patchUndici(): void {
  const requestSpanMap = new WeakMap<
    any,
    {
      spanAttributes: Record<string, string>;
      resolveTransactionPromise: () => void;
      rejectTransactionPromise: (e: Error) => void;
    }
  >();

  function onRequestCreated({
    request,
  }: {
    request: any;
    response: any;
  }): void {
    const hostname = new URL(request.origin + request.path).hostname;

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

    const log = getGlobalLogger().verbose;

    const spanAttributes = {
      'request.uri': request.path,
      'request.method': request.method,
      'request.headers.host': hostname,
    };

    void newrelic.startBackgroundTransaction(
      `ApiCall/${hostname}/${request.path}`,
      async () => {
        newrelic.addCustomAttributes(getCommonNewRelicAttributes());
        newrelic.addCustomAttributes(spanAttributes);

        console.log('async id in start', async_hooks.executionAsyncId());

        await new Promise<void>((resolve, reject) => {
          const resolveTransactionPromise = resolve;
          const rejectTransactionPromise = reject;

          requestSpanMap.set(request, {
            spanAttributes,
            resolveTransactionPromise,
            rejectTransactionPromise,
          });

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

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

    const { spanAttributes } = record;

    console.log('async id in headers', async_hooks.executionAsyncId());

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

      Object.assign(spanAttributes, additionalAttributes);
      newrelic.addCustomAttributes(spanAttributes);

      const log = getGlobalLogger().verbose;
      if (response.statusCode >= 400) {
        log(
          `Received error response from ${request.hostname}${request.path}`,
          spanAttributes,
        );
        newrelic.noticeError(
          new Error(
            `Received error response from ${request.hostname}${request.path}`,
          ),
          spanAttributes,
        );
      } else {
        log(
          `Received success response from ${request.ostname}${request.path}`,
          spanAttributes,
        );
      }
    }
  }

  function onDone({ request }: { request: any }): void {
    const record = requestSpanMap.get(request);
    if (!record) return;

    console.log('async id in done', async_hooks.executionAsyncId());

    const { spanAttributes, resolveTransactionPromise } = record;

    newrelic.addCustomAttributes(spanAttributes);

    resolveTransactionPromise();
    requestSpanMap.delete(request);
  }

  function onError({ request, error }: { request: any; error: any }): void {
    const record = requestSpanMap.get(request);
    if (!record) return;

    const { spanAttributes, rejectTransactionPromise } = record;
    spanAttributes['error.message'] = error.message;

    newrelic.noticeError(error, spanAttributes);

    rejectTransactionPromise(error);
    requestSpanMap.delete(request);
  }

  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);
  subscribeToChannel('undici:request:error', onError);
}

patchUndici();

[mcollina]

@jsumners can you help with this?

[jsumners]

The question should be asked on the newrelic/node-newrelic repo. Preferably using JavaScript.

[tobobo]

Thanks. I’m not suggesting there’s an issue with undici, I was just hoping to see if anyone else has experience instrumenting it in this way. For anyone who finds this via Google, I’m going to use the alternative approach of extending the Client class.

This comes with its own challenges but based on your response this isn’t a very receptive venue for community discussion, so I’ll seek help elsewhere.

[jsumners]

I'm not sure which response you are suggesting is prohibiting community discussion. Here's what I understand about this thread:

1. A question about how to instrument Undici with New Relic was posted
2. I was asked if I could help by getting tagged with my personal GitHub username
3. I will not see this thread in my work notifications because of 2
4. I asked that the question, about New Relic, be taken to the appropriate repo so that 2 will be solved and so that the whole team can participate

No one said to not ask the question.

[tobobo]

Thanks @jsumners . Including that context in the original comment would have made your intentions clearer.

[tobobo]

Here's my understanding of what's going on based on a bit more research and thought:

• It is expected that the diagnostic channels events occur in different async contexts, based on the behavior of the dispatcher
• You can work around this without too much ugliness by using Promise.withResolvers to pass data back to the original async context and make New Relic happy.

Here's a bare-bones JS example:

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 requestMeta = {};
      
      const { resolve: resolveTransactionPromise, reject: rejectTransactionPromise, promise: donePromise } = Promise.withResolvers();
      
      requestInstrumentationMap.set(request, {
        requestMeta,
        resolveTransactionPromise,
        rejectTransactionPromise,
      });
      
      try {
        await donePromise;
        console.log('request completed', `${hostname}${request.path}`);
      } catch (e) {
        console.log('request failed', `${hostname}${request.path}`, e);
        noticeError(e, requestMeta);
      } finally {
        newrelic.addCustomAttributes(requestMeta);
        requestInstrumentationMap.delete(request);
      }
    },
  );
}

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

function onDone({ request }) {
  const requestInstrumentation = requestInstrumentationMap.get(request);
  if (!requestInstrumentation) return;
  
  if (requestInstrumentation.requestMeta['http.statusCode'] >= 400) {
    requestInstrumentation.rejectTransactionPromise(new Error(`HTTP error ${requestInstrumentation.requestMeta['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');
  });

And here's my full implementation that I'm using in my app (using TypeScript):

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();