Support delayed sampling

.vscode/launch.json

@@ -4,7 +4,7 @@
     {
       "type": "node",
       "request": "launch",
-      "name": "Mocha Current File",
+      "name": "Debug Current File",
       "program": "${workspaceFolder}/node_modules/mocha/bin/_mocha",
       "args": ["--timeout", "999999", "--compilers", "js:babel-core/register", "--colors", "${file}"],
       "console": "integratedTerminal",

src/_flow/sampler.js

@@ -30,5 +30,5 @@ declare interface Sampler {
   onSetOperationName(span: Span, operationName: string): SamplingDecision;
   onSetTag(span: Span, key: string, value: any): SamplingDecision;
 
-  close(callback: ?Function): void;
+  close(callback: ?() => void): void;
 }

src/configuration.js

@@ -12,7 +12,7 @@
 
 import ConstSampler from './samplers/const_sampler';
 import ProbabilisticSampler from './samplers/probabilistic_sampler';
-import RateLimitingSampler from './samplers/ratelimiting_sampler';
+import RateLimitingSampler from './samplers/rate_limiting_sampler';
 import RemoteReporter from './reporters/remote_reporter';
 import CompositeReporter from './reporters/composite_reporter';
 import LoggingReporter from './reporters/logging_reporter';

src/index.js

@@ -19,7 +19,7 @@ import Tracer from './tracer';
 
 import ConstSampler from './samplers/const_sampler';
 import ProbabilisticSampler from './samplers/probabilistic_sampler';
-import RateLimitingSampler from './samplers/ratelimiting_sampler';
+import RateLimitingSampler from './samplers/rate_limiting_sampler';
 import RemoteSampler from './samplers/remote_sampler';
 
 import CompositeReporter from './reporters/composite_reporter';

src/reporters/composite_reporter.js

@@ -12,6 +12,7 @@
 // the License.
 
 import Span from '../span.js';
+import Utils from '../util.js';
 
 export default class CompositeReporter implements Reporter {
   _reporters: Array<Reporter>;
@@ -30,23 +31,9 @@ export default class CompositeReporter implements Reporter {
     });
   }
 
