Repository: honeycombio/beeline-nodejs
Pull request #56 (2018-06-05 05:51:51)

Deterministic sampling
Add a deterministic sampler (sha1 request id, first 4 bytes treated as a BE uint32, compare that number to the sample rate's cutoff).

The sampling decision is made in `startRequest`.  If we decide to sample the request, we create the context and everything works as normal.  If we decide _not_ to sample it, `startRequest` returns null, and we don't create the context.  All the instrumentation treats this as a no-op, so for the cost of a single sampler invocation, we can sample the entire trace.

toshok: > Only having four examples makes me a bit nervous, though, especially when two of them are the same (why is "hello" there twice anyway? is the sampler stateful?)

I took that to be checking if it was deterministic - the result should be the same for the same string.  That said, that test could be a lot ... fuller.
toshok: as the commit message alludes, I'd misread the golang slice in honeycombio/honeytail and was reading the uint32 from the end of the hash, not the beginning.  Additionally, I was reading the uint32 in LE, not BE form.

tbh, the slice in honeytail isn't necessary, from this line: https://github.com/honeycombio/honeytail/blob/master/sample/deterministic_sampler.go#L49



Patch:
diff --git a/README.md b/README.md
index 7188875..28b47cd 100644
--- a/README.md
+++ b/README.md
@@ -139,7 +139,7 @@ Instrumented only for context propagation
 
 ## express
 
-Adds columns with prefix `express.`
+Adds columns with prefix `request.`
 
 ### Configuration options
 
@@ -171,13 +171,31 @@ Adds columns with prefix `http.`
 
 Adds columns with prefix `https.`
 
+## mongoose
+
+Instrumented only for context propagation
+
+## mongodb
+
+Adds columns with prefix `db.`
+
+### Configuration options
+
+| Name                       | Type    |
+| -------------------------- | ------- |
+| `mongodb.includeDocuments` | boolean |
+
+#### `mongodb.includeDocuments`
+
+If true, documents in the api will be JSON serialized and included in the events sent to honeycomb.
+
 ## mpromise
 
 Instrumented only for context propagation
 
 ## mysql2
 
-Adds columns with prefix `mysql2.`
+Adds columns with prefix `db.`
 
 ## react-dom/server
 
diff --git a/lib/deterministic_sampler.js b/lib/deterministic_sampler.js
new file mode 100644
index 0000000..a5ad464
--- /dev/null
+++ b/lib/deterministic_sampler.js
@@ -0,0 +1,17 @@
+/* global require module */
+const createHash = require("crypto").createHash;
+
+const MAX_UINT32 = Math.pow(2, 32) - 1;
+
+module.exports = class DeterministicSampler {
+  constructor(sampleRate) {
+    this.upperBound = (MAX_UINT32 / sampleRate) >>> 0;
+  }
+
+  sample(determinant) {
+    let sum = createHash("sha1")
+      .update(determinant)
+      .digest();
+    return sum.readUInt32BE(0) <= this.upperBound;
+  }
+};
diff --git a/lib/event.js b/lib/event.js
index a0a119f..f586bdf 100644
--- a/lib/event.js
+++ b/lib/event.js
@@ -6,6 +6,7 @@ const libhoney = require("libhoney").default,
   uuidv4 = require("uuid/v4"),
   tracker = require("./async_tracker"),
   instrumentation = require("./instrumentation"),
+  DeterministicSampler = require("./deterministic_sampler"),
   schema = require("./schema"),
   pkg = require(path.join(__dirname, "..", "package.json")),
   debug = require("debug")(`${pkg.name}:event`);
@@ -95,6 +96,21 @@ class MockEventAPI {
 
 class LibhoneyEventAPI {
   constructor(opts) {
+    let sampleRate;
+    let optsWithoutSampleRate = {};
+    Object.keys(opts).forEach(k => {
+      if (k === "sampleRate") {
+        sampleRate = opts[k];
+      } else {
+        optsWithoutSampleRate[k] = opts[k];
+      }
+    });
+    if (typeof sampleRate === "number") {
+      this.ds = new DeterministicSampler(sampleRate);
+    } else if (typeof sampleRate !== "undefined") {
+      debug(".sampleRate must be a number.  ignoring.");
+    }
+
     this.honey = new libhoney(
       Object.assign(
         {
@@ -103,7 +119,7 @@ class LibhoneyEventAPI {
           dataset: process.env["HONEYCOMB_DATASET"] || defaultName,
           userAgentAddition: `honeycomb-beeline/${pkg.version}`,
         },
-        opts
+        optsWithoutSampleRate
       )
     );
     this.honey.addField(schema.HOSTNAME, os.hostname());
@@ -111,6 +127,12 @@ class LibhoneyEventAPI {
 
   startRequest(source, name, traceId) {
     let id = traceId || uuidv4();
+
+    if (this.ds && !this.ds.sample(id)) {
+      // don't create the context at all if this request isn't going to send a trace
+      return null;
+    }
+
     const requestContext = {
       id,
       customContext: {},
diff --git a/lib/instrumentation.js b/lib/instrumentation.js
index e5d5ccb..a976b72 100644
--- a/lib/instrumentation.js
+++ b/lib/instrumentation.js
@@ -13,6 +13,8 @@ const instrumentedModules = new Set([
   "mysql2",
   "http",
   "https",
+  "mongodb",
+  "mongoose",
   "mpromise",
   "bluebird",
   "sequelize",
diff --git a/lib/instrumentation/express.js b/lib/instrumentation/express.js
index 5b2424c..09d0433 100644
--- a/lib/instrumentation/express.js
+++ b/lib/instrumentation/express.js
@@ -85,6 +85,12 @@ const getMagicMiddleware = ({ userContext, traceIdSource, packageVersion }) => (
   let traceIdContext = getTraceId(traceIdSource, req);
   let ev = event.startRequest("express", "request", traceIdContext.traceId);
 
+  if (!ev) {
+    // sampler has decided that we shouldn't trace this request
+    next();
+    return;
+  }
+
   event.addContext({
     [schema.TRACE_ID_SOURCE]: traceIdContext.source,
     [schema.PACKAGE_VERSION]: packageVersion,
diff --git a/lib/instrumentation/mongodb.js b/lib/instrumentation/mongodb.js
new file mode 100644
index 0000000..e54cae5
--- /dev/null
+++ b/lib/instrumentation/mongodb.js
@@ -0,0 +1,345 @@
+/* global require, module */
+const shimmer = require("shimmer"),
+  tracker = require("../async_tracker"),
+  event = require("../event"),
+  schema = require("../schema");
+
+function startCollectionEvent(context, packageVersion, name) {
+  let ev = event.startEvent(context, "mongodb", `collection.${name}`);
+  event.addContext({
+    [schema.PACKAGE_VERSION]: packageVersion,
+  });
+  return ev;
+}
+
+function finishCollectionEvent(ev, name) {
+  event.finishEvent(ev, `collection.${name}`);
+}
+
+function prefixKeys(prefix, map = {}) {
+  let prefixed = {};
+  Object.keys(map).forEach(k => {
+    prefixed[`${prefix}.${k}`] = map[k];
+  });
+  return prefixed;
+}
+
+function twoArgs(options, callback) {
+  if (typeof options === "function") (callback = options), (options = {});
+  return [options, callback];
+}
+
+function threeArgs(arg1, options, callback) {
+  if (typeof options === "function") (callback = options), (options = {});
+  return [arg1, options, callback];
+}
+
+function fourArgs(arg1, arg2, options, callback) {
+  if (typeof options === "function") (callback = options), (options = {});
+  return [arg1, arg2, options, callback];
+}
+
+function instrumentMongodb(mongodb, opts = {}) {
+  function wrap(name, populateArgs, additionalContext) {
+    shimmer.wrap(mongodb.Collection.prototype, name, function(original) {
+      return function(...args) {
+        let context = tracker.getTracked();
+        if (!context) {
+          return original.apply(this, args);
+        }
+
+        // callback always come last, and there's always options just before it
+        let populatedArgs = populateArgs(...args);
+        let callback = populatedArgs.pop(); // we'll put this back on (wrapped) when we original.apply below
+        let options = populatedArgs[populatedArgs.length - 1];
+
+        // filled in below the callback
+        let ev;
+
+        let wrapped_cb = tracker.bindFunction(function(...cb_args) {
+          finishCollectionEvent(ev, name);
+          if (callback) {
+            return callback(...cb_args);
+          }
+        });
+
+        ev = startCollectionEvent(context, opts.packageVersion, name);
+        event.addContext(prefixKeys("db.options", options));
+        if (additionalContext) {
+          event.addContext(prefixKeys("db", additionalContext(...populatedArgs)));
+        }
+        return original.apply(this, [...populatedArgs, wrapped_cb]);
+      };
+    });
+  }
+
+  wrap(
+    "find",
+    (query, options, callback) => {
+      let selector = query;
+      if (typeof callback !== "function") {
+        if (typeof options === "function") {
+          callback = options;
+          options = undefined;
+        } else if (options == null) {
+          callback = typeof selector === "function" ? selector : undefined;
+          selector = typeof selector === "object" ? selector : undefined;
+        }
+      }
+      return [selector, options, callback];
+    },
+    query => ({ query })
+  );
+
+  wrap("insertOne", threeArgs, document => (opts.includeDocuments ? { document } : {}));
+  wrap(
+    "insertMany",
+    threeArgs,
+    documents =>
+      opts.includeDocuments
+        ? {
+            documents,
+            documents_count: documents.length,
+          }
+        : {}
+  );
+  wrap(
+    "bulkWrite",
+    threeArgs,
+    operations =>
+      opts.includeDocuments && Array.isArray(operations)
+        ? {
+            operations,
+            operations_count: operations.length,
+          }
+        : {}
+  );
+  // .insert is just a wrapper around .insertMany
+
+  wrap("updateOne", threeArgs, (filter, update) => ({
+    filter,
+    update,
+  }));
+  wrap(
+    "replaceOne",
+    threeArgs,
+    (filter, document) =>
+      opts.includeDocuments
+        ? {
+            filter,
+            document,
+          }
+        : {
+            filter,
+          }
+  );
+  wrap("updateMany", threeArgs, (filter, update) => ({
+    filter,
+    update,
+  }));
+  wrap(
+    "update",
+    threeArgs,
+    (selector, document) =>
+      opts.includeDocuments
+        ? {
+            selector,
+            document,
+          }
+        : {
+            selector,
+          }
+  );
+
+  wrap("deleteOne", threeArgs, filter => ({
+    filter,
+  }));
+  wrap("deleteMany", threeArgs, filter => ({
+    filter,
+  }));
+  wrap("remove", threeArgs, selector => ({
+    selector,
+  }));
+
+  wrap("save", threeArgs, document => (opts.includeDocuments ? { document } : {}));
+  wrap(
+    "findOne",
+    (query, options, callback) => {
+      if (typeof query === "function") (callback = query), (query = {}), (options = {});
+      if (typeof options === "function") (callback = options), (options = {});
+      query = query || {};
+      options = options || {};
+      return [query, options, callback];
+    },
+    query => ({ query })
+  );
+  wrap("rename", threeArgs, newName => ({ newName }));
+
+  wrap("drop", twoArgs);
+  wrap("options", twoArgs);
+  wrap("isCapped", twoArgs);
+
+  wrap("createIndex", threeArgs, fieldOrSpec => ({
+    fieldOrSpec,
+  }));
+  wrap("createIndexes", threeArgs, indexSpecs => ({
+    indexSpecs,
+  }));
+  wrap(
+    "dropIndex",
+    (indexName, ...args) => {
+      let callback = typeof args[args.length - 1] === "function" ? args.pop() : undefined;
+      let options = args.length ? args.shift() || {} : {};
+      return [indexName, options, callback];
+    },
+    indexName => ({ indexName })
+  );
+  wrap("dropIndexes", twoArgs);
+  wrap("reIndex", twoArgs);
+  // XXX(toshok) listIndexes doesn't take a callback, it returns a cursor (ugggh)
+  wrap("ensureIndex", threeArgs, fieldOrSpec => ({
+    fieldOrSpec,
+  }));
+  wrap("indexExists", threeArgs, indexes => ({ indexes }));
+  wrap("indexInformation", twoArgs);
+
+  wrap(
+    "count",
+    (...args) => {
+      let callback = typeof args[args.length - 1] === "function" ? args.pop() : undefined;
+      let query = args.length ? args.shift() || {} : {};
+      let options = args.length ? args.shift() || {} : {};
+      return [query, options, callback];
+    },
+    query => ({
+      query,
+    })
+  );
+  wrap(
+    "distinct",
+    (key, ...args) => {
+      let callback = typeof args[args.length - 1] === "function" ? args.pop() : undefined;
+      let query = args.length ? args.shift() || {} : {};
+      let options = args.length ? args.shift() || {} : {};
+      return [key, query, options, callback];
+    },
+    (key, query) => ({
+      key,
+      query,
+    })
+  );
+  wrap("indexes", twoArgs);
+  wrap("stats", twoArgs);
+
+  wrap("findOneAndDelete", threeArgs, filter => ({ filter }));
+  wrap("findOneAndReplace", fourArgs, (filter, replacement) => ({
+    filter,
+    replacement,
+  }));
+  wrap("findOneAndUpdate", fourArgs, (filter, update) => ({
+    filter,
+    update,
+  }));
+  wrap(
+    "findAndModify",
+    (query, ...args) => {
+      let callback = typeof args[args.length - 1] === "function" ? args.pop() : undefined;
+      let sort = args.length ? args.shift() || [] : [];
+      let doc = args.length ? args.shift() : null;
+      let options = args.length ? args.shift() || {} : {};
+      return [query, sort, doc, options, callback];
+    },
+    (query, sort, document) =>
+      opts.includeDocuments
+        ? {
+            document,
+            query,
+            sort,
+          }
+        : {
+            query,
+            sort,
+          }
+  );
+  wrap(
+    "findAndRemove",
+    (query, ...args) => {
+      let callback = typeof args[args.length - 1] === "function" ? args.pop() : undefined;
+      let sort = args.length ? args.shift() || [] : [];
+      let options = args.length ? args.shift() || {} : {};
+      return [query, sort, options, callback];
+    },
+    (query, sort) => ({
+      query,
+      sort,
+    })
+  );
+
+  wrap(
+    "aggregate",
+    (...args) => {
+      let pipeline, options, callback;
+      if (Array.isArray(args[0])) {
+        pipeline = args[0];
+        options = args[1];
+        callback = args[2];
+
+        // Set up callback if one is provided
+        if (typeof options === "function") {
+          callback = options;
+          options = {};
+        }
+
+        // If we have no options or callback we are doing
+        // a cursor based aggregation
+        if (options == null && callback == null) {
+          options = {};
+        }
+      } else {
+        callback = args.pop();
+        // Get the possible options object
+        var opts = args[args.length - 1];
+        // If it contains any of the admissible options pop it of the args
+        options =
+          opts &&
+          (opts.readPreference ||
+            opts.explain ||
+            opts.cursor ||
+            opts.out ||
+            opts.maxTimeMS ||
+            opts.hint ||
+            opts.allowDiskUse)
+            ? args.pop()
+            : {};
+        // Left over arguments is the pipeline
+        pipeline = args;
+      }
+      return [pipeline, options, callback];
+    },
+    pipeline => ({ pipeline, pipeline_count: pipeline.length })
+  );
+
+  wrap("parallelCollectionScan", (options, callback) => {
+    if (typeof options === "function") (callback = options), (options = { numCursors: 1 });
+    return [options, callback];
+  });
+  wrap(
+    "geoHaystackSearch",
+    (x, y, ...args) => {
+      let callback = typeof args[args.length - 1] === "function" ? args.pop() : undefined;
+      let options = args.length ? args.shift() || {} : {};
+      return [x, y, options, callback];
+    },
+    (x, y) => ({
+      x,
+      y,
+    })
+  );
+
+  // XXX(toshok) .group has a lot of args (and multiple callbacks)
+  // XXX(toshok) .map also has a few callbacks
+
+  return mongodb;
+}
+
+module.exports = instrumentMongodb;
diff --git a/lib/instrumentation/mongoose.js b/lib/instrumentation/mongoose.js
new file mode 100644
index 0000000..5e55198
--- /dev/null
+++ b/lib/instrumentation/mongoose.js
@@ -0,0 +1,15 @@
+/* global require, module */
+const shimmer = require("shimmer"),
+  tracker = require("../async_tracker");
+
+let instrumentMongoose = function(mongoose) {
+  shimmer.wrap(mongoose.Model, "$wrapCallback", function(original) {
+    return function(callback) {
+      return original.apply(this, [tracker.bindFunction(callback)]);
+    };
+  });
+
+  return mongoose;
+};
+
+module.exports = instrumentMongoose;
diff --git a/package-lock.json b/package-lock.json
index 79a82c6..96f682a 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -1,6 +1,6 @@
 {
   "name": "honeycomb-beeline",
-  "version": "1.0.3",
+  "version": "1.1.0",
   "lockfileVersion": 1,
   "requires": true,
   "dependencies": {
diff --git a/package.json b/package.json
index 5fcca0a..e200092 100644
--- a/package.json
+++ b/package.json
@@ -3,7 +3,7 @@
   "description": "automatic instrumentation for honeycomb.io",
   "author": "support@honeycomb.io",
   "license": "Apache-2.0",
-  "version": "1.0.3",
+  "version": "1.1.0",
   "repository": {
     "type": "git",
     "url": "https://github.com/honeycombio/beeline-nodejs.git"


Test patch:
diff --git a/lib/deterministic_sampler.test.js b/lib/deterministic_sampler.test.js
new file mode 100644
index 0000000..6f436a5
--- /dev/null
+++ b/lib/deterministic_sampler.test.js
@@ -0,0 +1,68 @@
+/* global require test expect */
+const DeterministicSampler = require("./deterministic_sampler");
+
+test("it works (sample datapoints)", () => {
+  // taken from honeytail's test TestDeterministicSamplerDatapoints
+  let s = new DeterministicSampler(17);
+  expect(s.sample("hello")).toBe(false);
+  expect(s.sample("hello")).toBe(false);
+  expect(s.sample("world")).toBe(false);
+  expect(s.sample("this5")).toBe(true);
+});
+
+const requestIDBytes = "abcdef0123456789";
+
+function getRandomInt(max) {
+  return Math.floor(Math.random() * Math.floor(max));
+}
+
+function randomRequestID() {
+  // create request ID roughly resembling something you would get from
+  // AWS ALB, e.g.,
+  //
+  // 1-5ababc0a-4df707925c1681932ea22a20
+  //
+  // The AWS docs say the middle bit is "time in seconds since epoch",
+  // (implying base 10) but the above represents an actual Root= ID from
+  // an ALB access log, so... yeah.
+  let reqID = "1-";
+  for (let i = 0; i < 8; i++) {
+    reqID += requestIDBytes[getRandomInt(requestIDBytes.length)];
+  }
+  reqID += "-";
+  for (let i = 0; i < 24; i++) {
+    reqID += requestIDBytes[getRandomInt(requestIDBytes.length)];
+  }
+  return reqID;
+}
+
+test("it works (stastically)", () => {
+  const nRequestIDs = 50000;
+  const acceptableMarginOfError = 0.05;
+
+  const testSampleRates = [1, 2, 10];
+
+  for (let sampleRate of testSampleRates) {
+    let ds = new DeterministicSampler(sampleRate);
+    let nSampled = 0;
+
+    for (let i = 0; i < nRequestIDs; i++) {
+      if (ds.sample(randomRequestID())) {
+        nSampled++;
+      }
+    }
+
+    let expectedNSampled = nRequestIDs / sampleRate;
+
+    // Sampling should be balanced across all request IDs
+    // regardless of sample rate. If we cross this threshold, flunk
+    // the test.
+    let unacceptableLowBound =
+      (expectedNSampled - expectedNSampled * acceptableMarginOfError) >>> 0;
+    let unacceptableHighBound =
+      (expectedNSampled + expectedNSampled * acceptableMarginOfError) >>> 0;
+
+    expect(nSampled).toBeGreaterThanOrEqual(unacceptableLowBound);
+    expect(nSampled).toBeLessThanOrEqual(unacceptableHighBound);
+  }
+});