Repository: honeycombio/beeline-nodejs
Pull request #125 (2019-03-02 18:08:53)

Add presend and sampler hooks implementation.
### Aim
To allow users of the library to define their own `samplerHook` function that takes an event's data and can make a sampling decision based on that. Also allow a `presendHook` function to be defined and called to enrich an event with further context.

### Reasoning (for samplerHook)
Dynamic sampling is something that some of the other Honeycomb libraries are capable of and it is really helpful!

### Discussion
Moved all sampling behaviour to `finishSpan()`, following the pattern demonstrated in the beelines-python library ([see here]). This will allow users to define custom sampler hooks that meet a standard interface. This change required refactor of existing deterministic sampler (and its tests) to fit with the new interface requirements.

### What's missing?
1. I haven't added any documentation. Usage docs don't appear to live in this repo.
2. There are no tests for what happens if you do not provide a samplerHook or a sampleRate. See roughly line 224 in `lib/api/libhoney.js` for more detail. I wasn't sure the best way to verify that `send()` was called rather than `sendPresampled()`.

[see here]: https://github.com/honeycombio/beeline-python/blob/master/beeline/__init__.py#L179

allyjweir: Thanks for the thorough review @toshok. I'm on GMT (just wrapping my day up, 6:30pm here) so I will address these things tomorrow morning and hopefully will be ready for you to take a look again.

As you can see, I've pushed up the addition of a `presendHook` implementation as well. This was something I needed to add and thought it aligned closely with the existing PR and easily reviewed together.
allyjweir: @toshok What further work needs done to get this across the line and merged?
toshok: nothing! :) 🎉 

Patch:
diff --git a/CONTRIBUTORS b/CONTRIBUTORS
index 7fac20f..74d5a4a 100644
--- a/CONTRIBUTORS
+++ b/CONTRIBUTORS
@@ -3,3 +3,4 @@ honeycombio/beeline-nodejs contributors:
 Chris Toshok
 Christine Yen
 Erwin van der Koogh
+Ally Weir
\ No newline at end of file
diff --git a/lib/api/libhoney.js b/lib/api/libhoney.js
index dc576b7..37c3915 100644
--- a/lib/api/libhoney.js
+++ b/lib/api/libhoney.js
@@ -7,7 +7,7 @@ const libhoney = require("libhoney").default,
   url = require("url"),
   tracker = require("../async_tracker"),
   instrumentation = require("../instrumentation"),
-  DeterministicSampler = require("../deterministic_sampler"),
+  deterministicSampler = require("../deterministic_sampler"),
   schema = require("../schema"),
   pkg = require(path.join(__dirname, "..", "..", "package.json")),
   debug = require("debug")(`${pkg.name}:event`);
@@ -18,21 +18,41 @@ const incr = (payload, key, val = 1) => {
   payload[key] = (payload[key] || 0) + val;
 };
 
+const getFilteredOptions = opts => {
+  let filteredOptions = {};
+  Object.keys(opts).forEach(optionKey => {
+    if (optionKey !== "samplerHook" || optionKey !== "presendHook") {
+      filteredOptions[optionKey] = opts[optionKey];
+    }
+  });
+  return filteredOptions;
+};
+
 module.exports = 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 opts.samplerHook === "function") {
+      this.samplerHook = opts.samplerHook;
+      debug("using custom samplerHook");
+    } else if (opts.samplerHook !== undefined) {
+      debug(".samplerHook must be a function. ignoring");
+    }
+
+    if (typeof opts.samplerHook !== "function") {
+      const sampleRate = opts.sampleRate;
+
+      if (typeof sampleRate === "number") {
+        this.samplerHook = deterministicSampler(sampleRate);
+        debug("using deterministic sampler with .sampleRate provided.");
+      } else if (typeof sampleRate !== "undefined") {
+        debug(".sampleRate must be a number.  ignoring.");
       }
-    });
-    if (typeof sampleRate === "number") {
-      this.ds = new DeterministicSampler(sampleRate);
-    } else if (typeof sampleRate !== "undefined") {
-      debug(".sampleRate must be a number.  ignoring.");
+    }
+
+    if (typeof opts.presendHook === "function") {
+      this.presendHook = opts.presendHook;
+      debug("using custom presendHook");
+    } else if (opts.presendHook !== undefined) {
+      debug(".presendHook must be a function. ignoring");
     }
 
     let apiHost = process.env["HONEYCOMB_API_HOST"] || "https://api.honeycomb.io";
