feat: add rootSpanNameOverride option (#826)

This PR introduces rootSpanNameOverride, which allows us to override root span names with either a string value or a mapping function.

src/config.ts

@@ -62,6 +62,14 @@ export interface Config {
    */
   enhancedDatabaseReporting?: boolean;
 
+  /**
+   * A value that can be used to override names of root spans. If specified as
+   * a string, the string will be used to replace all such span names; if
+   * specified as a function, the function will be invoked with the request path
+   * as an argument, and its return value will be used as the span name.
+   */
+  rootSpanNameOverride?: string|((name: string) => string);
+
   /**
    * The maximum number of characters reported on a label value. This value
    * cannot exceed 16383, the maximum value accepted by the service.
@@ -193,6 +201,7 @@ export const defaultConfig = {
   logLevel: 1,
   enabled: true,
   enhancedDatabaseReporting: false,
+  rootSpanNameOverride: (name: string) => name,
   maximumLabelValueSize: 512,
   plugins: {
     // enable all by default

src/constants.ts

@@ -19,7 +19,7 @@
 /** Constant values. */
 // tslint:disable-next-line:variable-name
 export const Constants = {
-  /** The metadata key under which trace context  */
+  /** The metadata key under which trace context is stored as a binary value. */
   TRACE_CONTEXT_GRPC_METADATA_NAME: 'grpc-trace-bin',
 
   /** Header that carries trace context across Google infrastructure. */

src/index.ts

@@ -76,6 +76,13 @@ function initConfig(projectConfig: Forceable<Config>):
       Constants.TRACE_SERVICE_LABEL_VALUE_LIMIT) {
     config.maximumLabelValueSize = Constants.TRACE_SERVICE_LABEL_VALUE_LIMIT;
   }
+  // Make rootSpanNameOverride a function if not already.
+  if (typeof config.rootSpanNameOverride === 'string') {
+    const spanName = config.rootSpanNameOverride;
+    config.rootSpanNameOverride = () => spanName;
+  } else if (typeof config.rootSpanNameOverride !== 'function') {
+    config.rootSpanNameOverride = (name: string) => name;
+  }
 
   // If the CLS mechanism is set to auto-determined, decide now what it should
   // be.

src/plugin-types.ts

@@ -18,6 +18,7 @@
 // tslint:disable:no-any
 
 import {Constants, SpanType} from './constants';
+import {StackdriverTracerConfig} from './trace-api';
 import {TraceLabels} from './trace-labels';
 import {TraceContext} from './util';
 
@@ -116,6 +117,12 @@ export interface Tracer {
    */
   enhancedDatabaseReportingEnabled(): boolean;
 
+  /**
+   * Gets the current configuration, or throws if it can't be retrieved
+   * because the Trace Agent was not disabled.
+   */
+  getConfig(): StackdriverTracerConfig;
+
   /**
    * Runs the given function in a root span corresponding to an incoming
    * request, passing it an object that exposes an interface for adding

src/trace-api.ts

@@ -35,6 +35,7 @@ export interface StackdriverTracerConfig extends
     TracingPolicy.TracePolicyConfig {
   enhancedDatabaseReporting: boolean;
   ignoreContextHeader: boolean;
+  rootSpanNameOverride: (path: string) => string;
 }
 
 interface IncomingTraceContext {
@@ -126,6 +127,13 @@ export class StackdriverTracer implements Tracer {
     return !!this.config && this.config.enhancedDatabaseReporting;
   }
 
+  getConfig(): StackdriverTracerConfig {
+    if (!this.config) {
+      throw new Error('Configuration is not available.');
+    }
+    return this.config;
+  }
+
   runInRootSpan<T>(options: RootSpanOptions, fn: (span: RootSpan) => T): T {
     if (!this.isActive()) {
       return fn(UNTRACED_ROOT_SPAN);
@@ -166,9 +174,10 @@ export class StackdriverTracer implements Tracer {
       const traceId =
           incomingTraceContext.traceId || (uuid.v4().split('-').join(''));
       const parentId = incomingTraceContext.spanId || '0';
+      const name = this.config!.rootSpanNameOverride(options.name);
       rootContext = new RootSpanData(
           {projectId: '', traceId, spans: []}, /* Trace object */
-          options.name,                        /* Span name */
+          name,                                /* Span name */
           parentId,                            /* Parent's span ID */
           options.skipFrames || 0);
     }

test/plugins/common.ts

@@ -73,6 +73,7 @@ shimmer.wrap(trace, 'start', function(original) {
     testTraceAgent.enable({
       enhancedDatabaseReporting: false,
       ignoreContextHeader: false,
+      rootSpanNameOverride: (name: string) => name,
       samplingRate: 0
     }, logger());
     testTraceAgent.policy = new TracingPolicy.TraceAllPolicy();

test/test-config-cls.ts → test/test-config.ts

@@ -22,6 +22,8 @@ import * as util from 'util';
 import {TraceCLSConfig, TraceCLSMechanism} from '../src/cls';
 
 import * as testTraceModule from './trace';
+import { NormalizedConfig } from '../src/tracing';
+import { StackdriverTracer } from '../src/trace-api';
 
 describe('Behavior set by config for context propagation mechanism', () => {
   const useAH = semver.satisfies(process.version, '>=8');
@@ -90,3 +92,64 @@ describe('Behavior set by config for context propagation mechanism', () => {
        });
   }
 });
+
+describe('Behavior set by config for overriding root span name', () => {
+  let capturedConfig: NormalizedConfig|null;
+
+  class CaptureConfigTestTracing extends testTraceModule.TestTracing {
+    constructor(config: NormalizedConfig, traceAgent: StackdriverTracer) {
+      super(config, traceAgent);
+      // Capture the config object passed into this constructor.
+      capturedConfig = config;
+    }
+  }
+
+  beforeEach(() => {
+    capturedConfig = null;
+  });
+
+  before(() => {
+    testTraceModule.setTracingForTest(CaptureConfigTestTracing);
+  });
+
+  after(() => {
+    testTraceModule.setTracingForTest(testTraceModule.TestTracing);
+  });
+
+  it('should convert a string to a function', () => {
+    testTraceModule.start({
+      rootSpanNameOverride: 'hello'
+    });
+    assert.ok(capturedConfig!);
+    // Avoid using the ! operator multiple times.
+    const config = capturedConfig!;
+    // If !config.enabled, then TSC does not permit access to other fields on
+    // config. So use this structure instead of assert.ok(config.enabled).
+    if (config.enabled) {
+      assert.strictEqual(typeof config.rootSpanNameOverride, 'function');
+      assert.strictEqual(config.rootSpanNameOverride(''), 'hello');
+    } else {
+      assert.fail('Configuration was not enabled.');
+    }
+  });
+
+  it('should convert a non-string, non-function to the identity fn', () => {
+    testTraceModule.start({
+      // We should make sure passing in unsupported values at least doesn't
+      // result in a crash.
+      // tslint:disable-next-line:no-any
+      rootSpanNameOverride: 2 as any
+    });
+    assert.ok(capturedConfig!);
+    // Avoid using the ! operator multiple times.
+    const config = capturedConfig!;
+    // If !config.enabled, then TSC does not permit access to other fields on
+    // config. So use this structure instead of assert.ok(config.enabled).
+    if (config.enabled) {
+      assert.strictEqual(typeof config.rootSpanNameOverride, 'function');
+      assert.strictEqual(config.rootSpanNameOverride('a'), 'a');
+    } else {
+      assert.fail('Configuration was not enabled.');
+    }
+  });
+});

test/test-plugin-loader.ts

@@ -47,6 +47,7 @@ describe('Trace Plugin Loader', () => {
               ignoreUrls: [],
               enhancedDatabaseReporting: false,
               ignoreContextHeader: false,
+              rootSpanNameOverride: (name: string) => name,
               projectId: '0'
             },
             config),

test/test-trace-api.ts

@@ -38,6 +38,7 @@ describe('Trace Interface', () => {
             {
               enhancedDatabaseReporting: false,
               ignoreContextHeader: false,
+              rootSpanNameOverride: (name: string) => name,
               samplingRate: 0
             },
             config),

test/test-trace-web-frameworks.ts

@@ -334,6 +334,27 @@ describe('Web framework tracing', () => {
               `span name ${serverSpan.name} includes query parameters`);
         });
       });
