feat: allow "disabling" cls, and relax requirements for creating root…

… spans (#728)

PR-URL: #728

.circleci/config.yml

@@ -74,7 +74,7 @@ postgres_service: &postgres_service
     POSTGRES_DB: test
 
 mysql_service: &mysql_service
-  image: mysql
+  image: mysql:5
   environment:
     MYSQL_ROOT_PASSWORD: Password12!
     MYSQL_DATABASE: test

src/cls.ts

@@ -26,6 +26,8 @@ import {SpanDataType} from './constants';
 import {Trace, TraceSpan} from './trace';
 import {Singleton} from './util';
 
+const asyncHooksAvailable = semver.satisfies(process.version, '>=8');
+
 export interface RealRootContext {
   readonly span: TraceSpan;
   readonly trace: Trace;
@@ -48,11 +50,34 @@ export interface PhantomRootContext {
  */
 export type RootContext = RealRootContext|PhantomRootContext;
 
-const asyncHooksAvailable = semver.satisfies(process.version, '>=8');
+/**
+ * An enumeration of the possible mechanisms for supporting context propagation
+ * through continuation-local storage.
+ */
+export enum TraceCLSMechanism {
+  /**
+   * Use the AsyncHooksCLS class to propagate root span context.
+   * Only available in Node 8+.
+   */
+  ASYNC_HOOKS = 'async-hooks',
+  /**
+   * Use the AsyncListenerCLS class to propagate root span context.
+   * Note that continuation-local-storage should be loaded as the first module.
+   */
+  ASYNC_LISTENER = 'async-listener',
+  /**
+   * Do not use any special mechanism to propagate root span context.
+   * Only a single root span can be open at a time.
+   */
+  NONE = 'none'
+}
 
-export interface TraceCLSConfig { mechanism: 'async-listener'|'async-hooks'; }
+/**
+ * Configuration options passed to the TraceCLS constructor.
+ */
+export interface TraceCLSConfig { mechanism: TraceCLSMechanism; }
 
-export interface CLSConstructor {
+interface CLSConstructor {
   new(defaultContext: RootContext): CLS<RootContext>;
 }
 
@@ -80,30 +105,29 @@ export class TraceCLS implements CLS<RootContext> {
   readonly rootSpanStackOffset: number;
 
   constructor(private readonly logger: Logger, config: TraceCLSConfig) {
-    const useAH = config.mechanism === 'async-hooks' && asyncHooksAvailable;
-    if (useAH) {
-      this.CLSClass = AsyncHooksCLS;
-      this.rootSpanStackOffset = 4;
-      this.logger.info(
-          'TraceCLS#constructor: Created [async-hooks] CLS instance.');
-    } else {
-      if (config.mechanism !== 'async-listener') {
-        if (config.mechanism === 'async-hooks') {
-          this.logger.error(
-              'TraceCLS#constructor: [async-hooks]-based context',
-              `propagation is not available in Node ${process.version}.`);
-        } else {
-          this.logger.error(
-              'TraceCLS#constructor: The specified CLS mechanism',
-              `[${config.mechanism}] was not recognized.`);
+    switch (config.mechanism) {
+      case TraceCLSMechanism.ASYNC_HOOKS:
+        if (!asyncHooksAvailable) {
+          throw new Error(`CLS mechanism [${
+              config.mechanism}] is not compatible with Node <8.`);
         }
-        throw new Error(`CLS mechanism [${config.mechanism}] is invalid.`);
-      }
-      this.CLSClass = AsyncListenerCLS;
-      this.rootSpanStackOffset = 8;
-      this.logger.info(
-          'TraceCLS#constructor: Created [async-listener] CLS instance.');
+        this.CLSClass = AsyncHooksCLS;
+        this.rootSpanStackOffset = 4;
+        break;
+      case TraceCLSMechanism.ASYNC_LISTENER:
+        this.CLSClass = AsyncListenerCLS;
+        this.rootSpanStackOffset = 8;
+        break;
+      case TraceCLSMechanism.NONE:
+        this.CLSClass = UniversalCLS;
+        this.rootSpanStackOffset = 4;
+        break;
+      default:
+        throw new Error(
+            `CLS mechanism [${config.mechanism}] was not recognized.`);
     }
+    this.logger.info(
+        `TraceCLS#constructor: Created [${config.mechanism}] CLS instance.`);
     this.currentCLS = new UniversalCLS(TraceCLS.UNTRACED);
     this.currentCLS.enable();
   }

src/cls/async-listener.ts

@@ -71,7 +71,6 @@ export class AsyncListenerCLS<Context extends {}> implements CLS<Context> {
 
   runWithNewContext<T>(fn: Func<T>): T {
     return this.getNamespace().runAndReturn(() => {
-      this.setContext(this.defaultContext);
       return fn();
     });
   }

src/cls/universal.ts

@@ -40,7 +40,7 @@ export class UniversalCLS<Context> implements CLS<Context> {
 
   disable(): void {
     this.enabled = false;
-    this.currentContext = this.defaultContext;
+    this.setContext(this.defaultContext);
   }
 
   getContext(): Context {

src/config.ts

@@ -19,8 +19,22 @@ import * as path from 'path';
 const pluginDirectory =
     path.join(path.resolve(__dirname, '..'), 'src', 'plugins');
 
+export type CLSMechanism = 'none'|'auto';
+
 /** Available configuration options. */
 export interface Config {
+  /**
+   * The trace context propagation mechanism to use. The following options are
+   * available:
+   * - 'auto' uses continuation-local-storage, unless async_hooks is available
+   *   _and_ the environment variable GCLOUD_TRACE_NEW_CONTEXT is set, in which
+   *   case async_hooks will be used instead.
+   * - 'none' disables CLS completely.
+   * The 'auto' mechanism is used by default if this configuration option is
+   * not explicitly set.
+   */
+  clsMechanism?: CLSMechanism;
+
   /**
    * Log levels: 0=disabled, 1=error, 2=warn, 3=info, 4=debug
    * The value of GCLOUD_TRACE_LOGLEVEL takes precedence over this value.
@@ -166,6 +180,7 @@ export interface Config {
  * user-provided value will be used to extend the default value.
  */
 export const defaultConfig = {
+  clsMechanism: 'auto' as CLSMechanism,
   logLevel: 1,
   enabled: true,
   enhancedDatabaseReporting: false,

src/index.ts

@@ -27,9 +27,9 @@ if (!useAH) {
 }
 
 import * as common from '@google-cloud/common';
-import {cls} from './cls';
+import {cls, TraceCLSConfig, TraceCLSMechanism} from './cls';
 import {Constants} from './constants';
-import {Config, defaultConfig} from './config';
+import {Config, defaultConfig, CLSMechanism} from './config';
 import * as extend from 'extend';
 import * as path from 'path';
 import * as PluginTypes from './plugin-types';
@@ -56,6 +56,7 @@ for (let i = 0; i < filesLoadedBeforeTrace.length; i++) {
 interface TopLevelConfig {
   enabled: boolean;
   logLevel: number;
+  clsMechanism: CLSMechanism;
 }
 
 // PluginLoaderConfig extends TraceAgentConfig
@@ -178,12 +179,14 @@ export function start(projectConfig?: Config): PluginTypes.TraceAgent {
 
   try {
     // Initialize context propagation mechanism.
-    // TODO(kjin): Publicly expose this field.
-    cls.create(logger, {
-         mechanism: useAH ? 'async-hooks' : 'async-listener',
-         [FORCE_NEW]: config[FORCE_NEW]
-       })
-        .enable();
+    const m = config.clsMechanism;
+    const clsConfig: Forceable<TraceCLSConfig> = {
+      mechanism: m === 'auto' ? (useAH ? TraceCLSMechanism.ASYNC_HOOKS :
+                                         TraceCLSMechanism.ASYNC_LISTENER) :
+                                m as TraceCLSMechanism,
+      [FORCE_NEW]: config[FORCE_NEW]
+    };
+    cls.create(logger, clsConfig).enable();
 
     traceWriter.create(logger, config).initialize((err) => {
       if (err) {

src/trace-api.ts

@@ -128,7 +128,8 @@ export class TraceAgent implements TraceAgentInterface {
 
     // TODO validate options
     // Don't create a root span if we are already in a root span
-    if (cls.get().getContext().type === SpanDataType.ROOT) {
+    const rootSpan = cls.get().getContext();
+    if (rootSpan.type === SpanDataType.ROOT && !rootSpan.span.endTime) {
       this.logger!.warn(`TraceApi#runInRootSpan: [${
           this.pluginName}] Cannot create nested root spans.`);
       return fn(UNCORRELATED_SPAN);

test/plugins/test-trace-google-gax.ts

@@ -59,11 +59,17 @@ describe('Tracing with google-gax', () => {
   let googleGax: GaxModule;
 
   before(() => {
+    trace.setCLS();
     trace.setPluginLoader();
     trace.start();
     googleGax = require('./fixtures/google-gax0.16');
   });
 
+  after(() => {
+    trace.setCLS(trace.TestCLS);
+    trace.setPluginLoader(trace.TestPluginLoader);
+  });
+
   it(`doesn't break context`, (done) => {
     const authPromise = Promise.resolve(
         ((args, metadata, opts, cb) => {

test/test-cls.ts

@@ -24,7 +24,7 @@ import {ITestDefinition} from 'mocha';
 import * as semver from 'semver';
 import {inspect} from 'util';
 
-import {TraceCLS, TraceCLSConfig} from '../src/cls';
+import {TraceCLS, TraceCLSConfig, TraceCLSMechanism} from '../src/cls';
 import {AsyncHooksCLS} from '../src/cls/async-hooks';
 import {AsyncListenerCLS} from '../src/cls/async-listener';
 import {CLS} from '../src/cls/base';
@@ -221,8 +221,15 @@ describe('Continuation-Local Storage', () => {
 
   describe('TraceCLS', () => {
     const validTestCases: TraceCLSConfig[] = asyncAwaitSupported ?
-        [{mechanism: 'async-hooks'}, {mechanism: 'async-listener'}] :
-        [{mechanism: 'async-listener'}];
+        [
+          {mechanism: TraceCLSMechanism.ASYNC_HOOKS},
+          {mechanism: TraceCLSMechanism.ASYNC_LISTENER},
+          {mechanism: TraceCLSMechanism.NONE}
+        ] :
+        [
+          {mechanism: TraceCLSMechanism.ASYNC_LISTENER},
+          {mechanism: TraceCLSMechanism.NONE}
+        ];
     for (const testCase of validTestCases) {
       describe(`with configuration ${inspect(testCase)}`, () => {
         const logger = new TestLogger();

test/test-config-cls.ts

@@ -0,0 +1,87 @@
+/**
+ * Copyright 2018 Google LLC
+ *
+ * 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
+ *
+ *      http://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 {Logger} from '@google-cloud/common';
+import * as assert from 'assert';
+import * as semver from 'semver';
+import * as util from 'util';
+
+import {TraceCLSConfig, TraceCLSMechanism} from '../src/cls';
+
+import * as trace from './trace';
+
+describe('Behavior set by config for context propagation mechanism', () => {
+  const useAH = semver.satisfies(process.version, '>=8') &&
+      !!process.env.GCLOUD_TRACE_NEW_CONTEXT;
+  const autoMechanism =
+      useAH ? TraceCLSMechanism.ASYNC_HOOKS : TraceCLSMechanism.ASYNC_LISTENER;
+  let capturedConfig: TraceCLSConfig|null;
+
+  class CaptureConfigTestCLS extends trace.TestCLS {
+    constructor(logger: Logger, config: TraceCLSConfig) {
+      super(logger, config);
+      // Capture the config object passed into this constructor.
+      capturedConfig = config;
+    }
+  }
+
+  beforeEach(() => {
+    capturedConfig = null;
+  });
+
+  before(() => {
+    trace.setCLS(CaptureConfigTestCLS);
+  });
+
+  after(() => {
+    trace.setCLS(trace.TestCLS);
+  });
+
+  const testCases: Array<
+      {tracingConfig: trace.Config, contextPropagationConfig: TraceCLSConfig}> =
+      [
+        {
+          tracingConfig: {clsMechanism: 'none'},
+          contextPropagationConfig: {mechanism: 'none'}
+        },
+        {
+          tracingConfig: {clsMechanism: 'auto'},
+          contextPropagationConfig: {mechanism: autoMechanism}
+        },
+        {
+          tracingConfig: {},
+          contextPropagationConfig: {mechanism: autoMechanism}
+        },
+        {
+          // tslint:disable:no-any
+          tracingConfig: {clsMechanism: 'unknown' as any},
+          contextPropagationConfig: {mechanism: 'unknown' as any}
+          // tslint:enable:no-any
+        }
+      ];
+
+  for (const testCase of testCases) {
+    it(`should be as expected for config: ${
+           util.inspect(testCase.tracingConfig)}`,
+       () => {
+         trace.start(testCase.tracingConfig);
+         assert.ok(capturedConfig);
+         assert.strictEqual(
+             capturedConfig!.mechanism,
+             testCase.contextPropagationConfig.mechanism);
+       });
+  }
+});