-  _compositeCallback(limit: number, callback: () => void): () => void {
-    let count = 0;
-    return () => {
-      count++;
-      if (count >= limit) {
-        callback();
-      }
-    };
-  }
-
   close(callback?: () => void): void {
-    const modifiedCallback = callback
-      ? this._compositeCallback(this._reporters.length, callback)
-      : function() {};
-    this._reporters.forEach(r => {
-      r.close(modifiedCallback);
-    });
+    const countdownCallback = Utils.countdownCallback(this._reporters.length, callback);
+    this._reporters.forEach(r => r.close(countdownCallback));
   }
 
   setProcess(serviceName: string, tags: Array<Tag>): void {

src/samplers/_adapt_sampler.js

@@ -15,7 +15,7 @@ import { SAMPLER_API_V2 } from './constants';
 import Span from '../span';
 import BaseSamplerV2 from './v2/base';
 
-function adaptSampler(sampler: any): ?Sampler {
+export function adaptSampler(sampler: any): ?Sampler {
   if (!sampler) {
     return null;
   }
@@ -38,7 +38,35 @@ export function adaptSamplerOrThrow(sampler: any): Sampler {
   return s;
 }
 
-export default adaptSampler;
+/**
+ * Convenience base class for simple samplers that implement isSampled() function
+ * that is not sensitive to its arguments.
+ */
+export default class LegacySamplerV1Base extends BaseSamplerV2 {
+  constructor(name: string) {
+    super(name);
+  }
+
+  isSampled(operationName: string, outTags: {}): boolean {
+    throw new Error('Subclass must override isSampled()');
+  }
+
+  onCreateSpan(span: Span): SamplingDecision {
+    const outTags = {};
+    const isSampled = this.isSampled(span.operationName, outTags);
+    return { sample: isSampled, retryable: false, tags: outTags };
+  }
+
+  onSetOperationName(span: Span, operationName: string): SamplingDecision {
+    const outTags = {};
+    const isSampled = this.isSampled(span.operationName, outTags);
+    return { sample: isSampled, retryable: false, tags: outTags };
+  }
+
+  onSetTag(span: Span, key: string, value: any): SamplingDecision {
+    return { sample: false, retryable: true, tags: null };
+  }
+}
 
 /**
  * Transforms legacy v1 sampler into V2.
@@ -50,28 +78,17 @@ export default adaptSampler;
  * where as onSetOperation() returns retryable=false, since that is what the tracer
  * used to do.
  */
-class LegacySamplerV1Adapter implements Sampler {
+class LegacySamplerV1Adapter extends LegacySamplerV1Base {
   apiVersion = SAMPLER_API_V2;
   _delegate: LegacySamplerV1;
 
   constructor(delegate: LegacySamplerV1) {
+    super(delegate.name());
     this._delegate = delegate;
   }
 
-  onCreateSpan(span: Span): SamplingDecision {
-    const outTags = {};
-    const isSampled = this._delegate.isSampled(span.operationName, outTags);
-    return { sample: isSampled, retryable: true, tags: outTags };
-  }
-
-  onSetOperationName(span: Span, operationName: string): SamplingDecision {
-    const outTags = {};
-    const isSampled = this._delegate.isSampled(span.operationName, outTags);
-    return { sample: isSampled, retryable: false, tags: outTags };
-  }
-
-  onSetTag(span: Span, key: string, value: any): SamplingDecision {
-    return { sample: false, retryable: true, tags: null };
+  isSampled(operationName: string, outTags: {}) {
+    return this._delegate.isSampled(operationName, outTags);
   }
 
   toString(): string {

src/samplers/const_sampler.js

@@ -11,12 +11,14 @@
 // or implied. See the License for the specific language governing permissions and limitations under
 // the License.
 
-import * as constants from '../constants.js';
+import * as constants from '../constants';
+import LegacySamplerV1Base from './_adapt_sampler';
 
-export default class ConstSampler implements LegacySamplerV1 {
+export default class ConstSampler extends LegacySamplerV1Base implements LegacySamplerV1 {
   _decision: boolean;
 
   constructor(decision: boolean) {
+    super('ConstSampler');
     this._decision = decision;
   }
 
@@ -47,10 +49,4 @@ export default class ConstSampler implements LegacySamplerV1 {
 
     return this.decision === other.decision;
   }
-
-  close(callback: ?Function): void {
-    if (callback) {
-      callback();
-    }
-  }
 }

src/samplers/experimental/priority_sampler.js

@@ -0,0 +1,103 @@
+// @flow
+// Copyright (c) 2019 Uber Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under
+// the License.
+
+import { adaptSamplerOrThrow } from '../_adapt_sampler';
+import BaseSamplerV2 from '../v2/base';
+import Span from '../../span';
+import Utils from '../../util';
+
+/**
+ * PrioritySamplerState keeps the state of all underlying samplers, specifically
+ * whether each of them has previously returned retryable=false, in which case
+ * those samplers are no longer invoked on future sampling calls.
+ */
+export class PrioritySamplerState {
+  samplerFired: Array<boolean>;
+
+  constructor(numDelegateSamplers: number) {
+    this.samplerFired = Array(numDelegateSamplers);
+    // TODO: for some reason Babel does not translate array.fill() to polyfil that works w/ Node 0.10
+    for (let i = 0; i < numDelegateSamplers; i++) {
+      this.samplerFired[i] = false;
+    }
+  }
+}
+
+/**
+ * PrioritySampler contains a list of samplers that it interrogates in order.
+ * Sampling methods return as soon as one of the samplers returns sample=true.
+ * The retryable state for each underlying sampler is stored in the extended context
+ * and once retryable=false is returned by one of the delegates it will never be
+ * called against.
+ */
+export default class PrioritySampler extends BaseSamplerV2 {
+  _delegates: Array<Sampler>;
+
+  constructor(samplers: Array<Sampler | LegacySamplerV1>) {
+    super('PrioritySampler');
+    this._delegates = samplers.map(s => adaptSamplerOrThrow(s));
+  }
+
+  _getOrCreateState(span: Span): PrioritySamplerState {
+    const store = span.context()._samplingState.extendedState();
+    const stateKey = this.uniqueName();
+    let state: ?PrioritySamplerState = store[stateKey];
+    if (!state) {
+      state = new PrioritySamplerState(this._delegates.length);
+      store[stateKey] = state;
+    }
+    return state;
+  }
+
+  _trySampling(span: Span, fn: Function): SamplingDecision {
+    const state = this._getOrCreateState(span);
+    let retryable = false;
+    for (let i = 0; i < this._delegates.length; i++) {
+      if (state.samplerFired[i]) {
+        continue;
+      }
+      const d = fn(this._delegates[i]);
+      retryable = retryable || d.retryable;
+      if (!d.retryable) {
+        state.samplerFired[i] = true;
+      }
+      if (d.sample) {
+        return d; // TODO do we want to alter out tags?
+      }
+    }
+    return { sample: false, retryable: retryable, tags: null };
+  }
+
+  onCreateSpan(span: Span): SamplingDecision {
+    return this._trySampling(span, function(delegate: Sampler): SamplingDecision {
+      return delegate.onCreateSpan(span);
+    });
+  }
+
+  onSetOperationName(span: Span, operationName: string): SamplingDecision {
+    return this._trySampling(span, function(delegate: Sampler): SamplingDecision {
+      return delegate.onSetOperationName(span, operationName);
+    });
+  }
+
+  onSetTag(span: Span, key: string, value: any): SamplingDecision {
+    return this._trySampling(span, function(delegate: Sampler): SamplingDecision {
+      return delegate.onSetTag(span, key, value);
+    });
+  }
+
+  close(callback: ?() => void): void {
+    const countdownCallback = Utils.countdownCallback(this._delegates.length, callback);
+    this._delegates.forEach(r => r.close(countdownCallback));
+  }
+}

src/samplers/experimental/tag_equals_sampler.js

@@ -0,0 +1,101 @@
+// @flow
+// Copyright (c) 2019 Uber Technologies, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
+// in compliance with the License. You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software distributed under the License
+// is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
+// or implied. See the License for the specific language governing permissions and limitations under
+// the License.
+
+import { adaptSamplerOrThrow } from '../_adapt_sampler';
+import BaseSamplerV2 from '../v2/base';
+import Span from '../../span';
+
+declare type Matcher = {
+  tagValue: string,
+  firehose: boolean,
+};
+
+export default class TagEqualsSampler extends BaseSamplerV2 {
+  _tagKey: string;
+  _matchers: { [string]: Matcher };
+  _undecided: SamplingDecision;
+
+  constructor(tagKey: string, matchers: Array<Matcher>) {
+    super('TagEqualsSampler');
+    this._tagKey = tagKey;
+    this._matchers = {};
+    matchers.forEach(m => {
+      this._matchers[m.tagValue] = m;
+    });
+    this._undecided = { sample: false, retryable: true, tags: null };
+  }
+
+  /**
+   * Creates the sampler from a JSON configuration of the following form:
+   * <code>
+   *   {
+   *     key: 'taKey',
+   *     values: {
+   *       'tagValue1': {
+   *         firehose: true,
+   *       },
+   *       'tagValue1: {
+   *         firehose: false,
+   *       },
+   *     },
+   *   }
+   * </code>
+   * @param {JSON} strategy
+   */
+  static fromStrategy(strategy: any): TagEqualsSampler {
+    let key = strategy.key;
+    let matchers: Array<Matcher> = [];
+    Object.keys(strategy.values).forEach(v => {
+      matchers.push({
+        tagValue: v,
+        firehose: Boolean(strategy.values[v].firehose),
+      });
+    });
+    return new TagEqualsSampler(key, matchers);
+  }
+
+  _createOutTags(tagValue: string): { [string]: string } {
+    return {
+      'sampler.type': 'TagEqualsSampler',
+      'sampler.param': tagValue,
+    };
+  }
+
+  _decide(span: Span, tagValue: any): SamplingDecision {
+    const match: ?Matcher = this._matchers[tagValue];
+    if (match) {
+      if (match.firehose) {
+        span._spanContext._setFirehose(true);
+      }
+      return { sample: true, retryable: false, tags: this._createOutTags(match.tagValue) };
+    }
+    return this._undecided;
+  }
+
+  onCreateSpan(span: Span): SamplingDecision {
+    // onCreateSpan is called on a brand new span that has no tags yet, so nothing to do here.
+    return this._undecided;
+  }
+
+  onSetOperationName(span: Span, operationName: string): SamplingDecision {
+    // this sampler is not sensitive to operationName, so nothing to do here.
+    return this._undecided;
+  }
+
+  onSetTag(span: Span, key: string, value: any): SamplingDecision {
+    if (key === this._tagKey) {
+      return this._decide(span, value);
+    }
+    return this._undecided;
+  }
+}

src/samplers/guaranteed_throughput_sampler.js

@@ -11,9 +11,9 @@
 // or implied. See the License for the specific language governing permissions and limitations under
 // the License.
 
-import * as constants from '../constants.js';
-import ProbabilisticSampler from './probabilistic_sampler.js';
-import RateLimitingSampler from './ratelimiting_sampler.js';
+import * as constants from '../constants';
+import ProbabilisticSampler from './probabilistic_sampler';
+import RateLimitingSampler from './rate_limiting_sampler';
 
 // GuaranteedThroughputProbabilisticSampler is a sampler that leverages both probabilisticSampler and
 // rateLimitingSampler. The rateLimitingSampler is used as a guaranteed lower bound sampler such that