Refactor public api to match plugin api (#393)

PR-URL: #393

Fixes #330

index.js

@@ -24,11 +24,12 @@ require('continuation-local-storage');
 
 var common = require('@google-cloud/common');
 var constants = require('./src/constants.js');
+var TraceLabels = require('./src/trace-labels.js');
 var gcpMetadata = require('gcp-metadata');
 var semver = require('semver');
 var traceUtil = require('./src/util.js');
-var SpanData = require('./src/span-data.js');
 var util = require('util');
+var PluginApi = require('./src/trace-plugin-interface.js');
 
 var modulesLoadedBeforeTrace = [];
 
@@ -43,31 +44,23 @@ for (var i = 0; i < filesLoadedBeforeTrace.length; i++) {
 var onUncaughtExceptionValues = ['ignore', 'flush', 'flushAndExit'];
 
 /**
- * Phantom implementation of the trace agent. This allows API users to decouple
+ * Phantom implementation of the trace api. This allows API users to decouple
  * the enable/disable logic from the calls to the tracing API. The phantom API
  * has a lower overhead than isEnabled checks inside the API functions.
  * @private
  */
 var phantomTraceAgent = {
-  startSpan: function() { return SpanData.nullSpan; },
-  endSpan: function(spanData) { spanData.close(); },
-  runInSpan: function(name, labels, fn) {
-    if (typeof(labels) === 'function') {
-      fn = labels;
-    }
-    fn(function() {});
-  },
-  runInRootSpan: function(name, labels, fn) {
-    if (typeof(labels) === 'function') {
-      fn = labels;
-    }
-    fn(function() {});
-  },
-  setTransactionName: function() {},
-  addTransactionLabel: function() {}
+  enhancedDatabaseReportingEnabled: function() { return false; },
+  runInRootSpan: function(opts, fn) { return fn(PluginApi.nullSpan); },
+  createChildSpan: function(opts) { return PluginApi.nullSpan; },
+  wrap: function(fn) { return fn; },
+  wrapEmitter: function(ee) { return ee; },
+  constants: constants,
+  labels: TraceLabels
 };
 
 /** @private */
+// This can be either a PhantomAgent or PluginAPI object.
 var agent = phantomTraceAgent;
 
 var initConfig = function(projectConfig) {
@@ -91,36 +84,36 @@ var initConfig = function(projectConfig) {
 };
 
 /**
- * The singleton public agent. This is the public API of the module.
+ * The singleton public api. This is the public API of the module.
  */
 var publicAgent = {
   isActive: function() {
     return agent !== phantomTraceAgent;
   },
 
-  startSpan: function(name, labels) {
-    return agent.startSpan(name, labels);
+  enhancedDatabaseReportingEnabled: function() {
+    return agent.enhancedDatabaseReportingEnabled();
   },
 
-  endSpan: function(spanData, labels) {
-    return agent.endSpan(spanData, labels);
+  runInRootSpan: function(opts, fn) {
+    return agent.runInRootSpan(opts, fn);
   },
 
-  runInSpan: function(name, labels, fn) {
-    return agent.runInSpan(name, labels, fn);
+  createChildSpan: function(opts) {
+    return agent.createChildSpan(opts);
   },
 
-  runInRootSpan: function(name, labels, fn) {
-    return agent.runInRootSpan(name, labels, fn);
+  wrap: function(fn) {
+    return agent.wrap(fn);
   },
 
-  setTransactionName: function(name) {
-    return agent.setTransactionName(name);
+  wrapEmitter: function(ee) {
+    return agent.wrapEmitter(ee);
   },
 
-  addTransactionLabel: function(key, value) {
-    return agent.addTransactionLabel(key, value);
-  },
+  constants: constants,
+
+  labels: TraceLabels,
 
   start: function(projectConfig) {
     var config = initConfig(projectConfig);
@@ -191,7 +184,7 @@ var publicAgent = {
             'service. Please provide a valid project number as an env. ' +
             'variable, or through config.projectId passed to start(). ' + err);
           if (that.isActive()) {
-            agent.stop();
+            agent.agent_.stop();
             agent = phantomTraceAgent;
           }
           return;
@@ -204,7 +197,7 @@ var publicAgent = {
       return this;
     }
 
-    agent = require('./src/trace-agent.js').get(config, logger);
+    agent = new PluginApi(require('./src/trace-agent.js').get(config, logger));
     return this; // for chaining
   },
 
@@ -219,7 +212,7 @@ var publicAgent = {
    * For use in tests only.
    * @private
    */
-  private_: function() { return agent; }
+  private_: function() { return agent.agent_; }
 };
 
 /**

src/span-data.js

@@ -156,13 +156,6 @@ function StackFrame(className, methodName, fileName, lineNumber, columnNumber) {
   }
 }
 
-SpanData.nullSpan = {
-  createChildSpanData: function() { return SpanData.nullSpan; },
-  addLabel: function() {},
-  addLabels: function() {},
-  close: function() {}
-};
-
 /**
  * Export SpanData.
  */

src/trace-agent.js

@@ -126,94 +126,6 @@ TraceAgent.prototype.endSpan = function(spanData, labels) {
   spanData.close();
 };
 
-/**
- * Run the provided function in a new span with the provided name.
- * If the provided function accepts a parameter, it is assumed to be
- * async and is given a continuation to terminate the span after its
- * work is completed.
- * @param {string} name The name of the resulting span.
- * @param {Object<string, string}>=} labels Labels to be attached
- *   to the resulting span. Non-object data types are silently ignored.
- * @param {function(function()=)} fn The function to trace.
- */
-TraceAgent.prototype.runInSpan = function(name, labels, fn) {
-  if (typeof(labels) === 'function') {
-    fn = labels;
-    labels = undefined;
-  }
-  var span = this.startSpan(name, labels, 1);
-  if (fn.length === 0) {
-    fn();
-    this.endSpan(span);
-  } else {
-    fn(this.endSpan.bind(this, span));
-  }
-};
-
-/**
- * Run the provided function in a new root span with the provided
- * name. As with runInSpan, if the provided function accepts a parameter,
- * it is assumed to be asynchronous, and a callback will be passed
- * to the function to terminate the root span.
- * @param {string} name The name of the resulting root span.
- * @param {Object<string, string}>=} labels Labels to be attached
- *   to the resulting span.
- * @param {function(function()=)} fn The function to trace.
- */
-TraceAgent.prototype.runInRootSpan = function(name, labels, fn) {
-  if (typeof(labels) === 'function') {
-    fn = labels;
-    labels = undefined;
-  }
-  var self = this;
-  this.namespace.run(function () {
-    if (cls.getRootContext()) {
-      self.logger.error('Can\'t nest root spans');
-      return;
-    }
-    var span = self.createRootSpanData(name, null, null, 3, 'SPAN_KIND_UNSPECIFIED');
-    span.addLabels(labels);
-    if (fn.length === 0) {
-      fn();
-      span.close();
-    } else {
-      fn(function () {
-        span.close();
-      });
-    }
-  });
-};
-
-/**
- * Set the name of the root transaction.
- * @param {string} name The new name for the current root transaction.
- */
-TraceAgent.prototype.setTransactionName = function(name) {
-  var rootSpan = cls.getRootContext();
-  if (rootSpan === SpanData.nullSpan) {
-    return;
-  }
-  if (rootSpan) {
-    rootSpan.span.name = name;
-  } else {
-    this.logger.error('Cannot set transaction name without an active transaction');
-  }
-};
-
-/**
- * Add a new key value label to the root transaction.
- * @param {string} key The key for the new label.
- * @param {string} value The value for the new label.
- */
-TraceAgent.prototype.addTransactionLabel = function(key, value) {
-  var rootSpan = cls.getRootContext();
-  if (rootSpan) {
-    rootSpan.addLabel(key, value);
-  } else {
-    this.logger.error('Cannot add label without an active transaction');
-  }
-};
-
 /**
  * Determines whether a trace of the given name should be recorded based
  * on the current tracing policy.

src/trace-labels.js

@@ -19,8 +19,7 @@
 /**
  * Well-known trace span label values.
  */
-function TraceLabels() {
-}
+var TraceLabels = {};
 
 
 /**

src/trace-plugin-interface.js

@@ -100,6 +100,12 @@ RootSpan.prototype.getTraceContext = function() {
   return this.serializedTraceContext_;
 };
 
+var nullSpan = {
+  addLabel: function(k, v) {},
+  endSpan: function() {},
+  getTraceContext: function() { return ''; }
+};
+
 /**
  * PluginAPI constructor. Don't call directly - a plugin object will be passed to
  * plugin themselves
@@ -130,43 +136,6 @@ PluginAPI.prototype.summarizeDatabaseResults = function(res) {
          this.agent_.config_.databaseResultReportingSize);
 };
 
-/**
- * Creates and returns a new RootSpan object corresponding to an incoming
- * request.
- * @param {object} options An object that specifies options for how the root
- * span is created and propogated.
- * @param {string} options.name The name to apply to the root span.
- * @param {?string} options.url A URL associated with the root span, if
- * applicable.
- * @param {?string} options.traceContext The serialized form of an object that
- * contains information about an existing trace context.
- * @param {?number} options.skipFrames The number of stack frames to skip when
- * collecting call stack information for the root span, starting from the top;
- * this should be set to avoid including frames in the plugin. Defaults to 0.
- * @returns A new RootSpan object, or null if the trace agent's policy has
- * disabled tracing for the given set of options.
- */
-PluginAPI.prototype.createRootSpan = function(options) {
-  var skipFrames = options.skipFrames ? options.skipFrames + 1 : 1;
-  return createRootSpan_(this, options, skipFrames);
-};
-
-/**
- * Returns a RootSpan object that corresponds to a root span started earlier
- * in the same context, or null if one doesn't exist.
- * @returns A new RootSpan object, or null if a root span doesn't exist in
- * the current context.
- */
-PluginAPI.prototype.getRootSpan = function() {
-  if (cls.getRootContext()) {
-    return new RootSpan(this.agent_, cls.getRootContext());
-  } else {
-    this.logger_.warn('Attempted to get root span when it doesn\'t' + 
-      ' exist');
-    return null;
-  }
-};
-
 /**
  * Runs the given function in a root span corresponding to an incoming request,
  * possibly passing it an object that exposes an interface for adding labels
@@ -187,6 +156,10 @@ PluginAPI.prototype.runInRootSpan = function(options, fn) {
       'root span.');
     return fn(null);
   }
+  if (cls.getRootContext()) {
+    this.logger_.warn('Trace agent: Cannot create nested root spans.');
+    return fn(null);
+  }
   return this.agent_.namespace.runAndReturn(function() {
     var skipFrames = options.skipFrames ? options.skipFrames + 2 : 2;
     var rootSpan = createRootSpan_(that, options, skipFrames);
@@ -195,15 +168,14 @@ PluginAPI.prototype.runInRootSpan = function(options, fn) {
 };
 
 /**
- * Creates and returns a new ChildSpan object nested within the root span object
- * returned by getRootSpan. If there is no current RootSpan object, this
- * function returns null.
+ * Creates and returns a new ChildSpan object nested within the root span. If
+ * there is no current RootSpan object, this function returns null.
  * @param {object} options An object that specifies options for how the child
  * span is created and propogated.
  * @returns A new ChildSpan object, or null if there is no active root span.
  */
 PluginAPI.prototype.createChildSpan = function(options) {
-  var rootSpan = this.getRootSpan();
+  var rootSpan = getRootSpan_(this);
   if (rootSpan) {
     options = options || {};
     var childContext = this.agent_.startSpan(options.name, {},
@@ -247,6 +219,8 @@ PluginAPI.prototype.constants = constants;
 
 PluginAPI.prototype.labels = TraceLabels;
 
+PluginAPI.nullSpan = nullSpan;
+
 module.exports = PluginAPI;
 
 // Module-private functions
@@ -270,3 +244,13 @@ function createRootSpan_(api, options, skipFrames) {
     skipFrames + 1);
   return new RootSpan(api.agent_, rootContext);
 }
+
+function getRootSpan_(api) {
+  if (cls.getRootContext()) {
+    return new RootSpan(api.agent_, cls.getRootContext());
+  } else {
+    api.logger_.warn('Attempted to get root span when it doesn\'t' +
+      ' exist');
+    return null;
+  }
+}

test/hooks/common.js

@@ -47,11 +47,18 @@ function init(agent) {
   };
 }
 
+function replaceFunction(target, prop, fn) {
+  var old = target[prop];
+  target[prop] = fn;
+  return old;
+}
+
 function replaceDebugLogger(agent, fn) {
-  var privateAgent = agent.private_();
-  var oldDebug = privateAgent.logger.debug;
-  privateAgent.logger.debug = fn;
-  return oldDebug;
+  return replaceFunction(agent.private_().logger, 'debug', fn);
+}
+
+function replaceTracingPolicy(agent, fn) {
+  return replaceFunction(agent.private_(), 'policy', fn);
 }
 
 /**
@@ -248,6 +255,7 @@ module.exports = {
   runInTransaction: runInTransaction,
   getShouldTraceArgs: getShouldTraceArgs,
   replaceDebugLogger: replaceDebugLogger,
+  replaceTracingPolicy: replaceTracingPolicy,
   createRootSpanData: createRootSpanData,
   clearNamespace: clearNamespace,
   getConfig: getConfig,