Add trace ID to log metadata when available (#2574)

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).
@@ -143,8 +150,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'.
@@ -181,9 +186,58 @@ 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);
 };
 
+/**
+ * Gets the current fully qualified trace ID when available from the
+ * @google-cloud/trace-agent library in the LogEntry.trace field format of:
+ * "projects/[PROJECT-ID]/traces/[TRACE-ID]".
+ */
+function getCurrentTraceFromAgent() {
+  var agent = global._google_trace_agent;
+  if (!agent || !agent.getCurrentContextId || !agent.getWriterProjectId) {
+    return null;
+  }
+
+  var traceId = agent.getCurrentContextId();
+  if (!traceId) {
+    return null;
+  }
+
+  var traceProjectId = agent.getWriterProjectId();
+  if (!traceProjectId) {
+    return null;
+  }
+
+  return `projects/${traceProjectId}/traces/${traceId}`;
+}
+
+/**
+ * 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 = 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

@@ -231,6 +231,135 @@ 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);
+
+        assert.strictEqual(encoding, 'encoding');
+        assert.strictEqual(callback, assert.ifError);
+        assert.strictEqual(this, loggingBunyan);
+        done();
+      };
+
+      loggingBunyan.write(recordWithoutTrace, 'encoding', assert.ifError);
+    });
+
+    it('should leave prefixed trace property as is if set', function(done) {
+      var oldTraceAgent = global._google_trace_agent;
+      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);
+        assert.strictEqual(encoding, '');
+        assert.strictEqual(callback, assert.ifError);
+        assert.strictEqual(this, loggingBunyan);
+        done();
+      };
+
+      loggingBunyan.write(recordWithTraceAlreadySet, '', assert.ifError);
+
+      global._google_trace_agent = oldTraceAgent;
+    });
+  });
+
+  it('should not set prefixed trace property if trace unavailable', function() {
+    FakeWritable.prototype.write = function(record, encoding, callback) {
+      assert.deepStrictEqual(record, RECORD);
+      assert.strictEqual(encoding, '');
+      assert.strictEqual(callback, assert.ifError);
+      assert.strictEqual(this, loggingBunyan);
+    };
+    var oldTraceAgent = global._google_trace_agent;
+
+    global._google_trace_agent = {};
+    loggingBunyan.write(RECORD, '', assert.ifError);
+
+    global._google_trace_agent = {
+      getCurrentContextId: function() { return null; },
+      getWriterProjectId: function() { return null; }
+    };
+    loggingBunyan.write(RECORD, '', assert.ifError);
+
+    global._google_trace_agent = {
+      getCurrentContextId: function() { return null; },
+      getWriterProjectId: function() { return 'project1'; }
+    };
+    loggingBunyan.write(RECORD, '', assert.ifError);
+
+    global._google_trace_agent = {
+      getCurrentContextId: function() { return 'trace1'; },
+      getWriterProjectId: function() { return null; }
+    };
+    loggingBunyan.write(RECORD, '', assert.ifError);
+
+    global._google_trace_agent = oldTraceAgent;
   });
 
   describe('_write', function() {
@@ -311,4 +440,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).
@@ -135,6 +140,30 @@ function LoggingWinston(options) {
 winston.transports.StackdriverLogging = LoggingWinston;
 util.inherits(LoggingWinston, winston.Transport);
 
+/**
+ * Gets the current fully qualified trace ID when available from the
+ * @google-cloud/trace-agent library in the LogEntry.trace field format of:
+ * "projects/[PROJECT-ID]/traces/[TRACE-ID]".
+ */
+function getCurrentTraceFromAgent() {
+  var agent = global._google_trace_agent;
+  if (!agent || !agent.getCurrentContextId || !agent.getWriterProjectId) {
+    return null;
+  }
+
+  var traceId = agent.getCurrentContextId();
+  if (!traceId) {
+    return null;
+  }
+
+  var traceProjectId = agent.getWriterProjectId();
+  if (!traceProjectId) {
+    return null;
+  }
+
+  return `projects/${traceProjectId}/traces/${traceId}`;
+}
+
 /**
  * Relay a log entry to the logging agent. This is normally called by winston.
  *
@@ -206,6 +235,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 = getCurrentTraceFromAgent();
+    if (trace) {
+      entryMetadata.trace = trace;
+    }
+  }
+
   var entry = this.log_.entry(entryMetadata, data);
   this.log_[stackdriverLevel](entry, callback);
 };

packages/logging-winston/test/index.js

@@ -318,6 +318,85 @@ 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 leave out trace metadata if trace unavailable', function() {
+      loggingWinston.log_.entry = function(entryMetadata, data) {
+        assert.deepStrictEqual(entryMetadata, {
+          resource: loggingWinston.resource_,
+        });
+        assert.deepStrictEqual(data, {
+          message: MESSAGE,
+          metadata: METADATA
+        });
+      };
+
+      var oldTraceAgent = global._google_trace_agent;
+
+      global._google_trace_agent = {};
+      loggingWinston.log(LEVEL, MESSAGE, METADATA, assert.ifError);
+
+      global._google_trace_agent = {
+        getCurrentContextId: function() { return null; },
+        getWriterProjectId: function() { return null; }
+      };
+      loggingWinston.log(LEVEL, MESSAGE, METADATA, assert.ifError);
+
+      global._google_trace_agent = {
+        getCurrentContextId: function() { return null; },
+        getWriterProjectId: function() { return 'project1'; }
+      };
+      loggingWinston.log(LEVEL, MESSAGE, METADATA, assert.ifError);
+
+      global._google_trace_agent = {
+        getCurrentContextId: function() { return 'trace1'; },
+        getWriterProjectId: function() { return null; }
+      };
+      loggingWinston.log(LEVEL, MESSAGE, METADATA, assert.ifError);
+
+      global._google_trace_agent = oldTraceAgent;
+    });
+
     it('should write to the log', function(done) {
       var entry = {};