@@ -58,7 +78,7 @@ module.exports = class LibhoneyEventAPI {
           dataset: process.env["HONEYCOMB_DATASET"] || defaultName,
           userAgentAddition: `honeycomb-beeline/${pkg.version}`,
         },
-        optsWithoutSampleRate
+        getFilteredOptions(opts)
       )
     );
     this.honey.add({
@@ -70,11 +90,6 @@ module.exports = class LibhoneyEventAPI {
   startTrace(metadataContext, traceId, parentSpanId) {
     const 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;
-    }
-
     tracker.setTracked({
       id,
       customContext: {},
@@ -205,9 +220,27 @@ module.exports = class LibhoneyEventAPI {
         [schema.BEELINE_VERSION]: pkg.version,
         [schema.NODE_VERSION]: process.version,
       });
-      if (this.ds) {
-        ev.sampleRate = this.ds.sampleRate;
+
+      if (this.samplerHook) {
+        debug(`executing sampler hook on event ev = ${JSON.stringify(ev.data)}`);
+        const samplerHookResponse = this.samplerHook(ev.data);
+        const keepEvent = samplerHookResponse.shouldSample;
+        if (!keepEvent) {
+          debug(`skipping event due to sampler hook sample ev = ${JSON.stringify(ev.data)}`);
+          return;
+        }
+
+        if (typeof samplerHookResponse.sampleRate === "number") {
+          ev.sampleRate = samplerHookResponse.sampleRate;
+        }
+      }
+
+      if (this.presendHook) {
+        debug(`executing presend hook on event ev = ${JSON.stringify(ev.data)}`);
+        this.presendHook(ev);
       }
+
+      debug(`enqueuing presampled event ev = ${JSON.stringify(ev.data)}`);
       ev.sendPresampled();
     });
   }
diff --git a/lib/deterministic_sampler.js b/lib/deterministic_sampler.js
index fdf0106..0b8f9d0 100644
--- a/lib/deterministic_sampler.js
+++ b/lib/deterministic_sampler.js
@@ -1,18 +1,19 @@
 /* eslint-env node */
 const createHash = require("crypto").createHash;
+const schema = require("./schema");
 
 const MAX_UINT32 = Math.pow(2, 32) - 1;
 
-module.exports = class DeterministicSampler {
-  constructor(sampleRate) {
-    this.sampleRate = sampleRate;
-    this.upperBound = (MAX_UINT32 / sampleRate) >>> 0;
-  }
-
-  sample(determinant) {
-    let sum = createHash("sha1")
-      .update(determinant)
+module.exports = function(sampleRate) {
+  return function(eventData) {
+    const sum = createHash("sha1")
+      .update(eventData[schema.TRACE_ID])
       .digest();
-    return sum.readUInt32BE(0) <= this.upperBound;
-  }
+    const upperBound = (MAX_UINT32 / sampleRate) >>> 0;
+
+    return {
+      shouldSample: sum.readUInt32BE(0) <= upperBound,
+      sampleRate: sampleRate,
+    };
+  };
 };


Test patch:
diff --git a/lib/api/index.test.js b/lib/api/index.test.js
index d85f4c8..52ec9b6 100644
--- a/lib/api/index.test.js
+++ b/lib/api/index.test.js
@@ -4,8 +4,11 @@ const path = require("path"),
   schema = require("../schema"),
   tracker = require("../async_tracker"),
   propagation = require("../propagation"),
+  deterministicSampler = require("../deterministic_sampler"),
   pkg = require(path.join(__dirname, "..", "..", "package.json"));
 
+jest.mock("../deterministic_sampler");
+
 beforeEach(() =>
   event.configure({ impl: "libhoney-event", transmission: "mock", writeKey: "abc123" })
 );
@@ -75,50 +78,199 @@ test("sample rates are propagated", () => {
   const honey = event._apiForTesting().honey;
   expect(tracker.getTracked()).toBeUndefined(); // context should be empty initially
 
-  let sampled = false;
-  while (!sampled) {
-    let requestPayload = event.startTrace({
-      [schema.EVENT_TYPE]: "source",
-      [schema.TRACE_SPAN_NAME]: "name",
+  let requestPayload = event.startTrace({
+    [schema.EVENT_TYPE]: "source",
+    [schema.TRACE_SPAN_NAME]: "name",
+  });
+
+  // starting a request creates a context
+  let context = tracker.getTracked();
+  expect(context).not.toBeUndefined();
+
+  // the stack consists solely of the request's event
+  expect(context.stack).toEqual([requestPayload]);
+  // and it should have these properties
+  expect(requestPayload[schema.EVENT_TYPE]).toBe("source");
+  let traceId = requestPayload[schema.TRACE_ID];
+  expect(traceId).not.toBeUndefined();
+  let startTime = requestPayload.startTime;
+  expect(startTime).not.toBeUndefined();
+
+  // libhoney shouldn't have been told to send anything yet
+  expect(honey.transmission.events).toEqual([]);
+
+  // finish the request
+  event.finishTrace(requestPayload);
+  expect(tracker.getTracked()).toBeUndefined(); // context should be cleared
+
+  // libhoney should have been told to send one event
+  expect(honey.transmission.events.length).toBe(1);
+  let sent = honey.transmission.events[0];
+  let postData = JSON.parse(sent.postData);
+  expect(sent.timestamp).toEqual(new Date(startTime));
+  expect(sent.sampleRate).toBe(10);
+  expect(postData[schema.EVENT_TYPE]).toBe("source");
+  expect(postData[schema.TRACE_SPAN_NAME]).toBe("name");
+  expect(postData[schema.TRACE_ID]).toBe(traceId);
+  expect(postData[schema.TRACE_PARENT_ID]).toBeUndefined();
+  expect(postData[schema.TRACE_SPAN_ID]).not.toBeUndefined();
+  expect(postData[schema.DURATION_MS]).not.toBeUndefined();
+});
+
+describe("sampler hook", () => {
+  describe("if samplerHook option provided", () => {
+    test("it enqueues event during finalSpan if shouldSample is true", () => {
+      const mockSamplerHook = jest.fn(() => ({
+        shouldSample: true,
+        sampleRate: 123,
+      }));
+      event._resetForTesting();
+      event.configure({
+        impl: "libhoney-event",
+        transmission: "mock",
+        writeKey: "abc123",
+        samplerHook: mockSamplerHook,
+      });
+
+      let eventPayload = event.startTrace({});
+
+      expect(mockSamplerHook).toHaveBeenCalledTimes(0);
+
+      event.finishTrace(eventPayload);
+
+      expect(mockSamplerHook).toHaveBeenCalledTimes(1);
+      expect(event._apiForTesting().honey.transmission.events).toHaveLength(1);
     });
-    if (requestPayload === null) {
-      continue;
-    }
-    sampled = true;
-
-    // starting a request creates a context
-    let context = tracker.getTracked();
-    expect(context).not.toBeUndefined();
-
-    // the stack consists solely of the request's event
-    expect(context.stack).toEqual([requestPayload]);
-    // and it should have these properties
-    expect(requestPayload[schema.EVENT_TYPE]).toBe("source");
-    let traceId = requestPayload[schema.TRACE_ID];
-    expect(traceId).not.toBeUndefined();
-    let startTime = requestPayload.startTime;
-    expect(startTime).not.toBeUndefined();
-
-    // libhoney shouldn't have been told to send anything yet
-    expect(honey.transmission.events).toEqual([]);
 
-    // finish the request
-    event.finishTrace(requestPayload);
-    expect(tracker.getTracked()).toBeUndefined(); // context should be cleared
-
-    // libhoney should have been told to send one event
-    expect(honey.transmission.events.length).toBe(1);
-    let sent = honey.transmission.events[0];
-    let postData = JSON.parse(sent.postData);
-    expect(sent.timestamp).toEqual(new Date(startTime));
-    expect(sent.sampleRate).toBe(10);
-    expect(postData[schema.EVENT_TYPE]).toBe("source");
-    expect(postData[schema.TRACE_SPAN_NAME]).toBe("name");
-    expect(postData[schema.TRACE_ID]).toBe(traceId);
-    expect(postData[schema.TRACE_PARENT_ID]).toBeUndefined();
-    expect(postData[schema.TRACE_SPAN_ID]).not.toBeUndefined();
-    expect(postData[schema.DURATION_MS]).not.toBeUndefined();
-  }
+    test("it does not enqueue event if shouldSample is false", () => {
+      const mockSamplerHook = jest.fn(() => ({
+        shouldSample: false,
+        sampleRate: 123,
+      }));
+      event._resetForTesting();
+      event.configure({
+        impl: "libhoney-event",
+        transmission: "mock",
+        writeKey: "abc123",
+        samplerHook: mockSamplerHook,
+      });
+
+      let eventPayload = event.startTrace({});
+      event.finishTrace(eventPayload);
+
+      expect(mockSamplerHook).toHaveBeenCalledTimes(1);
+      expect(event._apiForTesting().honey.transmission.events).toHaveLength(0);
+    });
+
+    test("it does not enqueue event if invalid samplerHook provided", () => {
+      const mockSamplerHook = jest.fn(() => ({
+        foo: "bar",
+        baz: false,
+      }));
+      event._resetForTesting();
+      event.configure({
+        impl: "libhoney-event",
+        transmission: "mock",
+        writeKey: "abc123",
+        samplerHook: mockSamplerHook,
+      });
+
+      let eventPayload = event.startTrace({});
+      event.finishTrace(eventPayload);
+
+      expect(mockSamplerHook).toHaveBeenCalledTimes(1);
+      expect(event._apiForTesting().honey.transmission.events).toHaveLength(0);
+    });
+
+    test("it falls back to deterministicSampler with sampleRate if non-function samplerHook provided", () => {
+      deterministicSampler.mockReset();
+      const mockSamplerHook = jest.fn(() => () => ({
+        shouldSample: true,
+        sampleRate: 999,
+      }));
+      deterministicSampler.mockImplementation(mockSamplerHook);
+      event._resetForTesting();
+      event.configure({
+        impl: "libhoney-event",
+        transmission: "mock",
+        writeKey: "abc123",
+        sampleRate: 12345,
+        samplerHook: "I am invalid as I am not a function",
+      });
+
+      let eventPayload = event.startTrace({});
+      event.finishTrace(eventPayload);
+
+      expect(deterministicSampler).toHaveBeenCalledTimes(1);
+      expect(deterministicSampler).toHaveBeenCalledWith(12345);
+    });
+
+    test("it updates the sampleRate of the event if shouldSample is true", () => {
+      const mockSamplerHook = jest.fn(() => ({
+        shouldSample: true,
+        sampleRate: 999,
+      }));
+      event._resetForTesting();
+      event.configure({
+        impl: "libhoney-event",
+        transmission: "mock",
+        writeKey: "abc123",
+        samplerHook: mockSamplerHook,
+      });
+
+      let eventPayload = event.startTrace({});
+      event.finishTrace(eventPayload);
+
+      expect(event._apiForTesting().honey.transmission.events[0].sampleRate).toEqual(999);
+    });
+  });
+
+  describe("if no samplerHook given but valid sampleRate provided", () => {
+    test("the default deterministicSampler hook initialised and used", () => {
+      deterministicSampler.mockReset();
+      const mockSamplerHook = jest.fn(() => () => ({
+        shouldSample: true,
+        sampleRate: 999,
+      }));
+      deterministicSampler.mockImplementation(mockSamplerHook);
+      event._resetForTesting();
+      event.configure({
+        impl: "libhoney-event",
+        transmission: "mock",
+        writeKey: "abc123",
+        sampleRate: 123,
+      });
+
+      let eventPayload = event.startTrace();
+
+      // initialisation
+      expect(deterministicSampler).toHaveBeenCalledTimes(1);
+      expect(deterministicSampler).toHaveBeenCalledWith(123);
+
+      event.finishTrace(eventPayload);
+
+      // verify initiliser not called again
+      expect(deterministicSampler).toHaveBeenCalledTimes(1);
+      // verify actual sampler hook returned from intiialiser called
+      expect(mockSamplerHook).toHaveBeenCalledTimes(1);
+    });
+  });
+
+  describe("if neither a samplerHook nor sampleRate given", () => {
+    test("send all events", () => {
+      event._resetForTesting();
+      event.configure({
+        impl: "libhoney-event",
+        transmission: "mock",
+        writeKey: "abc123",
+      });
+
+      let eventPayload = event.startTrace({});
+      event.finishTrace(eventPayload);
+
+      expect(event._apiForTesting().honey.transmission.events).toHaveLength(1);
+    });
+  });
 });
 
 test("ending events out of order isn't allowed", () => {
@@ -392,3 +544,44 @@ test("async spans share custom context", () => {
   let asyncData = JSON.parse(honey.transmission.events[0].postData);
   expect(asyncData["app.custom"]).toBe("value");
 });
+
+describe("presend hook", () => {
+  it("calls the presend hook function if provided", () => {
+    const mockPresendHook = jest.fn();
+    event._resetForTesting();
+    event.configure({
+      impl: "libhoney-event",
+      transmission: "mock",
+      writeKey: "abc123",
+      presendHook: mockPresendHook,
+    });
+
+    let eventPayload = event.startTrace({});
+
+    event.finishTrace(eventPayload);
+
+    expect(mockPresendHook).toHaveBeenCalledTimes(1);
+    expect(event._apiForTesting().honey.transmission.events).toHaveLength(1);
+  });
+
+  it("doesn't call the presend hook if the event is not being sampled", () => {
+    const mockPresendHook = jest.fn();
+    const mockSamplerHook = jest.fn(() => ({ shouldSample: false, sampleRate: 1 }));
+    event._resetForTesting();
+    event.configure({
+      impl: "libhoney-event",
+      transmission: "mock",
+      writeKey: "abc123",
+      samplerHook: mockSamplerHook,
+      presendHook: mockPresendHook,
+    });
+
+    let eventPayload = event.startTrace({});
+
+    event.finishTrace(eventPayload);
+
+    expect(mockSamplerHook).toHaveBeenCalledTimes(1);
+    expect(mockPresendHook).toHaveBeenCalledTimes(0);
+    expect(event._apiForTesting().honey.transmission.events).toHaveLength(0);
+  });
+});
diff --git a/lib/deterministic_sampler.test.js b/lib/deterministic_sampler.test.js
index 0ab0bbe..939c29e 100644
--- a/lib/deterministic_sampler.test.js
+++ b/lib/deterministic_sampler.test.js
@@ -1,13 +1,34 @@
 /* eslint-env node, jest */
-const DeterministicSampler = require("./deterministic_sampler");
+const deterministicSampler = require("./deterministic_sampler");
+const schema = require("./schema");
 
 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 testSampleRate = 17;
+  const testSampler = deterministicSampler(testSampleRate);
+
+  expect(
+    testSampler({
+      [schema.TRACE_ID]: "hello",
+    }).shouldSample
+  ).toBe(false);
+
+  expect(
+    testSampler({
+      [schema.TRACE_ID]: "hello",
+    }).shouldSample
+  ).toBe(false);
+
+  expect(
+    testSampler({
+      [schema.TRACE_ID]: "world",
+    }).shouldSample
+  ).toBe(false);
+
+  expect(
+    testSampler({
+      [schema.TRACE_ID]: "this5",
+    }).shouldSample
+  ).toBe(true);
 });
 
 const requestIDBytes = "abcdef0123456789";
@@ -43,11 +64,14 @@ test("it works (stastically)", () => {
   const testSampleRates = [1, 2, 10];
 
   for (let sampleRate of testSampleRates) {
-    let ds = new DeterministicSampler(sampleRate);
+    const testSampler = deterministicSampler(sampleRate);
     let nSampled = 0;
 
     for (let i = 0; i < nRequestIDs; i++) {
-      if (ds.sample(randomRequestID())) {
+      const testEventData = {
+        [schema.TRACE_ID]: randomRequestID(),
+      };
+      if (testSampler(testEventData).shouldSample) {
         nSampled++;
       }
     }