feat(instrumentation-undici): Add responseHook (#2356)

open-telemetry · Aug 9, 2024 · 60a99c9 · 60a99c9
1 parent 167dced
commit 60a99c9
Show file tree

Hide file tree

Showing 5 changed files with 71 additions and 7 deletions.
diff --git a/plugins/node/instrumentation-undici/README.md b/plugins/node/instrumentation-undici/README.md
@@ -51,11 +51,12 @@ Undici instrumentation has few options available to choose from. You can set the
 
 | Options                                                                                                                                                | Type                    | Description                                                                                                                                                                                                                                                  |
 | ------------------------------------------------------------------------------------------------------------------------------------------------------ | ----------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------ |
-| [`ignoreRequestHook`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L63)       | `IgnoreRequestFunction` | Undici instrumentation will not trace all incoming requests that matched with custom function.                                                                                                                                                               |
-| [`requestHook`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L65)             | `RequestHookFunction`   | Function for adding custom attributes before request is handled.                                                                                                                                                                                             |
-| [`startSpanHook`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L67)           | `StartSpanHookFunction` | Function for adding custom attributes before a span is started.                                                                                                                                                                                              |
-| [`requireParentforSpans`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L69)   | `Boolean`               | Require a parent span is present to create new span for outgoing requests.                                                                                                                                                                                   |
-| [`headersToSpanAttributes`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L71) | `Object`                | List of case insensitive HTTP headers to convert to span attributes. Headers will be converted to span attributes in the form of `http.{request\|response}.header.header-name` where the name is only lowercased, e.g. `http.response.header.content-length` |
+| [`ignoreRequestHook`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L73)       | `IgnoreRequestFunction` | Undici instrumentation will not trace all incoming requests that matched with custom function.                                                                                                                                                               |
+| [`requestHook`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L75)             | `RequestHookFunction`   | Function for adding custom attributes before request is handled.                                                                                                                                                                                             |
+| [`responseHook`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L77)            | `ResponseHookFunction`  | Function for adding custom attributes after the response headers are received.                                                                                                                                                                               |
+| [`startSpanHook`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L79)           | `StartSpanHookFunction` | Function for adding custom attributes before a span is started.                                                                                                                                                                                              |
+| [`requireParentforSpans`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L81)   | `Boolean`               | Require a parent span is present to create new span for outgoing requests.                                                                                                                                                                                   |
+| [`headersToSpanAttributes`](https://github.com/open-telemetry/opentelemetry-js-contrib/blob/main/plugins/node/instrumentation-undici/src/types.ts#L83) | `Object`                | List of case insensitive HTTP headers to convert to span attributes. Headers will be converted to span attributes in the form of `http.{request\|response}.header.header-name` where the name is only lowercased, e.g. `http.response.header.content-length` |
 
 ### Observations
 

diff --git a/plugins/node/instrumentation-undici/src/types.ts b/plugins/node/instrumentation-undici/src/types.ts
@@ -41,6 +41,7 @@ export interface UndiciRequest {
 export interface UndiciResponse {
   headers: Buffer[];
   statusCode: number;
+  statusText: string;
 }
 
 export interface IgnoreRequestFunction<T = UndiciRequest> {
@@ -51,18 +52,29 @@ export interface RequestHookFunction<T = UndiciRequest> {
   (span: Span, request: T): void;
 }
 
+export interface ResponseHookFunction<
+  RequestType = UndiciResponse,
+  ResponseType = UndiciResponse
+> {
+  (span: Span, info: { request: RequestType; response: ResponseType }): void;
+}
+
 export interface StartSpanHookFunction<T = UndiciRequest> {
   (request: T): Attributes;
 }
 
 // This package will instrument HTTP requests made through `undici` or  `fetch` global API
 // so it seems logical to have similar options than the HTTP instrumentation
-export interface UndiciInstrumentationConfig<RequestType = UndiciRequest>
-  extends InstrumentationConfig {
+export interface UndiciInstrumentationConfig<
+  RequestType = UndiciRequest,
+  ResponseType = UndiciResponse
+> extends InstrumentationConfig {
   /** Not trace all outgoing requests that matched with custom function */
   ignoreRequestHook?: IgnoreRequestFunction<RequestType>;
   /** Function for adding custom attributes before request is handled */
   requestHook?: RequestHookFunction<RequestType>;
+  /** Function called once response headers have been received */
+  responseHook?: ResponseHookFunction<RequestType, ResponseType>;
   /** Function for adding custom attributes before a span is started */
   startSpanHook?: StartSpanHookFunction<RequestType>;
   /** Require parent to create span for outgoing requests */

diff --git a/plugins/node/instrumentation-undici/src/undici.ts b/plugins/node/instrumentation-undici/src/undici.ts
@@ -354,6 +354,14 @@ export class UndiciInstrumentation extends InstrumentationBase<UndiciInstrumenta
     };
 
     const config = this.getConfig();
+
+    // Execute the response hook if defined
+    safeExecuteInTheMiddle(
+      () => config.responseHook?.(span, { request, response }),
+      e => e && this._diag.error('caught responseHook error: ', e),
+      true
+    );
+
     const headersToAttribs = new Set();
 
     if (config.headersToSpanAttributes?.responseHeaders) {

diff --git a/plugins/node/instrumentation-undici/test/fetch.test.ts b/plugins/node/instrumentation-undici/test/fetch.test.ts
@@ -142,6 +142,9 @@ describe('UndiciInstrumentation `fetch` tests', function () {
         requestHook: () => {
           throw new Error('requestHook error');
         },
+        responseHook: () => {
+          throw new Error('responseHook error');
+        },
         startSpanHook: () => {
           throw new Error('startSpanHook error');
         },
@@ -213,6 +216,12 @@ describe('UndiciInstrumentation `fetch` tests', function () {
             req.headers.push('x-requested-with', 'undici');
           }
         },
+        responseHook: (span, { response }) => {
+          span.setAttribute(
+            'test.response-hook.attribute',
+            response.statusText
+          );
+        },
         startSpanHook: request => {
           return {
             'test.hook.attribute': 'hook-value',
@@ -281,6 +290,11 @@ describe('UndiciInstrumentation `fetch` tests', function () {
         'hook-value',
         'startSpanHook is called'
       );
+      assert.strictEqual(
+        span.attributes['test.response-hook.attribute'],
+        'OK',
+        'responseHook is called'
+      );
     });
 
     it('should not create spans without parent if required in configuration', async function () {

diff --git a/plugins/node/instrumentation-undici/test/undici.test.ts b/plugins/node/instrumentation-undici/test/undici.test.ts
@@ -169,6 +169,12 @@ describe('UndiciInstrumentation `undici` tests', function () {
             req.headers.push('x-requested-with', 'undici');
           }
         },
+        responseHook: (span, { response }) => {
+          span.setAttribute(
+            'test.response-hook.attribute',
+            response.statusText
+          );
+        },
         startSpanHook: request => {
           return {
             'test.hook.attribute': 'hook-value',
@@ -357,6 +363,11 @@ describe('UndiciInstrumentation `undici` tests', function () {
         'hook-value',
         'startSpanHook is called'
       );
+      assert.strictEqual(
+        span.attributes['test.response-hook.attribute'],
+        'OK',
+        'responseHook is called'
+      );
     });
 
     it('should create valid spans for "fetch" method', async function () {
@@ -417,6 +428,11 @@ describe('UndiciInstrumentation `undici` tests', function () {
         'hook-value',
         'startSpanHook is called'
       );
+      assert.strictEqual(
+        span.attributes['test.response-hook.attribute'],
+        'OK',
+        'responseHook is called'
+      );
     });
 
     it('should create valid spans for "stream" method', async function () {
@@ -485,6 +501,11 @@ describe('UndiciInstrumentation `undici` tests', function () {
         'hook-value',
         'startSpanHook is called'
       );
+      assert.strictEqual(
+        span.attributes['test.response-hook.attribute'],
+        'OK',
+        'responseHook is called'
+      );
     });
 
     it('should create valid spans for "dispatch" method', async function () {
@@ -561,6 +582,11 @@ describe('UndiciInstrumentation `undici` tests', function () {
         'hook-value',
         'startSpanHook is called'
       );
+      assert.strictEqual(
+        span.attributes['test.response-hook.attribute'],
+        'OK',
+        'responseHook is called'
+      );
     });
 
     it('should create valid spans even if the configuration hooks fail', async function () {
@@ -576,6 +602,9 @@ describe('UndiciInstrumentation `undici` tests', function () {
         requestHook: () => {
           throw new Error('requestHook error');
         },
+        responseHook: () => {
+          throw new Error('responseHook error');
+        },
         startSpanHook: () => {
           throw new Error('startSpanHook error');
         },