Add trace ID to log metadata when available

packages/logging-bunyan/src/index.js

@@ -40,6 +40,13 @@ var BUNYAN_TO_STACKDRIVER = {
   10: 'DEBUG'
 };
 
+/**
+ * Key to use in the Bunyan payload to allow users to indicate a trace for the
+ * request, and to store as an intermediate value on the log entry before it
+ * gets written to the Stackdriver logging API.
+ */
+var LOGGING_TRACE_KEY = 'logging.googleapis.com/trace';
+
 /**
  * This module provides support for streaming your Bunyan logs to
  * [Stackdriver Logging](https://cloud.google.com/logging).
@@ -133,8 +140,6 @@ LoggingBunyan.prototype.formatEntry_ = function(record) {
     );
   }
 
-  record = extend({}, record);
-
   // Stackdriver Log Viewer picks up the summary line from the 'message' field
   // of the payload. Unless the user has provided a 'message' property also,
   // move the 'msg' to 'message'.
@@ -173,9 +178,34 @@ LoggingBunyan.prototype.formatEntry_ = function(record) {
     delete record.httpRequest;
   }
 
+  if (record[LOGGING_TRACE_KEY]) {
+    entryMetadata.trace = record[LOGGING_TRACE_KEY];
+    delete record[LOGGING_TRACE_KEY];
+  }
+
   return this.log_.entry(entryMetadata, record);
 };
 
+/**
+ * Intercept log entries as they are written so we can attempt to add the trace
+ * ID in the same continuation as the function that wrote the log, because the
+ * trace agent currently uses continuation local storage for the trace context.
+ *
+ * By the time the Writable stream buffer gets flushed and _write gets called
+ * we may well be in a different continuation.
+ */
+LoggingBunyan.prototype.write = function(record, encoding, callback) {
+  record = extend({}, record);
+  if (!record[LOGGING_TRACE_KEY]) {
+    var trace = logging.getCurrentTraceFromAgent();
+    if (trace) {
+      record[LOGGING_TRACE_KEY] = trace;
+    }
+  }
+
+  Writable.prototype.write.call(this, record, encoding, callback);
+}
+
 /**
  * Relay a log entry to the logging agent. This is called by bunyan through
  * Writable#write.

packages/logging-bunyan/test/index.js

@@ -223,6 +223,93 @@ describe('logging-bunyan', function() {
 
       loggingBunyan.formatEntry_(recordWithRequest);
     });
+
+    it('should promote prefixed trace property to metadata', function(done) {
+      var recordWithRequest = extend({
+        'logging.googleapis.com/trace': 'trace1'
+      }, RECORD);
+
+      loggingBunyan.log_.entry = function(entryMetadata, record) {
+        assert.deepStrictEqual(entryMetadata, {
+          resource: loggingBunyan.resource_,
+          timestamp: RECORD.time,
+          severity: LoggingBunyan.BUNYAN_TO_STACKDRIVER[RECORD.level],
+          trace: 'trace1'
+        });
+        assert.deepStrictEqual(record, RECORD);
+        done();
+      };
+
+      loggingBunyan.formatEntry_(recordWithRequest);
+    });
+  });
+
+  describe('write', function() {
+    var oldWritableWrite;
+    var oldTraceAgent;
+
+    beforeEach(function() {
+      oldWritableWrite = FakeWritable.prototype.write;
+      oldTraceAgent = global._google_trace_agent;
+    });
+
+    afterEach(function() {
+      FakeWritable.prototype.write = oldWritableWrite;
+      global._google_trace_agent = oldTraceAgent;
+    });
+
+    it('should not set trace property if trace unavailable', function(done) {
+      global._google_trace_agent = undefined;
+      FakeWritable.prototype.write = function(record, encoding, callback) {
+        assert.deepStrictEqual(record, RECORD);
+        assert.strictEqual(encoding, 'encoding');
+        assert.strictEqual(callback, assert.ifError);
+        assert.strictEqual(this, loggingBunyan);
+        done();
+      };
+
+      loggingBunyan.write(RECORD, 'encoding', assert.ifError);
+    });
+
+    it('should set prefixed trace property if trace available', function(done) {
+      global._google_trace_agent = {
+        getCurrentContextId: function () { return 'trace1'; },
+        getWriterProjectId: function () { return 'project1'; }
+      };
+      const recordWithoutTrace = extend({}, RECORD);
+      const recordWithTrace = extend({
+        'logging.googleapis.com/trace': 'projects/project1/traces/trace1'
+      }, RECORD);
+
+      FakeWritable.prototype.write = function(record, encoding, callback) {
+        // Check that trace field added to record before calling Writable.write
+        assert.deepStrictEqual(record, recordWithTrace);
+
+        // Check that the original record passed in was not mutated
+        assert.deepStrictEqual(recordWithoutTrace, RECORD);
+
+        done();
+      };
+
+      loggingBunyan.write(recordWithoutTrace, '', assert.ifError);
+    });
+
+    it('should leave prefixed trace property as is if set', function(done) {
+      global._google_trace_agent = {
+        getCurrentContextId: function () { return 'trace-from-agent'; },
+        getWriterProjectId: function () { return 'project1'; }
+      };
+      const recordWithTraceAlreadySet = extend({
+        'logging.googleapis.com/trace': 'trace-already-set'
+      }, RECORD);
+
+      FakeWritable.prototype.write = function(record, encoding, callback) {
+        assert.deepStrictEqual(record, recordWithTraceAlreadySet);
+        done();
+      };
+
+      loggingBunyan.write(recordWithTraceAlreadySet, '', assert.ifError);
+    });
   });
 
   describe('_write', function() {
@@ -303,4 +390,4 @@ describe('logging-bunyan', function() {
       });
     });
   });
-});
+});

packages/logging-winston/src/index.js

@@ -59,6 +59,11 @@ var STACKDRIVER_LOGGING_LEVEL_CODE_TO_NAME = {
   7: 'debug'
 };
 
+/**
+ * Log entry data key to allow users to indicate a trace for the request.
+ */
+var LOGGING_TRACE_KEY = 'logging.googleapis.com/trace';
+
 /**
  * This module provides support for streaming your winston logs to
  * [Stackdriver Logging](https://cloud.google.com/logging).
@@ -198,6 +203,16 @@ LoggingWinston.prototype.log = function(levelName, msg, metadata, callback) {
     }
   }
 
+  if (metadata && metadata[LOGGING_TRACE_KEY]) {
+    entryMetadata.trace = metadata[LOGGING_TRACE_KEY];
+    delete data.metadata[LOGGING_TRACE_KEY];
+  } else {
+    var trace = logging.getCurrentTraceFromAgent();
+    if (trace) {
+      entryMetadata.trace = trace;
+    }
+  }
+
   var entry = this.log_.entry(entryMetadata, data);
   this.log_[stackdriverLevel](entry, callback);
 };

packages/logging-winston/test/index.js

@@ -308,6 +308,48 @@ describe('logging-winston', function() {
       loggingWinston.log(LEVEL, MESSAGE, metadataWithRequest, assert.ifError);
     });
 
+    it('should promote prefixed trace property to metadata', function(done) {
+      const metadataWithTrace = extend({
+        'logging.googleapis.com/trace': 'trace1'
+      }, METADATA);
+
+      loggingWinston.log_.entry = function(entryMetadata, data) {
+        assert.deepStrictEqual(entryMetadata, {
+          resource: loggingWinston.resource_,
+          trace: 'trace1'
+        });
+        assert.deepStrictEqual(data, {
+          message: MESSAGE,
+          metadata: METADATA
+        });
+        done();
+      };
+      loggingWinston.log(LEVEL, MESSAGE, metadataWithTrace, assert.ifError);
+    });
+
+    it('should set trace metadata from agent if available', function(done) {
+      var oldTraceAgent = global._google_trace_agent;
+      global._google_trace_agent = {
+        getCurrentContextId: function () { return 'trace1'; },
+        getWriterProjectId: function () { return 'project1'; }
+      };
+      loggingWinston.log_.entry = function(entryMetadata, data) {
+        assert.deepStrictEqual(entryMetadata, {
+          resource: loggingWinston.resource_,
+          trace: 'projects/project1/traces/trace1'
+        });
+        assert.deepStrictEqual(data, {
+          message: MESSAGE,
+          metadata: METADATA
+        });
+        done();
+      };
+
+      loggingWinston.log(LEVEL, MESSAGE, METADATA, assert.ifError);
+
+      global._google_trace_agent = oldTraceAgent;
+    });
+
     it('should write to the log', function(done) {
       var entry = {};
 

packages/logging/src/index.js

@@ -773,6 +773,41 @@ Logging.prototype.setAclForTopic_ = function(name, config, callback) {
   });
 };
 
+/**
+ * Attempts to get the current fully qualified trace ID from the
+ * @google-cloud/trace-agent (if installed) in the format
+ * "projects/[PROJECT-ID]/traces/[TRACE-ID]".
+ *
+ * That is the specified format to set the LogEntry.trace field so that the
+ * Cloud Console logs and trace viewer can correlate log entries and associate
+ * them with traces.
+ *
+ * This returns null if the trace agent is not installed, there is no available
+ * trace context or no trace writer project ID.
+ */
+function getCurrentTraceFromAgent() {
+  var traceAgent = global._google_trace_agent;
+  if (!traceAgent || !traceAgent.getCurrentContextId ||
+      !traceAgent.getWriterProjectId) {
+    // Trace agent is not installed or is an older version
+    return null;
+  }
+
+  var traceId = traceAgent.getCurrentContextId();
+  if (!traceId) {
+    // No trace context because request not sampled or context lost
+    return null;
+  }
+
+  var traceProjectId = traceAgent.getWriterProjectId();
+  if (!traceProjectId) {
+    // No project, possibly because async project auto-discovery not yet done
+    return null;
+  }
+
+  return `projects/${traceProjectId}/traces/${traceId}`;
+}
+
 /*! Developer Documentation
  *
  * These methods can be auto-paginated.
@@ -797,6 +832,7 @@ Logging.Entry = Entry;
 Logging.Log = Log;
 Logging.Logging = Logging;
 Logging.Sink = Sink;
+Logging.getCurrentTraceFromAgent = getCurrentTraceFromAgent;
 
 module.exports = Logging;
 module.exports.v2 = v2;

packages/logging/test/index.js

@@ -1379,4 +1379,51 @@ describe('Logging', function() {
       });
     });
   });
+
+  describe('getCurrentTraceFromAgent', function() {
+    var oldTraceAgent;
+
+    beforeEach(function() {
+      oldTraceAgent = global._google_trace_agent;
+    });
+
+    afterEach(function() {
+      global._google_trace_agent = oldTraceAgent;
+    });
+
+    it('returns null if there is no trace agent', function() {
+      global._google_trace_agent = undefined;
+      assert.strictEqual(Logging.getCurrentTraceFromAgent(), null);
+    });
+
+    it('returns null if agent does not support context methods', function() {
+      global._google_trace_agent = {};
+      assert.strictEqual(Logging.getCurrentTraceFromAgent(), null);
+    });
+
+    it('returns null if context id is missing', function() {
+      global._google_trace_agent = {
+        getCurrentContextId: function () { return null; },
+        getWriterProjectId: function () { return 'project1'; }
+      };
+      assert.strictEqual(Logging.getCurrentTraceFromAgent(), null);
+    });
+
+    it('returns null if project id is missing', function() {
+      global._google_trace_agent = {
+        getCurrentContextId: function () { return 'trace1'; },
+        getWriterProjectId: function () { return null; }
+      };
+      assert.strictEqual(Logging.getCurrentTraceFromAgent(), null);
+    });
+
+    it('returns trace path if project and context available', function() {
+      global._google_trace_agent = {
+        getCurrentContextId: function () { return 'trace1'; },
+        getWriterProjectId: function () { return 'project1'; }
+      };
+      assert.strictEqual(Logging.getCurrentTraceFromAgent(),
+          'projects/project1/traces/trace1');
+    });
+  });
 });