fix: context activation for HTTP instrumentation's outgoing requests

packages/opentelemetry-instrumentation-http/src/http.ts

@@ -16,6 +16,7 @@
 import {
   SpanStatusCode,
   context,
+  Context,
   propagation,
   Span,
   SpanKind,
@@ -43,6 +44,7 @@ import {
   ParsedRequestOptions,
   ResponseEndArgs,
 } from './types';
+import { bindEmitter } from './patch';
 import * as utils from './utils';
 import { VERSION } from './version';
 import {
@@ -276,7 +278,9 @@ export class HttpInstrumentation extends InstrumentationBase<Http> {
     component: 'http' | 'https',
     request: http.ClientRequest,
     options: ParsedRequestOptions,
-    span: Span
+    span: Span,
+    requestContext: Context,
+    parentContext: Context
   ): http.ClientRequest {
     const hostname =
       options.hostname ||
@@ -291,7 +295,7 @@ export class HttpInstrumentation extends InstrumentationBase<Http> {
       this._callRequestHook(span, request);
     }
 
-    request.on(
+    request.prependListener(
       'response',
       (response: http.IncomingMessage & { aborted?: boolean }) => {
         const responseAttributes = utils.getOutgoingRequestAttributesOnResponse(
@@ -303,7 +307,10 @@ export class HttpInstrumentation extends InstrumentationBase<Http> {
           this._callResponseHook(span, response);
         }
 
-        context.bind(response);
+        bindEmitter(response, event =>
+          event === 'data' ? requestContext : parentContext
+        );
+
         diag.debug('outgoingRequest on response()');
         response.on('end', () => {
           diag.debug('outgoingRequest on end()');
@@ -427,6 +434,7 @@ export class HttpInstrumentation extends InstrumentationBase<Http> {
             response.end = originalEnd;
             // Cannot pass args of type ResponseEndArgs,
             const returned = safeExecuteInTheMiddle(
+              // eslint-disable-next-line @typescript-eslint/no-explicit-any
               () => response.end.apply(this, arguments as any),
               error => {
                 if (error) {
@@ -527,34 +535,46 @@ export class HttpInstrumentation extends InstrumentationBase<Http> {
       const spanOptions: SpanOptions = {
         kind: SpanKind.CLIENT,
       };
+
       const span = instrumentation._startHttpSpan(operationName, spanOptions);
+
+      const parentContext = context.active();
+      const requestContext = setSpan(parentContext, span);
+
       if (!optionsParsed.headers) {
         optionsParsed.headers = {};
       }
-      propagation.inject(
-        setSpan(context.active(), span),
-        optionsParsed.headers
-      );
+      propagation.inject(requestContext, optionsParsed.headers);
 
-      const request: http.ClientRequest = safeExecuteInTheMiddle(
-        () => original.apply(this, [optionsParsed, ...args]),
-        error => {
-          if (error) {
-            utils.setSpanWithError(span, error);
-            instrumentation._closeHttpSpan(span);
-            throw error;
-          }
+      return context.with(requestContext, () => {
+        const cb = args[args.length - 1];
+        if (typeof cb === 'function') {
+          args[args.length - 1] = context.bind(cb);
         }
-      );
 
-      diag.debug('%s instrumentation outgoingRequest', component);
-      context.bind(request);
-      return instrumentation._traceClientRequest(
-        component,
-        request,
-        optionsParsed,
-        span
-      );
+        const request: http.ClientRequest = safeExecuteInTheMiddle(
+          () => original.apply(this, [optionsParsed, ...args]),
+          error => {
+            if (error) {
+              utils.setSpanWithError(span, error);
+              instrumentation._closeHttpSpan(span);
+              throw error;
+            }
+          }
+        );
+
+        context.bind(request);
+
+        diag.debug('%s instrumentation outgoingRequest', component);
+        return instrumentation._traceClientRequest(
+          component,
+          request,
+          optionsParsed,
+          span,
+          requestContext,
+          parentContext
+        );
+      });
     };
   }
 

packages/opentelemetry-instrumentation-http/src/patch.ts

@@ -0,0 +1,99 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+import { context, Context } from '@opentelemetry/api';
+import { EventEmitter } from 'events';
+
+const LISTENER_SYMBOL = Symbol('OtHttpInstrumentationListeners');
+
+const ADD_LISTENER_METHODS = [
+  'addListener' as const,
+  'on' as const,
+  'once' as const,
+  'prependListener' as const,
+  'prependOnceListener' as const,
+];
+
+interface ListenerMap {
+  [name: string]: WeakMap<Function, Function>;
+}
+
+function getListenerMap(emitter: EventEmitter): ListenerMap {
+  let map: ListenerMap = (emitter as never)[LISTENER_SYMBOL];
+  if (map === undefined) {
+    map = {};
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any
+    (emitter as any)[LISTENER_SYMBOL] = map;
+  }
+  return map;
+}
+
+export function bindEmitter<E extends EventEmitter>(
+  emitter: E,
+  getContext: (event: string) => Context
+): E {
+  function getPatchedAddFunction(original: Function) {
+    return function (this: never, event: string, listener: Function) {
+      const listenerMap = getListenerMap(emitter);
+      let listeners = listenerMap[event];
+      if (listeners === undefined) {
+        listeners = new WeakMap();
+        listenerMap[event] = listeners;
+      }
+      const patchedListener = context.bind(listener, getContext(event));
+      listeners.set(listener, patchedListener);
+      return original.call(this, event, patchedListener);
+    };
+  }
+
+  function getPatchedRemoveListener(original: Function) {
+    return function (this: never, event: string, listener: Function) {
+      const listeners = getListenerMap(emitter)[event];
+      if (listeners === undefined) {
+        return original.call(this, event, listener);
+      }
+      return original.call(this, event, listeners.get(listener) || listener);
+    };
+  }
+
+  function getPatchedRemoveAllListeners(original: Function) {
+    return function (this: never, event: string) {
+      delete getListenerMap(emitter)[event];
+      return original.call(this, event);
+    };
+  }
+
+  for (const method of ADD_LISTENER_METHODS) {
+    if (emitter[method] !== undefined) {
+      emitter[method] = getPatchedAddFunction(emitter[method]);
+    }
+  }
+
+  if (typeof emitter.removeListener === 'function') {
+    emitter.removeListener = getPatchedRemoveListener(emitter.removeListener);
+  }
+
+  if (typeof emitter.off === 'function') {
+    emitter.off = getPatchedRemoveListener(emitter.off);
+  }
+
+  if (typeof emitter.removeAllListeners === 'function') {
+    emitter.removeAllListeners = getPatchedRemoveAllListeners(
+      emitter.removeAllListeners
+    );
+  }
+
+  return emitter;
+}

packages/opentelemetry-instrumentation-http/src/types.ts

@@ -44,7 +44,7 @@ export type ParsedRequestOptions =
   | http.RequestOptions;
 export type Http = typeof http;
 export type Https = typeof https;
-/* tslint:disable-next-line:no-any */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
 export type Func<T> = (...args: any[]) => T;
 export type ResponseEndArgs =
   | [((() => void) | undefined)?]

packages/opentelemetry-instrumentation-http/test/functionals/http-enable.test.ts

@@ -19,7 +19,7 @@ import {
   propagation,
   Span as ISpan,
   SpanKind,
-  getSpan,
+  getSpanContext,
   setSpan,
 } from '@opentelemetry/api';
 import { NodeTracerProvider } from '@opentelemetry/node';
@@ -181,6 +181,7 @@ describe('HttpInstrumentation', () => {
         );
       });
     });
+
     describe('with good instrumentation options', () => {
       beforeEach(() => {
         memoryExporter.reset();
@@ -701,11 +702,62 @@ describe('HttpInstrumentation', () => {
         );
       });
 
-      it('should not set span as active in context for outgoing request', done => {
-        assert.deepStrictEqual(getSpan(context.active()), undefined);
-        http.get(`${protocol}://${hostname}:${serverPort}/test`, res => {
-          assert.deepStrictEqual(getSpan(context.active()), undefined);
-          done();
+      it('should set span as active when receiving response or data callbacks when using .get', done => {
+        const span = provider.getTracer('test').startSpan('parentSpan');
+        const parentSpanId = span.context().spanId;
+        assert.notEqual(parentSpanId, undefined);
+
+        context.with(setSpan(context.active(), span), () => {
+          http.get(`${protocol}://${hostname}:${serverPort}/test`, res => {
+            const requestSpanId = getSpanContext(context.active())?.spanId;
+            assert.notEqual(requestSpanId, parentSpanId);
+
+            res.on('data', chunk => {
+              const dataSpanId = getSpanContext(context.active())?.spanId;
+              assert.notEqual(dataSpanId, parentSpanId);
+              assert.equal(dataSpanId, requestSpanId);
+            });
+
+            res.on('end', () => {
+              assert.equal(
+                getSpanContext(context.active())?.spanId,
+                parentSpanId
+              );
+              done();
+            });
+          });
+        });
+      });
+
+      it('should set span as active when receiving response or data callbacks when using .request', done => {
+        const span = provider.getTracer('test').startSpan('parentSpan');
+        const parentSpanId = span.context().spanId;
+        assert.notEqual(parentSpanId, undefined);
+
+        context.with(setSpan(context.active(), span), () => {
+          http
+            .request(
+              { host: hostname, port: serverPort, path: '/test' },
+              res => {
+                const requestSpanId = getSpanContext(context.active())?.spanId;
+                assert.notEqual(requestSpanId, parentSpanId);
+
+                res.on('data', chunk => {
+                  const dataSpanId = getSpanContext(context.active())?.spanId;
+                  assert.notEqual(dataSpanId, parentSpanId);
+                  assert.equal(dataSpanId, requestSpanId);
+                });
+
+                res.on('end', () => {
+                  assert.equal(
+                    getSpanContext(context.active())?.spanId,
+                    parentSpanId
+                  );
+                  done();
+                });
+              }
+            )
+            .end();
         });
       });
     });