fix: class-ify cls implementations (#708)

PR-URL: #708
googleapis · Apr 19, 2018 · 132db9b · 132db9b
1 parent 395a0c7
commit 132db9b
Show file tree

Hide file tree

Showing 18 changed files with 862 additions and 375 deletions.
diff --git a/src/cls-ah.ts b/src/cls-ah.ts
diff --git a/src/cls.ts b/src/cls.ts
@@ -1,5 +1,5 @@
 /**
- * Copyright 2015 Google Inc. All Rights Reserved.
+ * 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.
@@ -14,12 +14,17 @@
  * limitations under the License.
  */
 
-import * as CLS from 'continuation-local-storage';
+import {Logger} from '@google-cloud/common';
+import {EventEmitter} from 'events';
 import * as semver from 'semver';
 
+import {AsyncHooksCLS} from './cls/async-hooks';
+import {AsyncListenerCLS} from './cls/async-listener';
+import {CLS, Func} from './cls/base';
+import {UniversalCLS} from './cls/universal';
 import {SpanDataType} from './constants';
-import {UNCORRELATED_SPAN, UNTRACED_SPAN} from './span-data';
 import {Trace, TraceSpan} from './trace';
+import {Singleton} from './util';
 
 export interface RealRootContext {
   readonly span: TraceSpan;
@@ -43,78 +48,109 @@ export interface PhantomRootContext {
  */
 export type RootContext = RealRootContext|PhantomRootContext;
 
-export type Namespace = CLS.Namespace;
-export type Func<T> = CLS.Func<T>;
+const asyncHooksAvailable = semver.satisfies(process.version, '>=8');
 
-const useAsyncHooks: boolean = semver.satisfies(process.version, '>=8') &&
-    !!process.env.GCLOUD_TRACE_NEW_CONTEXT;
+export interface TraceCLSConfig { mechanism: 'async-listener'|'async-hooks'; }
 
-const cls: typeof CLS =
-    useAsyncHooks ? require('./cls-ah') : require('continuation-local-storage');
-
-const TRACE_NAMESPACE = 'com.google.cloud.trace';
+export interface CLSConstructor {
+  new(defaultContext: RootContext): CLS<RootContext>;
+}
 
 /**
- * Stack traces are captured when a root span is started. Because the stack
- * trace height varies on the context propagation mechanism, to keep published
- * stack traces uniform we need to remove the top-most frames when using the
- * c-l-s module. Keep track of this number here.
+ * An implementation of continuation-local storage for the Trace Agent.
+ * In addition to the underlying API, there is a guarantee that when an instance
+ * of this class is disabled, all context-manipulation methods will either be
+ * no-ops or pass-throughs.
  */
-export const ROOT_SPAN_STACK_OFFSET = useAsyncHooks ? 0 : 2;
+export class TraceCLS implements CLS<RootContext> {
+  private currentCLS: CLS<RootContext>;
+  // tslint:disable-next-line:variable-name CLSClass is a constructor.
+  private CLSClass: CLSConstructor;
+  private enabled = false;
+
+  private static UNCORRELATED: RootContext = {type: SpanDataType.UNCORRELATED};
+  private static UNTRACED: RootContext = {type: SpanDataType.UNTRACED};
+
+  /**
+   * Stack traces are captured when a root span is started. Because the stack
+   * trace height varies on the context propagation mechanism, to keep published
+   * stack traces uniform we need to remove the top-most frames when using the
+   * c-l-s module. Keep track of this number here.
+   */
+  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.`);
+        }
+        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.currentCLS = new UniversalCLS(TraceCLS.UNTRACED);
+    this.currentCLS.enable();
+  }
 
-export function createNamespace(): CLS.Namespace {
-  return cls.createNamespace(TRACE_NAMESPACE);
-}
+  isEnabled(): boolean {
+    return this.enabled;
+  }
 
-export function destroyNamespace(): void {
-  cls.destroyNamespace(TRACE_NAMESPACE);
-}
+  enable(): void {
+    if (!this.enabled) {
+      this.logger.info('TraceCLS#enable: Enabling CLS.');
+      this.enabled = true;
+      this.currentCLS.disable();
+      this.currentCLS = new this.CLSClass(TraceCLS.UNCORRELATED);
+      this.currentCLS.enable();
+    }
+  }
 
-export function getNamespace(): CLS.Namespace {
-  return cls.getNamespace(TRACE_NAMESPACE);
-}
+  disable(): void {
+    if (this.enabled) {
+      this.logger.info('TraceCLS#disable: Disabling CLS.');
+      this.enabled = false;
+      this.currentCLS.disable();
+      this.currentCLS = new UniversalCLS(TraceCLS.UNTRACED);
+      this.currentCLS.enable();
+    }
+  }
 
-/**
- * Get a RootContext object from continuation-local storage.
- */
-export function getRootContext(): RootContext {
-  // First getNamespace check is necessary in case any
-  // patched closures escaped before the agent was stopped and the
-  // namespace was destroyed.
-  const namespace = getNamespace();
-  if (namespace) {
-    // A few things can be going on here:
-    // 1. setRootContext has been called earlier to store a real root span
-    //    in continuation-local storage, so retrieve it.
-    // 2. setRootContext has been called earlier to explicitly specify that
-    //    the request corresponding to this continuation is _not_ being traced
-    //    (by being passed UNTRACED_SPAN), so retrieve it as well.
-    // 3. setRootContext has _never_ been called in this continuation. This
-    //    indicates that context was lost, and namespace.get('root') will
-    //    return null. Therefore, explicitly return UNCORRELATED_SPAN to
-    //    indicate that context was lost.
-    return namespace.get('root') || UNCORRELATED_SPAN;
-  } else {
-    // No namespace indicates that the Trace Agent is disabled. This is a
-    // special case where _all_ requests are explicitly not being traced,
-    // so return UNTRACED_SPAN to be consistent with that.
-    return UNTRACED_SPAN;
+  getContext(): RootContext {
+    return this.currentCLS.getContext();
   }
-}
 
-/**
- * Store a RootContext object in continuation-local storage.
- * @param rootContext Either a root span or UNTRACED_SPAN. It doesn't make
- * sense to pass UNCORRELATED_SPAN, which is a value specifically reserved for
- * when getRootContext is known to give an unusable value.
- */
-export function setRootContext(rootContext: RootContext): void {
-  getNamespace().set('root', rootContext);
-}
+  setContext(value: RootContext): void {
+    this.currentCLS.setContext(value);
+  }
 
-// This is only used in tests (and is temporary), so it doesn't apply in the
-// comment in getRootContext about the possible values of namespace.get('root').
-// It's functionally identical to setRootContext(null).
-export function clearRootContext(): void {
-  setRootContext(UNCORRELATED_SPAN);
+  runWithNewContext<T>(fn: Func<T>): T {
+    return this.currentCLS.runWithNewContext(fn);
+  }
+
+  bindWithCurrentContext<T>(fn: Func<T>): Func<T> {
+    return this.currentCLS.bindWithCurrentContext(fn);
+  }
+
+  patchEmitterToPropagateContext<T>(ee: EventEmitter): void {
+    this.currentCLS.patchEmitterToPropagateContext(ee);
+  }
 }
+
+export const cls = new Singleton(TraceCLS);