[XRay Sampler] Fix - Ensure all XRay Sampler functionality is under ParentBased logic (#80)

Author: jj22ee

aws-distro-opentelemetry-node-autoinstrumentation/src/sampler/aws-xray-remote-sampler.ts

@@ -22,12 +22,38 @@ const DEFAULT_RULES_POLLING_INTERVAL_SECONDS: number = 5 * 60;
 // Default endpoint for awsproxy : https://aws-otel.github.io/docs/getting-started/remote-sampling#enable-awsproxy-extension
 const DEFAULT_AWS_PROXY_ENDPOINT: string = 'http://localhost:2000';
 
+// Wrapper class to ensure that all XRay Sampler Functionality in _AwsXRayRemoteSampler
+// uses ParentBased logic to respect the parent span's sampling decision
 export class AwsXRayRemoteSampler implements Sampler {
+  private _root: ParentBasedSampler;
+  constructor(samplerConfig: AwsXRayRemoteSamplerConfig) {
+    this._root = new ParentBasedSampler({ root: new _AwsXRayRemoteSampler(samplerConfig) });
+  }
+  public shouldSample(
+    context: Context,
+    traceId: string,
+    spanName: string,
+    spanKind: SpanKind,
+    attributes: Attributes,
+    links: Link[]
+  ): SamplingResult {
+    return this._root.shouldSample(context, traceId, spanName, spanKind, attributes, links);
+  }
+
+  public toString(): string {
+    return `AwsXRayRemoteSampler{root=${this._root.toString()}`;
+  }
+}
+
+// _AwsXRayRemoteSampler contains all core XRay Sampler Functionality,
+// however it is NOT Parent-based (e.g. Sample logic runs for each span)
+// Not intended for external use, use Parent-based `AwsXRayRemoteSampler` instead.
+export class _AwsXRayRemoteSampler implements Sampler {
   private rulePollingIntervalMillis: number;
   private targetPollingInterval: number;
   private awsProxyEndpoint: string;
   private ruleCache: RuleCache;
-  private fallbackSampler: ParentBasedSampler;
+  private fallbackSampler: FallbackSampler;
   private samplerDiag: DiagLogger;
   private rulePoller: NodeJS.Timer | undefined;
   private targetPoller: NodeJS.Timer | undefined;
@@ -53,8 +79,8 @@ export class AwsXRayRemoteSampler implements Sampler {
     this.targetPollingJitterMillis = (Math.random() / 10) * 1000;
 
     this.awsProxyEndpoint = samplerConfig.endpoint ? samplerConfig.endpoint : DEFAULT_AWS_PROXY_ENDPOINT;
-    this.fallbackSampler = new ParentBasedSampler({ root: new FallbackSampler() });
-    this.clientId = AwsXRayRemoteSampler.generateClientId();
+    this.fallbackSampler = new FallbackSampler();
+    this.clientId = _AwsXRayRemoteSampler.generateClientId();
     this.ruleCache = new RuleCache(samplerConfig.resource);
 
     this.samplingClient = new AwsXraySamplingClient(this.awsProxyEndpoint, this.samplerDiag);
@@ -95,7 +121,9 @@ export class AwsXRayRemoteSampler implements Sampler {
   }
 
   public toString(): string {
-    return 'AwsXRayRemoteSampler{remote sampling with AWS X-Ray}';
+    return `_AwsXRayRemoteSampler{awsProxyEndpoint=${
+      this.awsProxyEndpoint
+    }, rulePollingIntervalMillis=${this.rulePollingIntervalMillis.toString()}}`;
   }
 
   private startSamplingRulesPoller(): void {

aws-distro-opentelemetry-node-autoinstrumentation/src/sampler/sampling-rule-applier.ts

@@ -3,12 +3,7 @@
 
 import { AttributeValue, Attributes, Context, Link, SpanKind } from '@opentelemetry/api';
 import { Resource } from '@opentelemetry/resources';
-import {
-  ParentBasedSampler,
-  SamplingDecision,
-  SamplingResult,
-  TraceIdRatioBasedSampler,
-} from '@opentelemetry/sdk-trace-base';
+import { SamplingDecision, SamplingResult, TraceIdRatioBasedSampler } from '@opentelemetry/sdk-trace-base';
 import {
   ATTR_CLIENT_ADDRESS,
   ATTR_HTTP_REQUEST_METHOD,
@@ -39,20 +34,20 @@ const MAX_DATE_TIME_MILLIS: number = new Date(8_640_000_000_000_000).getTime();
 
 export class SamplingRuleApplier {
   public samplingRule: SamplingRule;
-  private reservoirSampler: ParentBasedSampler;
-  private fixedRateSampler: ParentBasedSampler;
+  private reservoirSampler: RateLimitingSampler;
+  private fixedRateSampler: TraceIdRatioBasedSampler;
   private statistics: Statistics;
   private borrowingEnabled: boolean;
   private reservoirExpiryTimeInMillis: number;
 
   constructor(samplingRule: ISamplingRule, statistics: Statistics = new Statistics(), target?: SamplingTargetDocument) {
     this.samplingRule = new SamplingRule(samplingRule);
 
-    this.fixedRateSampler = new ParentBasedSampler({ root: new TraceIdRatioBasedSampler(this.samplingRule.FixedRate) });
+    this.fixedRateSampler = new TraceIdRatioBasedSampler(this.samplingRule.FixedRate);
     if (samplingRule.ReservoirSize > 0) {
-      this.reservoirSampler = new ParentBasedSampler({ root: new RateLimitingSampler(1) });
+      this.reservoirSampler = new RateLimitingSampler(1);
     } else {
-      this.reservoirSampler = new ParentBasedSampler({ root: new RateLimitingSampler(0) });
+      this.reservoirSampler = new RateLimitingSampler(0);
     }
 
     this.reservoirExpiryTimeInMillis = MAX_DATE_TIME_MILLIS;
@@ -63,7 +58,7 @@ export class SamplingRuleApplier {
     if (target) {
       this.borrowingEnabled = false;
       if (target.ReservoirQuota) {
-        this.reservoirSampler = new ParentBasedSampler({ root: new RateLimitingSampler(target.ReservoirQuota) });
+        this.reservoirSampler = new RateLimitingSampler(target.ReservoirQuota);
       }
 
       if (target.ReservoirQuotaTTL) {
@@ -73,7 +68,7 @@ export class SamplingRuleApplier {
       }
 
       if (target.FixedRate) {
-        this.fixedRateSampler = new ParentBasedSampler({ root: new TraceIdRatioBasedSampler(target.FixedRate) });
+        this.fixedRateSampler = new TraceIdRatioBasedSampler(target.FixedRate);
       }
     }
   }
@@ -156,7 +151,7 @@ export class SamplingRuleApplier {
     }
 
     if (result.decision === SamplingDecision.NOT_RECORD) {
-      result = this.fixedRateSampler.shouldSample(context, traceId, spanName, spanKind, attributes, links);
+      result = this.fixedRateSampler.shouldSample(context, traceId);
     }
 
     this.statistics.SampleCount += result.decision !== SamplingDecision.NOT_RECORD ? 1 : 0;

aws-distro-opentelemetry-node-autoinstrumentation/test/aws-opentelemetry-configurator.test.ts

@@ -19,6 +19,7 @@ import expect from 'expect';
 import * as sinon from 'sinon';
 import { AlwaysRecordSampler } from '../src/always-record-sampler';
 import { AttributePropagatingSpanProcessor } from '../src/attribute-propagating-span-processor';
+import { AwsBatchUnsampledSpanProcessor } from '../src/aws-batch-unsampled-span-processor';
 import { AwsMetricAttributesSpanExporter } from '../src/aws-metric-attributes-span-exporter';
 import {
   ApplicationSignalsExporterProvider,
@@ -27,12 +28,11 @@ import {
   customBuildSamplerFromEnv,
 } from '../src/aws-opentelemetry-configurator';
 import { AwsSpanMetricsProcessor } from '../src/aws-span-metrics-processor';
+import { OTLPUdpSpanExporter } from '../src/otlp-udp-exporter';
 import { setAwsDefaultEnvironmentVariables } from '../src/register';
 import { AwsXRayRemoteSampler } from '../src/sampler/aws-xray-remote-sampler';
 import { AwsXraySamplingClient } from '../src/sampler/aws-xray-sampling-client';
 import { GetSamplingRulesResponse } from '../src/sampler/remote-sampler.types';
-import { OTLPUdpSpanExporter } from '../src/otlp-udp-exporter';
-import { AwsBatchUnsampledSpanProcessor } from '../src/aws-batch-unsampled-span-processor';
 
 // Tests AwsOpenTelemetryConfigurator after running Environment Variable setup in register.ts
 describe('AwsOpenTelemetryConfiguratorTest', () => {
@@ -73,6 +73,7 @@ describe('AwsOpenTelemetryConfiguratorTest', () => {
       const startTimeSec: number = Math.floor(new Date().getTime() / 1000.0);
       const span: Span = tracer.startSpan('test');
       const traceId: string = span.spanContext().traceId;
+      span.end();
       const traceId4ByteHex: string = traceId.substring(0, 8);
       const traceId4ByteNumber: number = Number(`0x${traceId4ByteHex}`);
       expect(traceId4ByteNumber).toBeGreaterThanOrEqual(startTimeSec);
@@ -118,11 +119,11 @@ describe('AwsOpenTelemetryConfiguratorTest', () => {
     const sampler = customBuildSamplerFromEnv(Resource.empty());
 
     expect(sampler).toBeInstanceOf(AwsXRayRemoteSampler);
-    expect((sampler as any).awsProxyEndpoint).toEqual('http://localhost:2000');
-    expect((sampler as any).rulePollingIntervalMillis).toEqual(300000); // ms
+    expect((sampler as any)._root._root.awsProxyEndpoint).toEqual('http://localhost:2000');
+    expect((sampler as any)._root._root.rulePollingIntervalMillis).toEqual(300000); // ms
 
-    clearInterval((sampler as any).rulePoller);
-    clearInterval((sampler as any).targetPoller);
+    clearInterval((sampler as any)._root._root.rulePoller);
+    clearInterval((sampler as any)._root._root.targetPoller);
   });
 
   it('ImportXRaySamplerWhenSamplerArgsSet', () => {
@@ -133,17 +134,17 @@ describe('AwsOpenTelemetryConfiguratorTest', () => {
     const sampler = customBuildSamplerFromEnv(Resource.empty());
 
     expect(sampler).toBeInstanceOf(AwsXRayRemoteSampler);
-    expect((sampler as any).awsProxyEndpoint).toEqual('http://asdfghjkl:2000');
-    expect((sampler as any).rulePollingIntervalMillis).toEqual(600000); // ms
-    expect(((sampler as any).samplingClient as any).getSamplingRulesEndpoint).toEqual(
+    expect((sampler as any)._root._root.awsProxyEndpoint).toEqual('http://asdfghjkl:2000');
+    expect((sampler as any)._root._root.rulePollingIntervalMillis).toEqual(600000); // ms
+    expect(((sampler as any)._root._root.samplingClient as any).getSamplingRulesEndpoint).toEqual(
       'http://asdfghjkl:2000/GetSamplingRules'
     );
-    expect(((sampler as any).samplingClient as any).samplingTargetsEndpoint).toEqual(
+    expect(((sampler as any)._root._root.samplingClient as any).samplingTargetsEndpoint).toEqual(
       'http://asdfghjkl:2000/SamplingTargets'
     );
 
-    clearInterval((sampler as any).rulePoller);
-    clearInterval((sampler as any).targetPoller);
+    clearInterval((sampler as any)._root._root.rulePoller);
+    clearInterval((sampler as any)._root._root.targetPoller);
   });
 
   it('ImportXRaySamplerWithInvalidPollingIntervalSet', () => {
@@ -156,17 +157,17 @@ describe('AwsOpenTelemetryConfiguratorTest', () => {
     const sampler = customBuildSamplerFromEnv(Resource.empty());
 
     expect(sampler).toBeInstanceOf(AwsXRayRemoteSampler);
-    expect((sampler as any).awsProxyEndpoint).toEqual('http://asdfghjkl:2000');
-    expect((sampler as any).rulePollingIntervalMillis).toEqual(300000); // default value
-    expect(((sampler as any).samplingClient as any).getSamplingRulesEndpoint).toEqual(
+    expect((sampler as any)._root._root.awsProxyEndpoint).toEqual('http://asdfghjkl:2000');
+    expect((sampler as any)._root._root.rulePollingIntervalMillis).toEqual(300000); // default value
+    expect(((sampler as any)._root._root.samplingClient as any).getSamplingRulesEndpoint).toEqual(
       'http://asdfghjkl:2000/GetSamplingRules'
     );
-    expect(((sampler as any).samplingClient as any).samplingTargetsEndpoint).toEqual(
+    expect(((sampler as any)._root._root.samplingClient as any).samplingTargetsEndpoint).toEqual(
       'http://asdfghjkl:2000/SamplingTargets'
     );
 
-    clearInterval((sampler as any).rulePoller);
-    clearInterval((sampler as any).targetPoller);
+    clearInterval((sampler as any)._root._root.rulePoller);
+    clearInterval((sampler as any)._root._root.targetPoller);
   });
 
   // test_import_xray_sampler_with_invalid_environment_arguments
@@ -188,12 +189,12 @@ describe('AwsOpenTelemetryConfiguratorTest', () => {
     let sampler = customBuildSamplerFromEnv(Resource.empty());
 
     expect(sampler).toBeInstanceOf(AwsXRayRemoteSampler);
-    expect((sampler as any).awsProxyEndpoint).toEqual('http://lo=cal=host=:2000');
-    expect((sampler as any).rulePollingIntervalMillis).toEqual(600000);
-    expect(((sampler as any).samplingClient as any).getSamplingRulesEndpoint).toEqual(
+    expect((sampler as any)._root._root.awsProxyEndpoint).toEqual('http://lo=cal=host=:2000');
+    expect((sampler as any)._root._root.rulePollingIntervalMillis).toEqual(600000);
+    expect(((sampler as any)._root._root.samplingClient as any).getSamplingRulesEndpoint).toEqual(
       'http://lo=cal=host=:2000/GetSamplingRules'
     );
-    expect(((sampler as any).samplingClient as any).samplingTargetsEndpoint).toEqual(
+    expect(((sampler as any)._root._root.samplingClient as any).samplingTargetsEndpoint).toEqual(
       'http://lo=cal=host=:2000/SamplingTargets'
     );
 
@@ -202,12 +203,12 @@ describe('AwsOpenTelemetryConfiguratorTest', () => {
     sampler = customBuildSamplerFromEnv(Resource.empty());
 
     expect(sampler).toBeInstanceOf(AwsXRayRemoteSampler);
-    expect((sampler as any).awsProxyEndpoint).toEqual('http://localhost:2000');
-    expect((sampler as any).rulePollingIntervalMillis).toEqual(550000);
-    expect(((sampler as any).samplingClient as any).getSamplingRulesEndpoint).toEqual(
+    expect((sampler as any)._root._root.awsProxyEndpoint).toEqual('http://localhost:2000');
+    expect((sampler as any)._root._root.rulePollingIntervalMillis).toEqual(550000);
+    expect(((sampler as any)._root._root.samplingClient as any).getSamplingRulesEndpoint).toEqual(
       'http://localhost:2000/GetSamplingRules'
     );
-    expect(((sampler as any).samplingClient as any).samplingTargetsEndpoint).toEqual(
+    expect(((sampler as any)._root._root.samplingClient as any).samplingTargetsEndpoint).toEqual(
       'http://localhost:2000/SamplingTargets'
     );