feat(sampler): use a reservoir sampler

Author: gergelyke

lib/agent/index.js

@@ -8,10 +8,11 @@ var assign = require('lodash.assign')
 var CollectorApi = require('./api')
 var Metrics = require('./metrics')
 var consts = require('../consts')
+var ReservoirSampler = require('./reservoir_sampler')
 
 var REQUEST_ID = 'request-id'
 var SPAN_ID = 'span-id'
-var MUST_COLLECT_LIMIT = 20
+var MUST_COLLECT_LIMIT = 100
 
 function Agent (options) {
   debug('Agent is initializing...')
@@ -27,7 +28,6 @@ function Agent (options) {
 
   // init required variables
   this.partials = {}
-  this.spans = []
 
   this.apmMetrics = Metrics.apm.create({
     collectorApi: this.collectorApi,
@@ -57,6 +57,8 @@ function Agent (options) {
     _this.serviceKey = serviceKey
     _this.start()
   })
+
+  this.reservoirSampler = new ReservoirSampler(MUST_COLLECT_LIMIT)
 }
 
 Agent.prototype.start = function () {
@@ -132,11 +134,7 @@ Agent.prototype.serverSend = function (data) {
   span.isForceSampled = span.isForceSampled || data.mustCollect === consts.MUST_COLLECT.ERROR
 
   if (span.isForceSampled) {
-    this.mustCollectCount += 1
-  }
-
-  if (span.isForceSampled && (this.mustCollectCount <= MUST_COLLECT_LIMIT)) {
-    this.spans.push(span)
+    this.reservoirSampler.addReturnsSuccess(span)
   }
 
   this.rpmMetrics.addResponseTime(data.responseTime)
@@ -251,7 +249,7 @@ Agent.prototype.onCrash = function (data) {
     }
   })
 
-  this.spans.push(span)
+  this.reservoirSampler.addReturnsSuccess(span)
 
   this._send({
     isSync: true
@@ -386,19 +384,18 @@ Agent.prototype.getMicrotime = function () {
 
 Agent.prototype._send = function (options) {
   debug('sending logs to the trace service')
-  if (this.spans.length > 0) {
+  var spans = this.reservoirSampler.getItems()
+  if (spans.length > 0) {
     var dataBag = {
       sample: {
         rate: 1,
         totalRequestCount: 1
       },
-      spans: this.spans
+      spans: spans
     }
 
     this.totalRequestCount = 0
-
-    this.spans = []
-    this.mustCollectCount = 0
+    this.reservoirSampler = new ReservoirSampler(MUST_COLLECT_LIMIT)
     this.collectorApi.sendSamples(dataBag, options && options.isSync)
   }
 }

lib/agent/index.spec.js

@@ -151,7 +151,7 @@ describe('The Trace agent', function () {
     expect(rpmMetrics.addStatusCode).to.have.been.calledWith(statusCode)
     expect(agent.totalRequestCount).to.eql(1)
     expect(agent.partials).to.eql({})
-    expect(agent.spans).to.eql([{
+    expect(agent.reservoirSampler.getItems()).to.eql([{
       requestId: transactionId,
       isForceSampled: true,
       isSampled: false,
@@ -195,7 +195,7 @@ describe('The Trace agent', function () {
     })
 
     expect(agent.totalRequestCount).to.eql(1)
-    expect(agent.spans).to.eql([])
+    expect(agent.reservoirSampler.getItems()).to.eql([])
     expect(agent.partials[transactionId]).to.eql({
       requestId: transactionId,
       isForceSampled: false,
@@ -237,7 +237,7 @@ describe('The Trace agent', function () {
     this.sandbox.stub(agent, 'getSpanId').returns(spanId)
     agent.report('obviously', { a: 'mock' })
     expect(agent.totalRequestCount).to.eql(1)
-    expect(agent.spans).to.eql([])
+    expect(agent.reservoirSampler.getItems()).to.eql([])
     expect(agent.partials[transactionId]).to.eql({
       requestId: transactionId,
       isForceSampled: false,
@@ -281,7 +281,7 @@ describe('The Trace agent', function () {
 
     agent.reportError('mocking a', new Error('with style'))
     expect(agent.totalRequestCount).to.eql(1)
-    expect(agent.spans).to.eql([])
+    expect(agent.reservoirSampler.getItems()).to.eql([])
     expect(agent.partials[transactionId]).to.eql({
       requestId: transactionId,
       isForceSampled: false,
@@ -311,17 +311,13 @@ describe('The Trace agent', function () {
   })
 
   it('passes sample and span data to the API client', function () {
-    agent.spans = [
-      1,
-      2,
-      3
-    ]
+    agent.reservoirSampler.addReturnsSuccess(1)
     agent.totalRequestCount = 6
     agent._send()
 
     expect(agent.totalRequestCount).to.eql(0)
     expect(collectorApi.sendSamples).to.be.calledWith({
-      spans: [1, 2, 3],
+      spans: [1],
       sample: {
         rate: 1,
         totalRequestCount: 1
@@ -331,7 +327,9 @@ describe('The Trace agent', function () {
 
   it('limits must-collect count to 20', function () {
     var statusCode = 200
-    agent.mustCollectCount = 20
+    for (var i = 0; i < 100; i += 1) {
+      agent.reservoirSampler.addReturnsSuccess(1)
+    }
     agent.serverReceive({
       id: transactionId,
       spanId: spanId,
@@ -349,7 +347,7 @@ describe('The Trace agent', function () {
       mustCollect: '1'
     })
 
-    expect(agent.spans).to.eql([])
-    expect(agent.mustCollectCount).to.eql(21)
+    expect(agent.reservoirSampler.getItems().length).to.eql(100)
+    expect(agent.reservoirSampler.itemsSeen).to.eql(101)
   })
 })

lib/agent/reservoir_sampler.js

@@ -0,0 +1,33 @@
+function Reservoir (limit) {
+  this.MAX_ITEMS = limit || 10
+
+  this.itemsSeen = 0
+  this.data = []
+}
+
+Reservoir.prototype.addReturnsSuccess = function (item) {
+  var inserted = true
+  if (this.itemsSeen < this.MAX_ITEMS) {
+    this.data.push(item)
+    inserted = true
+  } else {
+    inserted = this._replaceSampledItemReturnsSuccess(item)
+  }
+  this.itemsSeen++
+  return inserted
+}
+
+Reservoir.prototype.getItems = function () {
+  return this.data
+}
+
+Reservoir.prototype._replaceSampledItemReturnsSuccess = function (item) {
+  var toReplaceIndex = Math.floor(Math.random() * (this.itemsSeen + 2))
+  if (toReplaceIndex < this.MAX_ITEMS) {
+    this.data[toReplaceIndex] = item
+    return true
+  }
+  return false
+}
+
+module.exports = Reservoir

lib/agent/reservoir_sampler.spec.js

@@ -0,0 +1,40 @@
+var expect = require('chai').expect
+var ReservoirSampler = require('./reservoir_sampler')
+
+describe('The Reservoir Sampler', function () {
+  var itemToAdd = {
+    name: 'John Coin'
+  }
+
+  describe('when seen count is below the size limit', function () {
+    it('adds the item', function () {
+      var sampler = new ReservoirSampler(3)
+      var isAdded = sampler.addReturnsSuccess(itemToAdd)
+
+      expect(isAdded).to.eql(true)
+      expect(sampler.getItems()).to.eql([itemToAdd])
+    })
+  })
+
+  describe('when seen count is above the size limit', function () {
+    it('adds the item if selected', function () {
+      this.sandbox.stub(Math, 'random', function () {
+        return 0
+      })
+      var sampler = new ReservoirSampler(1)
+      sampler.itemsSeen = 1
+      var isAdded = sampler.addReturnsSuccess(itemToAdd)
+      expect(isAdded).to.eql(true)
+    })
+
+    it('discards the item if not selected', function () {
+      this.sandbox.stub(Math, 'random', function () {
+        return 1
+      })
+      var sampler = new ReservoirSampler(1)
+      sampler.itemsSeen = 1
+      var isAdded = sampler.addReturnsSuccess(itemToAdd)
+      expect(isAdded).to.eql(false)
+    })
+  })
+})