+
+      it('uses the span name override option', async () => {
+        const oldSpanNameOverride =
+            testTraceModule.get().getConfig().rootSpanNameOverride;
+        testTraceModule.get().getConfig().rootSpanNameOverride =
+            (path: string) => `${path}-goodbye`;
+        try {
+          await testTraceModule.get().runInRootSpan(
+              {name: 'outer'}, async (span) => {
+                assert.ok(testTraceModule.get().isRealSpan(span));
+                await axios.get(`http://localhost:${port}/hello`);
+                span!.endSpan();
+              });
+          assert.strictEqual(testTraceModule.getSpans().length, 3);
+          const serverSpan = testTraceModule.getOneSpan(isServerSpan);
+          assert.strictEqual(serverSpan.name, '/hello-goodbye');
+        } finally {
+          testTraceModule.get().getConfig().rootSpanNameOverride =
+              oldSpanNameOverride;
+        }
+      });
     });
   });
 });

test/trace.ts

@@ -44,6 +44,7 @@ import {cls, TraceCLS, TraceCLSMechanism} from '../src/cls';
 import {Trace, TraceSpan} from '../src/trace';
 import {PluginLoader, pluginLoader} from '../src/trace-plugin-loader';
 import {TraceWriter, traceWriter, TraceWriterConfig} from '../src/trace-writer';
+import {tracing, Tracing} from '../src/tracing';
 import {FORCE_NEW} from '../src/util';
 
 import {TestLogger} from './logger';
@@ -85,10 +86,15 @@ export class TestPluginLoader extends PluginLoader {
   }
 }
 
+// Make no modifications to Tracing. This class exists for symmetricality
+// purposes.
+export class TestTracing extends Tracing {}
+
 setCLSForTest(TestCLS);
 setLoggerForTest(TestLogger);
 setTraceWriterForTest(TestTraceWriter);
 setPluginLoaderForTest(TestPluginLoader);
+setTracingForTest(TestTracing);
 
 export type Predicate<T> = (value: T) => boolean;
 
@@ -124,6 +130,10 @@ export function setPluginLoaderForTest(impl?: typeof PluginLoader) {
   pluginLoader['implementation'] = impl || PluginLoader;
 }
 
+export function setTracingForTest(impl?: typeof Tracing) {
+  tracing['implementation'] = impl || Tracing;
+}
+
 export function getTraces(predicate?: Predicate<Trace>): Trace[] {
   if (!predicate) {
     predicate = () => true;

test/utils.ts

@@ -37,7 +37,7 @@ export const ASSERT_SPAN_TIME_TOLERANCE_MS = 40;
  */
 
 export function isServerSpan(span: TraceSpan) {
-  return span.kind === 'RPC_SERVER' && span.name !== 'outer';
+  return span.kind === 'RPC_SERVER' && !span.name.startsWith('outer');
 }
 
 // Convenience function that, when awaited, stalls for a given duration of time