feat: add singular cls option (#748)

kjin · web-flow · commit 000643fe21ec · 2018-05-14T14:35:28.000-07:00
PR-URL: #748
diff --git a/src/cls.ts b/src/cls.ts
@@ -22,6 +22,7 @@ import {AsyncHooksCLS} from './cls/async-hooks';
 import {AsyncListenerCLS} from './cls/async-listener';
 import {CLS, Func} from './cls/base';
 import {NullCLS} from './cls/null';
+import {SingularCLS} from './cls/singular';
 import {SpanDataType} from './constants';
 import {SpanData, SpanOptions} from './plugin-types';
 import {Trace, TraceSpan} from './trace';
@@ -71,6 +72,11 @@ export enum TraceCLSMechanism {
    * Do not use any special mechanism to propagate root span context.
    * Only a single root span can be open at a time.
    */
+  SINGULAR = 'singular',
+  /**
+   * Do not write root span context; in other words, querying the current root
+   * span context will always result in a default value.
+   */
   NONE = 'none'
 }
 
@@ -120,6 +126,10 @@ export class TraceCLS implements CLS<RootContext> {
         this.CLSClass = AsyncListenerCLS;
         this.rootSpanStackOffset = 8;
         break;
+      case TraceCLSMechanism.SINGULAR:
+        this.CLSClass = SingularCLS;
+        this.rootSpanStackOffset = 4;
+        break;
       case TraceCLSMechanism.NONE:
         this.CLSClass = NullCLS;
         this.rootSpanStackOffset = 4;
diff --git a/src/cls/singular.ts b/src/cls/singular.ts
@@ -0,0 +1,66 @@
+/**
+ * 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 {EventEmitter} from 'events';
+
+import {CLS, Func} from './base';
+
+/**
+ * A trivial implementation of continuation-local storage where everything is
+ * in the same continuation. Therefore, only one unique value can be stored at
+ * a time.
+ */
+export class SingularCLS<Context> implements CLS<Context> {
+  private enabled = false;
+  private currentContext: Context;
+
+  constructor(private readonly defaultContext: Context) {
+    this.currentContext = this.defaultContext;
+  }
+
+  isEnabled(): boolean {
+    return this.enabled;
+  }
+
+  enable(): void {
+    this.enabled = true;
+  }
+
+  disable(): void {
+    this.enabled = false;
+    this.setContext(this.defaultContext);
+  }
+
+  getContext(): Context {
+    return this.currentContext;
+  }
+
+  setContext(value: Context): void {
+    if (this.enabled) {
+      this.currentContext = value;
+    }
+  }
+
+  runWithNewContext<T>(fn: Func<T>): T {
+    return fn();
+  }
+
+  bindWithCurrentContext<T>(fn: Func<T>): Func<T> {
+    return fn;
+  }
+
+  patchEmitterToPropagateContext(ee: EventEmitter): void {}
+}
diff --git a/src/config.ts b/src/config.ts
@@ -19,7 +19,7 @@ import * as path from 'path';
 const pluginDirectory =
     path.join(path.resolve(__dirname, '..'), 'src', 'plugins');
 
-export type CLSMechanism = 'none'|'auto';
+export type CLSMechanism = 'auto'|'none'|'singular';
 
 /** Available configuration options. */
 export interface Config {
@@ -30,6 +30,10 @@ export interface Config {
    *   _and_ the environment variable GCLOUD_TRACE_NEW_CONTEXT is set, in which
    *   case async_hooks will be used instead.
    * - 'none' disables CLS completely.
+   * - 'singular' allows one root span to exist at a time. This option is meant
+   *   to be used internally by Google Cloud Functions, or in any other
+   *   environment where it is guaranteed that only one request is being served
+   *   at a time.
    * The 'auto' mechanism is used by default if this configuration option is
    * not explicitly set.
    */
diff --git a/test/test-cls.ts b/test/test-cls.ts
@@ -25,6 +25,7 @@ import {AsyncHooksCLS} from '../src/cls/async-hooks';
 import {AsyncListenerCLS} from '../src/cls/async-listener';
 import {CLS} from '../src/cls/base';
 import {NullCLS} from '../src/cls/null';
+import {SingularCLS} from '../src/cls/singular';
 import {SpanDataType} from '../src/constants';
 import {createStackTrace, FORCE_NEW} from '../src/util';
 
@@ -236,36 +237,44 @@ describe('Continuation-Local Storage', () => {
            });
       });
     }
+
+    describe('SingularCLS', () => {
+      it('uses a single global context', async () => {
+        const cls = new SingularCLS('default');
+        cls.enable();
+        cls.runWithNewContext(() => {
+          cls.setContext('modified');
+        });
+        await Promise.resolve();
+        cls.runWithNewContext(() => {
+          assert.strictEqual(cls.getContext(), 'modified');
+        });
+      });
+    });
   });
 
   describe('TraceCLS', () => {
     const validTestCases:
-        Array<{config: TraceCLSConfig, expectedDefaultType: SpanDataType}> =
-            asyncAwaitSupported ?
-        [
-          {
-            config: {mechanism: TraceCLSMechanism.ASYNC_HOOKS},
-            expectedDefaultType: SpanDataType.UNCORRELATED
-          },
+        Array<{config: TraceCLSConfig, expectedDefaultType: SpanDataType}> = [
           {
             config: {mechanism: TraceCLSMechanism.ASYNC_LISTENER},
             expectedDefaultType: SpanDataType.UNCORRELATED
           },
           {
-            config: {mechanism: TraceCLSMechanism.NONE},
-            expectedDefaultType: SpanDataType.UNTRACED
-          }
-        ] :
-        [
-          {
-            config: {mechanism: TraceCLSMechanism.ASYNC_LISTENER},
+            config: {mechanism: TraceCLSMechanism.SINGULAR},
             expectedDefaultType: SpanDataType.UNCORRELATED
           },
           {
             config: {mechanism: TraceCLSMechanism.NONE},
             expectedDefaultType: SpanDataType.UNTRACED
           }
         ];
+    if (asyncAwaitSupported) {
+      validTestCases.push({
+        config: {mechanism: TraceCLSMechanism.ASYNC_HOOKS},
+        expectedDefaultType: SpanDataType.UNCORRELATED
+      });
+    }
     for (const testCase of validTestCases) {
       describe(`with configuration ${inspect(testCase)}`, () => {
         const logger = new TestLogger();
diff --git a/test/test-config-cls.ts b/test/test-config-cls.ts
@@ -67,6 +67,10 @@ describe('Behavior set by config for context propagation mechanism', () => {
           tracingConfig: {},
           contextPropagationConfig: {mechanism: autoMechanism}
         },
+        {
+          tracingConfig: {clsMechanism: 'singular'},
+          contextPropagationConfig: {mechanism: 'singular'}
+        },
         {
           // tslint:disable:no-any
           tracingConfig: {clsMechanism: 'unknown' as any},
