feat: add getProjectId and getCurrentRootSpan (#782)

kjin · web-flow · commit f7ae770b34be · 2018-06-22T10:57:48.000-07:00
diff --git a/src/cls.ts b/src/cls.ts
@@ -24,7 +24,8 @@ import {CLS, Func} from './cls/base';
 import {NullCLS} from './cls/null';
 import {SingularCLS} from './cls/singular';
 import {SpanType} from './constants';
-import {Span, SpanOptions} from './plugin-types';
+import {RootSpan} from './plugin-types';
+import {UNCORRELATED_ROOT_SPAN, UNTRACED_ROOT_SPAN} from './span-data';
 import {Trace, TraceSpan} from './trace';
 import {Singleton} from './util';
 
@@ -33,7 +34,6 @@ const asyncHooksAvailable = semver.satisfies(process.version, '>=8');
 export interface RealRootContext {
   readonly span: TraceSpan;
   readonly trace: Trace;
-  createChildSpan(options: SpanOptions): Span;
   readonly type: SpanType.ROOT;
 }
 
@@ -51,7 +51,7 @@ export interface PhantomRootContext {
  * When we store an actual root span, the only information we need is its
  * current trace/span fields.
  */
-export type RootContext = RealRootContext|PhantomRootContext;
+export type RootContext = RootSpan&(RealRootContext|PhantomRootContext);
 
 /**
  * An enumeration of the possible mechanisms for supporting context propagation
@@ -101,8 +101,8 @@ export class TraceCLS implements CLS<RootContext> {
   private CLSClass: CLSConstructor;
   private enabled = false;
 
-  static UNCORRELATED: RootContext = {type: SpanType.UNCORRELATED};
-  static UNTRACED: RootContext = {type: SpanType.UNTRACED};
+  static UNCORRELATED: RootContext = UNCORRELATED_ROOT_SPAN;
+  static UNTRACED: RootContext = UNTRACED_ROOT_SPAN;
 
   /**
    * Stack traces are captured when a root span is started. Because the stack
diff --git a/src/plugin-types.ts b/src/plugin-types.ts
@@ -125,6 +125,16 @@ export interface TraceAgent {
    */
   runInRootSpan<T>(options: RootSpanOptions, fn: (span: RootSpan) => T): T;
 
+  /**
+   * Gets the active root span for the current context. This method is
+   * guaranteed to return an object with the surface of a RootSpan object, but
+   * it may not represent a real root span if we are not in one. Use isRealSpan
+   * or check the `type` field to determine whether this is a real or phantom
+   * span.
+   * @returns An object that represents either a real or phantom root span.
+   */
+  getCurrentRootSpan(): RootSpan;
+
   /**
    * Returns a unique identifier for the currently active context. This can be
    * used to uniquely identify the current root span. If there is no current,
@@ -135,6 +145,12 @@ export interface TraceAgent {
    */
   getCurrentContextId(): string|null;
 
+  /**
+   * Returns the projectId that was either configured or auto-discovered by the
+   * TraceWriter.
+   */
+  getProjectId(): Promise<string>;
+
   /**
    * Returns the projectId that was either configured or auto-discovered by the
    * TraceWriter. Note that the auto-discovery is done asynchronously, so this
diff --git a/src/trace-api.ts b/src/trace-api.ts
@@ -16,15 +16,14 @@
 
 import {Logger} from '@google-cloud/common';
 import * as is from 'is';
-import * as semver from 'semver';
 import * as uuid from 'uuid';
 
 import {cls, RootContext} from './cls';
 import {Constants, SpanType} from './constants';
 import {Func, RootSpan, RootSpanOptions, Span, SpanOptions, TraceAgent as TraceAgentInterface} from './plugin-types';
-import {ChildSpanData, RootSpanData, UNCORRELATED_CHILD_SPAN, UNCORRELATED_ROOT_SPAN, UNTRACED_CHILD_SPAN, UNTRACED_ROOT_SPAN} from './span-data';
-import {SpanKind, Trace} from './trace';
+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 * as TracingPolicy from './tracing-policy';
 import * as util from './util';
 
@@ -35,7 +34,6 @@ import * as util from './util';
 export interface TraceAgentConfig extends TracingPolicy.TracePolicyConfig {
   enhancedDatabaseReporting: boolean;
   ignoreContextHeader: boolean;
-  projectId?: string;
 }
 
 interface IncomingTraceContext {
@@ -173,21 +171,33 @@ export class TraceAgent implements TraceAgentInterface {
     }, rootContext);
   }
 
-  getCurrentContextId(): string|null {
+  getCurrentRootSpan(): RootSpan {
     if (!this.isActive()) {
-      return null;
+      return UNTRACED_ROOT_SPAN;
     }
+    return cls.get().getContext();
+  }
 
-    const rootSpan = cls.get().getContext();
-    if (rootSpan.type === SpanType.ROOT) {
-      return rootSpan.trace.traceId;
+  getCurrentContextId(): string|null {
+    // In v3, this will be deprecated for getCurrentRootSpan.
+    const traceContext = this.getCurrentRootSpan().getTraceContext();
+    const parsedTraceContext = util.parseContextFromHeader(traceContext);
+    return parsedTraceContext ? parsedTraceContext.traceId : null;
+  }
+
+  getProjectId(): Promise<string> {
+    if (traceWriter.exists() && traceWriter.get().isActive) {
+      return traceWriter.get().getProjectId();
+    } else {
+      return Promise.reject(
+          new Error('The Project ID could not be retrieved.'));
     }
-    return null;
   }
 
   getWriterProjectId(): string|null {
-    if (this.config) {
-      return this.config.projectId || null;
+    // In v3, this will be deprecated for getProjectId.
+    if (traceWriter.exists() && traceWriter.get().isActive) {
+      return traceWriter.get().projectId;
     } else {
       return null;
     }
diff --git a/src/trace-writer.ts b/src/trace-writer.ts
@@ -19,7 +19,6 @@ import {AxiosError} from 'axios';
 import * as gcpMetadata from 'gcp-metadata';
 import {OutgoingHttpHeaders} from 'http';
 import * as os from 'os';
-import * as util from 'util';
 
 import {Constants} from './constants';
 import {SpanKind, Trace} from './trace';
@@ -28,6 +27,9 @@ import {Singleton} from './util';
 
 const pjson = require('../../package.json');
 
+// TODO(kjin): This value should be exported from @g-c/c.
+const NO_PROJECT_ID_TOKEN = '{{projectId}}';
+
 const onUncaughtExceptionValues = ['ignore', 'flush', 'flushAndExit'];
 
 const headers: OutgoingHttpHeaders = {};
@@ -52,9 +54,6 @@ export interface LabelObject { [key: string]: string; }
  * A class representing a service that publishes traces in the background.
  */
 export class TraceWriter extends common.Service {
-  // TODO(kjin): Make public members private (they're public for testing)
-  private logger: common.Logger;
-  private config: TraceWriterConfig;
   /** Stringified traces to be published */
   buffer: string[];
   /** Default labels to be attached to written spans */
@@ -71,7 +70,9 @@ export class TraceWriter extends common.Service {
    * @param logger The Trace Agent's logger object.
    * @constructor
    */
-  constructor(config: TraceWriterConfig, logger: common.Logger) {
+  constructor(
+      private readonly config: TraceWriterConfig,
+      private readonly logger: common.Logger) {
     super(
         {
           packageJson: pjson,
@@ -82,9 +83,6 @@ export class TraceWriter extends common.Service {
         config);
 
     this.logger = logger;
-    // Clone the config object
-    this.config = {...config};
-    this.config.serviceContext = {...this.config.serviceContext};
     this.buffer = [];
     this.defaultLabels = {};
 
@@ -214,13 +212,13 @@ export class TraceWriter extends common.Service {
   }
 
   getProjectId() {
-    if (this.config.projectId) {
-      return Promise.resolve(this.config.projectId);
+    // super.getProjectId writes to projectId, but doesn't check it first
+    // before going through the flow of obtaining it. So we add that logic
+    // first.
+    if (this.projectId !== NO_PROJECT_ID_TOKEN) {
+      return Promise.resolve(this.projectId);
     }
-    return super.getProjectId().then((projectId) => {
-      this.config.projectId = projectId;
-      return projectId;
-    });
+    return super.getProjectId();
   }
 
   /**
@@ -269,8 +267,8 @@ export class TraceWriter extends common.Service {
     // Any test that doesn't mock the Trace Writer will assume that traces get
     // buffered synchronously. We need to refactor those tests to remove that
     // assumption before we can make this fix.
-    if (this.config.projectId) {
-      afterProjectId(this.config.projectId);
+    if (this.projectId !== NO_PROJECT_ID_TOKEN) {
+      afterProjectId(this.projectId);
     } else {
       this.getProjectId().then(afterProjectId, (err: Error) => {
         // Because failing to get a project ID means that the trace agent will
@@ -329,7 +327,7 @@ export class TraceWriter extends common.Service {
    */
   publish(json: string) {
     const uri = `https://cloudtrace.googleapis.com/v1/projects/${
-        this.config.projectId}/traces`;
+        this.projectId}/traces`;
     const options = {method: 'PATCH', uri, body: json, headers};
     this.logger.info('TraceWriter#publish: Publishing to ' + uri);
     this.request(options, (err, body?, response?) => {
diff --git a/test/test-index.ts b/test/test-index.ts
@@ -22,14 +22,12 @@ import { SpanType } from '../src/constants';
 import { FORCE_NEW } from '../src/util';
 
 var assert = require('assert');
-var nock = require('nock');
-var nocks = require('./nocks'/*.js*/);
 var trace = require('../..');
 
 var disabledAgent: TraceAgent = trace.get();
 
 describe('index.js', function() {
-  it('should get a disabled agent with `Trace.get`', function() {
+  it('should get a disabled agent with `Trace.get`', async function() {
     assert.ok(!disabledAgent.isActive()); // ensure it's disabled first
     let ranInRootSpan = false;
     disabledAgent.runInRootSpan({ name: '' }, (span) => {
@@ -40,6 +38,10 @@ describe('index.js', function() {
     assert.strictEqual(disabledAgent.enhancedDatabaseReportingEnabled(), false);
     assert.strictEqual(disabledAgent.getCurrentContextId(), null);
     assert.strictEqual(disabledAgent.getWriterProjectId(), null);
+    assert.strictEqual(disabledAgent.getCurrentRootSpan().type, SpanType.UNTRACED);
+    // getting project ID should reject.
+    await disabledAgent.getProjectId().then(
+        () => Promise.reject(new Error()), () => Promise.resolve());
     assert.strictEqual(disabledAgent.createChildSpan({ name: '' }).type, SpanType.UNTRACED);
     assert.strictEqual(disabledAgent.getResponseTraceContext('', false), '');
     const fn = () => {};
diff --git a/test/test-trace-api.ts b/test/test-trace-api.ts
@@ -51,7 +51,10 @@ describe('Trace Interface', () => {
     testTraceModule.setCLSForTest(TraceCLS);
     cls.create({mechanism: TraceCLSMechanism.ASYNC_LISTENER}, logger).enable();
     traceWriter
-        .create(Object.assign({[FORCE_NEW]: true}, defaultConfig), logger)
+        .create(
+            Object.assign(
+                {[FORCE_NEW]: true, projectId: 'project-1'}, defaultConfig),
+            logger)
         .initialize(done);
   });
 
@@ -141,6 +144,17 @@ describe('Trace Interface', () => {
       assert.strictEqual(testTraceModule.getTraces().length, 1);
     });
 
+    it('should return a root span when getCurrentRootSpan is called', () => {
+      const traceAPI = createTraceAgent();
+      // When a root span isn't running, return UNCORRELATED.
+      assert.strictEqual(
+          traceAPI.getCurrentRootSpan().type, SpanType.UNCORRELATED);
+      traceAPI.runInRootSpan({name: 'root'}, (rootSpan) => {
+        assert.strictEqual(traceAPI.getCurrentRootSpan(), rootSpan);
+        rootSpan.endSpan();
+      });
+    });
+
     it('should return null context id when one does not exist', () => {
       const traceAPI = createTraceAgent();
       assert.strictEqual(traceAPI.getCurrentContextId(), null);
@@ -156,9 +170,14 @@ describe('Trace Interface', () => {
       });
     });
 
-    it('should return get the project ID if set in config', () => {
-      const config = {projectId: 'project-1'};
-      const traceApi = createTraceAgent(null /* policy */, config);
+    it('should return the project ID from the Trace Writer (promise api)',
+       async () => {
+         const traceApi = createTraceAgent();
+         assert.equal(await traceApi.getProjectId(), 'project-1');
+       });
+
+    it('should return get the project ID from the Trace Writer', () => {
+      const traceApi = createTraceAgent();
       assert.equal(traceApi.getWriterProjectId(), 'project-1');
     });
 
diff --git a/test/test-trace-uncaught-exception.ts b/test/test-trace-uncaught-exception.ts
@@ -74,7 +74,7 @@ describe('Trace Writer', () => {
      (done) => {
        const restoreOriginalUncaughtExceptionListeners =
            removeAllUncaughtExceptionListeners();
-       const traceApi = trace.start({onUncaughtException: 'flush'});
+       trace.start({onUncaughtException: 'flush', projectId: '0'});
        setImmediate(() => {
          setImmediate(() => {
            removeAllUncaughtExceptionListeners();
diff --git a/test/test-trace-writer.ts b/test/test-trace-writer.ts
@@ -91,7 +91,12 @@ describe('Trace Writer', () => {
     nock.disableNetConnect();
     oauth2Scope = oauth2().persist();
     shimmer.wrap(
-        Service.prototype, 'getProjectId', () => () => getProjectIdOverride());
+        Service.prototype, 'getProjectId', () => function(this: Service) {
+          return getProjectIdOverride().then(projectId => {
+            this.projectId = projectId;
+            return projectId;
+          });
+        });
   });
 
   after(() => {
@@ -117,7 +122,7 @@ describe('Trace Writer', () => {
       getProjectIdOverride = () => Promise.resolve('my-project');
       writer.initialize(err => {
         assert.ifError(err);
-        assert.strictEqual(writer.getConfig().projectId, 'my-project');
+        assert.strictEqual(writer.projectId, 'my-project');
         writer.stop();
         done();
       });
@@ -129,7 +134,7 @@ describe('Trace Writer', () => {
       getProjectIdOverride = () => Promise.resolve('my-different-project');
       writer.initialize(err => {
         assert.ifError(err);
-        assert.strictEqual(writer.getConfig().projectId, 'my-project');
+        assert.strictEqual(writer.projectId, 'my-project');
         writer.stop();
         done();
       });
