fix(pg): update requireParentSpan to skip instrumentation when parent…

… not present (open-telemetry#1343) * fix(pg): update requireParentSpan to skip instrumentation when parent not present * Remove unnecessary memoryExporter.reset() --------- Co-authored-by: Amir Blum <[email protected]>
dynatrace-oss-contrib · Feb 7, 2023 · d23c329 · d23c329
1 parent 52136d8
commit d23c329
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 76 deletions.
diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts b/plugins/node/opentelemetry-instrumentation-pg/src/instrumentation.ts
@@ -20,7 +20,14 @@ import {
   safeExecuteInTheMiddle,
 } from '@opentelemetry/instrumentation';
 
-import { context, diag, trace, Span, SpanStatusCode } from '@opentelemetry/api';
+import {
+  context,
+  diag,
+  trace,
+  Span,
+  SpanStatusCode,
+  SpanKind,
+} from '@opentelemetry/api';
 import type * as pgTypes from 'pg';
 import type * as pgPoolTypes from 'pg-pool';
 import {
@@ -39,7 +46,6 @@ import {
   DbSystemValues,
 } from '@opentelemetry/semantic-conventions';
 import { VERSION } from './version';
-import { startSpan } from './utils';
 
 const PG_POOL_COMPONENT = 'pg-pool';
 
@@ -131,13 +137,18 @@ export class PgInstrumentation extends InstrumentationBase {
         this: pgTypes.Client,
         callback?: PgErrorCallback
       ) {
-        const span = startSpan(
-          plugin.tracer,
-          plugin.getConfig(),
+        if (utils.shouldSkipInstrumentation(plugin.getConfig())) {
+          return original.call(this, callback);
+        }
+
+        const span = plugin.tracer.startSpan(
           `${PgInstrumentation.COMPONENT}.connect`,
           {
-            [SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL,
-            ...utils.getSemanticAttributesFromConnection(this),
+            kind: SpanKind.CLIENT,
+            attributes: {
+              [SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL,
+              ...utils.getSemanticAttributesFromConnection(this),
+            },
           }
         );
 
@@ -168,6 +179,10 @@ export class PgInstrumentation extends InstrumentationBase {
         `Patching ${PgInstrumentation.COMPONENT}.Client.prototype.query`
       );
       return function query(this: PgClientExtended, ...args: unknown[]) {
+        if (utils.shouldSkipInstrumentation(plugin.getConfig())) {
+          return original.apply(this, args as never);
+        }
+
         // client.query(text, cb?), client.query(text, values, cb?), and
         // client.query(configObj, cb?) are all valid signatures. We construct
         // a queryConfig obj from all (valid) signatures to build the span in a
@@ -348,19 +363,21 @@ export class PgInstrumentation extends InstrumentationBase {
     const plugin = this;
     return (originalConnect: typeof pgPoolTypes.prototype.connect) => {
       return function connect(this: PgPoolExtended, callback?: PgPoolCallback) {
+        if (utils.shouldSkipInstrumentation(plugin.getConfig())) {
+          return originalConnect.call(this, callback as any);
+        }
+
         // setup span
-        const span = startSpan(
-          plugin.tracer,
-          plugin.getConfig(),
-          `${PG_POOL_COMPONENT}.connect`,
-          {
+        const span = plugin.tracer.startSpan(`${PG_POOL_COMPONENT}.connect`, {
+          kind: SpanKind.CLIENT,
+          attributes: {
             [SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL,
             ...utils.getSemanticAttributesFromConnection(this.options),
             [AttributeNames.IDLE_TIMEOUT_MILLIS]:
               this.options.idleTimeoutMillis,
             [AttributeNames.MAX_CLIENT]: this.options.maxClient,
-          }
-        );
+          },
+        });
 
         if (callback) {
           const parentSpan = trace.getSpan(context.active());

diff --git a/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts b/plugins/node/opentelemetry-instrumentation-pg/src/utils.ts
@@ -22,8 +22,6 @@ import {
   Tracer,
   SpanKind,
   diag,
-  INVALID_SPAN_CONTEXT,
-  Attributes,
   defaultTextMapSetter,
   ROOT_CONTEXT,
 } from '@opentelemetry/api';
@@ -116,24 +114,13 @@ export function getSemanticAttributesFromConnection(
   };
 }
 
-export function startSpan(
-  tracer: Tracer,
-  instrumentationConfig: PgInstrumentationConfig,
-  name: string,
-  attributes: Attributes
-): Span {
-  // If a parent span is required but not present, use a noop span to propagate
-  // context without recording it. Adapted from opentelemetry-instrumentation-http:
-  // https://github.com/open-telemetry/opentelemetry-js/blob/597ea98e58a4f68bcd9aec5fd283852efe444cd6/experimental/packages/opentelemetry-instrumentation-http/src/http.ts#L660
-  const currentSpan = trace.getSpan(context.active());
-  if (instrumentationConfig.requireParentSpan && currentSpan === undefined) {
-    return trace.wrapSpanContext(INVALID_SPAN_CONTEXT);
-  }
-
-  return tracer.startSpan(name, {
-    kind: SpanKind.CLIENT,
-    attributes,
-  });
+export function shouldSkipInstrumentation(
+  instrumentationConfig: PgInstrumentationConfig
+) {
+  return (
+    instrumentationConfig.requireParentSpan === true &&
+    trace.getSpan(context.active()) === undefined
+  );
 }
 
 // Create a span from our normalized queryConfig object,
@@ -149,9 +136,12 @@ export function handleConfigQuery(
   const dbName = connectionParameters.database;
 
   const spanName = getQuerySpanName(dbName, queryConfig);
-  const span = startSpan(tracer, instrumentationConfig, spanName, {
-    [SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL, // required
-    ...getSemanticAttributesFromConnection(connectionParameters),
+  const span = tracer.startSpan(spanName, {
+    kind: SpanKind.CLIENT,
+    attributes: {
+      [SemanticAttributes.DB_SYSTEM]: DbSystemValues.POSTGRESQL, // required
+      ...getSemanticAttributesFromConnection(connectionParameters),
+    },
   });
 
   if (!queryConfig) {

diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/pg-pool.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/pg-pool.test.ts
@@ -269,6 +269,23 @@ describe('pg-pool', () => {
         assert.strictEqual(resNoPromise, undefined, 'No promise is returned');
       });
     });
+
+    it('should not generate traces when requireParentSpan=true is specified', async () => {
+      // The pool gets shared between tests. We need to create a separate one
+      // to test cold start, which can cause nested spans
+      const newPool = new pgPool(CONFIG);
+      create({
+        requireParentSpan: true,
+      });
+      const client = await newPool.connect();
+      try {
+        await client.query('SELECT NOW()');
+      } finally {
+        client.release();
+      }
+      const spans = memoryExporter.getFinishedSpans();
+      assert.strictEqual(spans.length, 0);
+    });
   });
 
   describe('#pool.query()', () => {

diff --git a/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts b/plugins/node/opentelemetry-instrumentation-pg/test/utils.test.ts
@@ -128,58 +128,35 @@ describe('utils.ts', () => {
     });
   });
 
-  describe('.startSpan()', () => {
-    it('starts real span when requireParentSpan=false', async () => {
-      const span = utils.startSpan(tracer, instrumentationConfig, 'spanName', {
-        key: 'value',
-      });
-      span.end();
-
-      const readableSpan = getLatestSpan();
-
-      assert.strictEqual(readableSpan.name, 'spanName');
-      assert.strictEqual(readableSpan.attributes['key'], 'value');
-      assert.notDeepStrictEqual(readableSpan.spanContext, INVALID_SPAN_CONTEXT);
+  describe('.shouldSkipInstrumentation()', () => {
+    it('returns false when requireParentSpan=false', async () => {
+      assert.strictEqual(
+        utils.shouldSkipInstrumentation(instrumentationConfig),
+        false
+      );
     });
 
-    it('starts real span when requireParentSpan=true and there is a parent span', async () => {
+    it('returns false requireParentSpan=true and there is a parent span', async () => {
       const parent = tracer.startSpan('parentSpan');
       context.with(trace.setSpan(context.active(), parent), () => {
-        const childSpan = utils.startSpan(
-          tracer,
-          {
+        assert.strictEqual(
+          utils.shouldSkipInstrumentation({
             ...instrumentationConfig,
             requireParentSpan: true,
-          },
-          'childSpan',
-          { key: 'value' }
-        );
-        childSpan.end();
-
-        const readableSpan = getLatestSpan();
-        assert.strictEqual(readableSpan.name, 'childSpan');
-        assert.strictEqual(readableSpan.attributes['key'], 'value');
-        assert.notDeepStrictEqual(
-          readableSpan.spanContext,
-          INVALID_SPAN_CONTEXT
+          }),
+          false
         );
       });
     });
 
-    it('creates placeholder span when requireParentSpan=true and there is no parent span', async () => {
-      const span = utils.startSpan(
-        tracer,
-        {
+    it('returns true when requireParentSpan=true and there is no parent span', async () => {
+      assert.strictEqual(
+        utils.shouldSkipInstrumentation({
           ...instrumentationConfig,
           requireParentSpan: true,
-        },
-        'spanName',
-        { key: 'value' }
+        }),
+        true
       );
-      span.end();
-
-      const readableSpan = getLatestSpan();
-      assert.strictEqual(readableSpan, undefined);
     });
   });