From 9a51618fb5084b3e88dcbcb24b38c123ac22e434 Mon Sep 17 00:00:00 2001
From: EddeCCC <eduard.schander@gmx.de>
Date: Wed, 11 Oct 2023 21:31:42 +0200
Subject: [PATCH] refactor Tracer instrumentation

---
 src/impl/index.ts                             | 28 ++++++++++++-------
 .../documentLoadInstrumentation.ts            | 13 +++++++--
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/src/impl/index.ts b/src/impl/index.ts
index 2a50bdc..e95cae2 100644
--- a/src/impl/index.ts
+++ b/src/impl/index.ts
@@ -1,4 +1,4 @@
-import api, { context, trace, Span } from '@opentelemetry/api';
+import api, { context, trace, Span, SpanOptions, Context } from '@opentelemetry/api';
 import {
   AlwaysOnSampler,
   AlwaysOffSampler,
@@ -27,7 +27,10 @@ import { B3InjectEncoding, B3Propagator } from '@opentelemetry/propagator-b3';
 import { PluginProperties, ContextFunction, PropagationHeader } from '../types';
 import { patchExporter, patchExporterClass } from './patchCollectorPrototype';
 import { MultiSpanProcessor, CustomSpanProcessor } from './spanProcessing';
-import { CustomDocumentLoadInstrumentation, patchTracer } from './instrumentation/documentLoadInstrumentation';
+import {
+  CustomDocumentLoadInstrumentation,
+  patchTracerForTransactions
+} from './instrumentation/documentLoadInstrumentation';
 import { CustomIdGenerator } from './transaction/transactionIdGeneration';
 import { TransactionSpanManager } from './transaction/transactionSpanManager';
 import { CustomXMLHttpRequestInstrumentation } from './instrumentation/xmlHttpRequestInstrumentation';
@@ -181,10 +184,8 @@ export default class OpenTelemetryTracingImpl {
     // store the webtracer
     this.traceProvider = providerWithZone;
 
-    // If recordTransaction is enabled, patch the Tracer to always use the transaction span as root span
-    // and initialize the transaction data storage
+    // If recordTransaction is enabled, initialize the transaction manager
     if(this.isTransactionRecordingEnabled()) {
-      patchTracer();
       const delay = this.props.plugins_config?.instrument_document_load?.exporterDelay;
       TransactionSpanManager.initialize(true, this.customIdGenerator);
 
@@ -260,15 +261,22 @@ export default class OpenTelemetryTracingImpl {
    */
   private instrumentTracerClass = () => {
     const { commonAttributes, serviceName } = this.props;
-    // don't patch the function if no attributes are defined
-    if (Object.keys(commonAttributes).length <= 0) {
+
+    let startSpanFunction: (name: string, options?: SpanOptions, context?: Context) => (Span);
+
+    // If recordTransaction is enabled, patch the Tracer to always use the transaction span as root span
+    if(this.isTransactionRecordingEnabled)
+      startSpanFunction = patchTracerForTransactions();
+    else
+      startSpanFunction = Tracer.prototype.startSpan;
+
+    // don't patch the function if no attributes are defined AND no serviceName is defined
+    if (!serviceName && Object.keys(commonAttributes).length <= 0) {
       return;
     }
 
-    const originalStartSpanFunction = Tracer.prototype.startSpan;
-
     Tracer.prototype.startSpan = function () {
-      const span: Span = originalStartSpanFunction.apply(this, arguments);
+      const span: Span = startSpanFunction.apply(this, arguments);
 
       // add common attributes to each span
       if (commonAttributes) {
diff --git a/src/impl/instrumentation/documentLoadInstrumentation.ts b/src/impl/instrumentation/documentLoadInstrumentation.ts
index d3b08c0..eec8a95 100644
--- a/src/impl/instrumentation/documentLoadInstrumentation.ts
+++ b/src/impl/instrumentation/documentLoadInstrumentation.ts
@@ -10,6 +10,7 @@ import { sanitizeAttributes } from '@opentelemetry/core/build/src/common/attribu
 import { TransactionSpanManager } from '../transaction/transactionSpanManager';
 import { addUrlParams } from './urlParams';
 import { GlobalInstrumentationConfig, RequestParameterConfig } from '../../types';
+import { Context, SpanOptions } from '@opentelemetry/api';
 
 export interface CustomDocumentLoadInstrumentationConfig extends InstrumentationConfig {
   recordTransaction?: boolean;
@@ -18,12 +19,17 @@ export interface CustomDocumentLoadInstrumentationConfig extends Instrumentation
 
 /**
  * Patch the Tracer class to use the transaction span as root span
+ * For any additional instrumentation of the startSpan() function, you have to use the
+ * new returned function
+ *
  * OpenTelemetry version: 0.25.0
+ *
+ * @return new startSpan() function
  */
-export function patchTracer() {
+export function patchTracerForTransactions(): (name: string, options?: SpanOptions, context?: Context) => (api.Span) {
   // Overwrite startSpan() in Tracer class
   // Copy of the original startSpan()-function with additional logic inside the function to determine the parentContext
-  Tracer.prototype.startSpan = function (
+  const overwrittenFunction = function (
     name: string,
     options: api.SpanOptions = {},
     context = api.context.active()
@@ -115,6 +121,9 @@ export function patchTracer() {
     span.setAttributes(Object.assign(attributes, samplingResult.attributes));
     return span;
   }
+
+  Tracer.prototype.startSpan = overwrittenFunction;
+  return overwrittenFunction;
 }
 
 type PerformanceEntriesWithServerTiming = PerformanceEntries & {serverTiming?: ReadonlyArray<({name: string, duration: number, description: string})>}