feat!: allow users to specify a trace policy impl (#1027)

BREAKING CHANGE: contextHeaderBehavior and ignoreContextHeader now
act independently of one another. The former controls how a sampling
decision is made based on incoming context header, and the latter
controls whether trace context is propagated to the current
request.

src/config.ts

@@ -24,6 +24,34 @@ export type CLSMechanism =
 
 export type ContextHeaderBehavior = 'default'|'ignore'|'require';
 
+export interface RequestDetails {
+  /**
+   * The request timestamp.
+   */
+  timestamp: number;
+  /**
+   * The request URL.
+   */
+  url: string;
+  /**
+   * The request method.
+   */
+  method: string;
+  /**
+   * The parsed trace context, if it exists.
+   */
+  traceContext: {traceId: string; spanId: string; options: number;}|null;
+  /**
+   * The original options object used to create the root span that corresponds
+   * to this request.
+   */
+  options: {};
+}
+
+export interface TracePolicy {
+  shouldTrace: (requestDetails: RequestDetails) => boolean;
+}
+
 /**
  * Available configuration options. All fields are optional. See the
  * defaultConfig object defined in this file for default assigned values.
@@ -144,25 +172,29 @@ export interface Config {
   samplingRate?: number;
 
   /**
-   * Specifies how to use incoming trace context headers. The following options
-   * are available:
-   *   'default' -- Trace context will be propagated for incoming requests that
-   *   contain the context header. A new trace will be created for requests
-   *   without trace context headers. All traces are still subject to local
-   *   sampling and url filter policies.
-   *   'require' -- Same as default, but traces won't be created for requests
-   *   without trace context headers. This should not be set for end user-facing
-   *   services, as this header is usually set by other traced services rather
-   *   than by users.
-   *   'ignore' -- Trace context headers will always be ignored, so a new trace
-   *   with a unique ID will be created for every request. This means that a
-   *   sampling decision specified on an incoming request will be ignored.
+   * Specifies whether to trace based on the 'traced' bit specified on incoming
+   * trace context headers. The following options are available:
+   *   'default' -- Don't trace incoming requests that have a trace context
+   *   header with its 'traced' bit set to 0.
+   *   'require' -- Don't trace incoming requests that have a trace context
+   *   header with its 'traced' bit set to 0, or incoming requests without a
+   *   trace context header.
+   *   'ignore' -- The 'traced' bit will be ignored. In other words, the context
+   *   header isn't used to determine whether a request will be traced at all.
    *   This might be useful for aggregating traces generated by different cloud
    *   platform projects.
-   * All traces are still subject to local tracing policy.
    */
   contextHeaderBehavior?: ContextHeaderBehavior;
 
+  /**
+   * For advanced usage only.
+   * If specified, overrides the built-in trace policy object.
+   * Note that if any of ignoreUrls, ignoreMethods, samplingRate, or
+   * contextHeaderBehavior is specified, an error will be thrown when start()
+   * is called.
+   */
+  tracePolicy?: TracePolicy;
+
   /**
    * Buffer the captured traces for `flushDelaySeconds` seconds before
    * publishing to the Stackdriver Trace API, unless the buffer fills up first.
@@ -199,14 +231,6 @@ export interface Config {
    */
   onUncaughtException?: string;
 
-  /**
-   * Setting this to true or false is the same as setting contextHeaderBehavior
-   * to 'ignore' or 'default' respectively. If both are explicitly set,
-   * contextHeaderBehavior will be prioritized over this value.
-   * Deprecated: This option will be removed in a future release.
-   */
-  ignoreContextHeader?: boolean;
-
   /**
    * The ID of the Google Cloud Platform project with which traces should
    * be associated. The value of GCLOUD_PROJECT takes precedence over this

src/index.ts

@@ -19,15 +19,17 @@ const filesLoadedBeforeTrace = Object.keys(require.cache);
 // This file's top-level imports must not transitively depend on modules that
 // do I/O, or continuation-local-storage will not work.
 import * as semver from 'semver';
-import {Config, defaultConfig} from './config';
+import {Config, defaultConfig, TracePolicy} from './config';
 import * as extend from 'extend';
 import * as path from 'path';
 import * as PluginTypes from './plugin-types';
 import {Tracing, TopLevelConfig} from './tracing';
 import {FORCE_NEW, Forceable, lastOf} from './util';
 import {Constants} from './constants';
-import {StackdriverTracer, TraceContextHeaderBehavior} from './trace-api';
 import {TraceCLSMechanism} from './cls';
+import {StackdriverTracer} from './trace-api';
+import {BuiltinTracePolicy, TraceContextHeaderBehavior} from './tracing-policy';
+import {config} from './plugins/types/bluebird_3';
 
 export {Config, PluginTypes};
 
@@ -55,6 +57,22 @@ function initConfig(userConfig: Forceable<Config>): Forceable<TopLevelConfig> {
       extend(true, {}, defaultConfig, envSetConfig, userConfig);
   const forceNew = userConfig[FORCE_NEW];
 
+  // Throw for improper configurations.
+  const userSetKeys =
+      new Set([...Object.keys(envSetConfig), ...Object.keys(userConfig)]);
+  if (userSetKeys.has('tracePolicy')) {
+    // If the user specified tracePolicy, they should not have also set these
+    // other fields.
+    const forbiddenKeys =
+        ['ignoreUrls', 'ignoreMethods', 'samplingRate', 'contextHeaderBehavior']
+            .filter(key => userSetKeys.has(key))
+            .map(key => `config.${key}`);
+    if (forbiddenKeys.length > 0) {
+      throw new Error(`config.tracePolicy and any of [${
+          forbiddenKeys.join('\, ')}] can't be specified at the same time.`);
+    }
+  }
+
   const getInternalClsMechanism = (clsMechanism: string): TraceCLSMechanism => {
     // If the CLS mechanism is set to auto-determined, decide now
     // what it should be.
@@ -115,27 +133,20 @@ function initConfig(userConfig: Forceable<Config>): Forceable<TopLevelConfig> {
       plugins: {...mergedConfig.plugins},
       tracerConfig: {
         enhancedDatabaseReporting: mergedConfig.enhancedDatabaseReporting,
-        contextHeaderBehavior: lastOf<TraceContextHeaderBehavior>(
-            defaultConfig.contextHeaderBehavior as TraceContextHeaderBehavior,
-            // Internally, ignoreContextHeader is no longer being used, so
-            // convert the user's value into a value for contextHeaderBehavior.
-            // But let this value be overridden by the user's explicitly set
-            // value for contextHeaderBehavior.
-            mergedConfig.ignoreContextHeader ?
-                TraceContextHeaderBehavior.IGNORE :
-                TraceContextHeaderBehavior.DEFAULT,
-            userConfig.contextHeaderBehavior as TraceContextHeaderBehavior),
         rootSpanNameOverride:
             getInternalRootSpanNameOverride(mergedConfig.rootSpanNameOverride),
         spansPerTraceHardLimit: mergedConfig.spansPerTraceHardLimit,
-        spansPerTraceSoftLimit: mergedConfig.spansPerTraceSoftLimit,
-        tracePolicyConfig: {
-          samplingRate: mergedConfig.samplingRate,
-          ignoreMethods: mergedConfig.ignoreMethods,
-          ignoreUrls: mergedConfig.ignoreUrls
-        }
+        spansPerTraceSoftLimit: mergedConfig.spansPerTraceSoftLimit
       }
-    }
+    },
+    tracePolicyConfig: {
+      samplingRate: mergedConfig.samplingRate,
+      ignoreMethods: mergedConfig.ignoreMethods,
+      ignoreUrls: mergedConfig.ignoreUrls,
+      contextHeaderBehavior: mergedConfig.contextHeaderBehavior as
+          TraceContextHeaderBehavior
+    },
+    overrides: {tracePolicy: mergedConfig.tracePolicy}
   };
 }
 

src/trace-api.ts

@@ -19,53 +19,25 @@ import * as is from 'is';
 import * as uuid from 'uuid';
 
 import {cls, RootContext} from './cls';
+import {TracePolicy} from './config';
 import {Constants, SpanType} from './constants';
 import {Logger} from './logger';
 import {Func, RootSpan, RootSpanOptions, Span, SpanOptions, Tracer} from './plugin-types';
 import {RootSpanData, UNCORRELATED_CHILD_SPAN, UNCORRELATED_ROOT_SPAN, UNTRACED_CHILD_SPAN, UNTRACED_ROOT_SPAN} from './span-data';
 import {TraceLabels} from './trace-labels';
 import {traceWriter} from './trace-writer';
-import {TracePolicy, TracePolicyConfig} from './tracing-policy';
+import {neverTrace} from './tracing-policy';
 import * as util from './util';
 
-/**
- * An enumeration of the different possible types of behavior when dealing with
- * incoming trace context. Requests are still subject to local tracing policy.
- */
-export enum TraceContextHeaderBehavior {
-  /**
-   * Respect the trace context header if it exists; otherwise, trace the
-   * request as a new trace.
-   */
-  DEFAULT = 'default',
-  /**
-   * Respect the trace context header if it exists; otherwise, treat the
-   * request as unsampled and don't trace it.
-   */
-  REQUIRE = 'require',
-  /**
-   * Trace every request as a new trace, even if trace context exists.
-   */
-  IGNORE = 'ignore'
-}
-
 /**
  * An interface describing configuration fields read by the StackdriverTracer
  * object. This includes fields read by the trace policy.
  */
 export interface StackdriverTracerConfig {
   enhancedDatabaseReporting: boolean;
-  contextHeaderBehavior: TraceContextHeaderBehavior;
   rootSpanNameOverride: (path: string) => string;
   spansPerTraceSoftLimit: number;
   spansPerTraceHardLimit: number;
-  tracePolicyConfig: TracePolicyConfig;
-}
-
-interface IncomingTraceContext {
-  traceId?: string;
-  spanId?: string;
-  options: number;
 }
 
 /**
@@ -117,10 +89,10 @@ export class StackdriverTracer implements Tracer {
    * @param logger A logger object.
    * @private
    */
-  enable(config: StackdriverTracerConfig, logger: Logger) {
+  enable(config: StackdriverTracerConfig, policy: TracePolicy, logger: Logger) {
     this.logger = logger;
     this.config = config;
-    this.policy = new TracePolicy(config.tracePolicyConfig);
+    this.policy = policy;
     this.enabled = true;
   }
 
@@ -134,7 +106,7 @@ export class StackdriverTracer implements Tracer {
     // never generates traces allows persisting wrapped methods (either because
     // they are already instantiated or the plugin doesn't unpatch them) to
     // short-circuit out of trace generation logic.
-    this.policy = TracePolicy.never();
+    this.policy = neverTrace();
     this.enabled = false;
   }
 
@@ -175,48 +147,49 @@ export class StackdriverTracer implements Tracer {
     }
 
     // Attempt to read incoming trace context.
-    const incomingTraceContext: IncomingTraceContext = {options: 1};
-    let parsedContext: util.TraceContext|null = null;
-    if (isString(options.traceContext) &&
-        this.config!.contextHeaderBehavior !==
-            TraceContextHeaderBehavior.IGNORE) {
-      parsedContext = util.parseContextFromHeader(options.traceContext);
-    }
-    if (parsedContext) {
-      if (parsedContext.options === undefined) {
-        // If there are no incoming option flags, default to 0x1.
-        parsedContext.options = 1;
+    const parseContext = (stringifiedTraceContext?: string|null) => {
+      const parsedContext = isString(stringifiedTraceContext) ?
+          util.parseContextFromHeader(stringifiedTraceContext) :
+          null;
+      if (parsedContext) {
+        if (parsedContext.options === undefined) {
+          // If there are no incoming option flags, default to 0x1.
+          parsedContext.options = 1;
+        }
       }
-      Object.assign(incomingTraceContext, parsedContext);
-    } else if (
-        this.config!.contextHeaderBehavior ===
-        TraceContextHeaderBehavior.REQUIRE) {
-      incomingTraceContext.options = 0;
-    }
+      return parsedContext as Required<util.TraceContext>| null;
+    };
+    const traceContext = parseContext(options.traceContext);
 
     // Consult the trace policy.
-    const locallyAllowed = this.policy!.shouldTrace({
+    const shouldTrace = this.policy!.shouldTrace({
       timestamp: Date.now(),
       url: options.url || '',
-      method: options.method || ''
+      method: options.method || '',
+      traceContext,
+      options
     });
-    const remotelyAllowed = !!(
-        incomingTraceContext.options & Constants.TRACE_OPTIONS_TRACE_ENABLED);
 
     let rootContext: RootSpan&RootContext;
+
     // Don't create a root span if the trace policy disallows it.
-    if (!locallyAllowed || !remotelyAllowed) {
+    if (!shouldTrace) {
       rootContext = UNTRACED_ROOT_SPAN;
     } else {
       // Create a new root span, and invoke fn with it.
-      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 */
-          name,                                /* Span name */
-          parentId,                            /* Parent's span ID */
+          // Trace object
+          {
+            projectId: '',
+            traceId: traceContext ? traceContext.traceId :
+                                    uuid.v4().split('-').join(''),
+            spans: []
+          },
+          // Span name
+          this.config!.rootSpanNameOverride(options.name),
+          // Parent span ID
+          traceContext ? traceContext.spanId : '0',
+          // Number of stack frames to skip
           options.skipFrames || 0);
     }