feat(opentelemetry-sampler-aws-xray): add rule cache and rule matching

packages/opentelemetry-sampler-aws-xray/src/fallback-sampler.ts

@@ -0,0 +1,42 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * 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
+ *
+ *      https://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 {
+  Sampler,
+  SamplingResult,
+  AlwaysOnSampler,
+} from '@opentelemetry/sdk-trace-base';
+import { Context, Link, Attributes, SpanKind } from '@opentelemetry/api';
+
+export class FallbackSampler implements Sampler {
+  private _alwaysOnSampler = new AlwaysOnSampler();
+
+  shouldSample(
+    context: Context,
+    traceId: string,
+    spanName: string,
+    spanKind: SpanKind,
+    attributes: Attributes,
+    links: Link[]
+  ): SamplingResult {
+    // This will be updated to be a rate limiting sampler in the next PR
+    return this._alwaysOnSampler.shouldSample();
+  }
+
+  public toString = (): string => {
+    return 'FallbackSampler';
+  };
+}

packages/opentelemetry-sampler-aws-xray/src/remote-sampler.ts

@@ -16,16 +16,22 @@
 
 import {
   Sampler,
-  SamplingDecision,
   SamplingResult,
 } from '@opentelemetry/sdk-trace-base';
-import { diag, DiagLogger } from '@opentelemetry/api';
 import {
-  SamplingRule,
-  AWSXRaySamplerConfig,
-  SamplingRuleRecord,
-} from './types';
+  diag,
+  DiagLogger,
+  Context,
+  Link,
+  Attributes,
+  SpanKind,
+} from '@opentelemetry/api';
+import { SamplingRule } from './sampling-rule';
+import { Resource } from '@opentelemetry/resources';
+import { AWSXRaySamplerConfig, SamplingRuleRecord } from './types';
 import axios from 'axios';
+import { RuleCache } from './rule-cache';
+import { FallbackSampler } from './fallback-sampler';
 
 // 5 minute interval on sampling rules fetch (default polling interval)
 const DEFAULT_INTERVAL_MS = 5 * 60 * 1000;
@@ -37,6 +43,9 @@ const SAMPLING_RULES_PATH = '/GetSamplingRules';
 export class AWSXRayRemoteSampler implements Sampler {
   private _pollingInterval: number;
   private _awsProxyEndpoint: string;
+  private _resource: Resource;
+  private _ruleCache: RuleCache;
+  private _fallBackSampler: FallbackSampler;
   private _samplerDiag: DiagLogger;
 
   constructor(samplerConfig: AWSXRaySamplerConfig) {
@@ -45,6 +54,9 @@ export class AWSXRayRemoteSampler implements Sampler {
     this._awsProxyEndpoint = samplerConfig.endpoint
       ? samplerConfig.endpoint
       : DEFAULT_AWS_PROXY_ENDPOINT;
+    this._resource = samplerConfig.resource;
+    this._ruleCache = new RuleCache();
+    this._fallBackSampler = new FallbackSampler();
 
     if (this._pollingInterval <= 0) {
       throw new TypeError('pollingInterval must be a positive integer');
@@ -58,9 +70,46 @@ export class AWSXRayRemoteSampler implements Sampler {
     this.startRulePoller();
   }
 
-  shouldSample(): SamplingResult {
-    // Implementation to be added
-    return { decision: SamplingDecision.NOT_RECORD };
+  shouldSample(
+    context: Context,
+    traceId: string,
+    spanName: string,
+    spanKind: SpanKind,
+    attributes: Attributes,
+    links: Link[]
+  ): SamplingResult {
+    if (this._ruleCache.isExpired()) {
+      // go to fallback sampler .shouldSample
+      return this._fallBackSampler.shouldSample(
+        context,
+        traceId,
+        spanName,
+        spanKind,
+        attributes,
+        links
+      );
+    }
+
+    const matchedRule = this._ruleCache.getMatchedRule(
+      attributes,
+      this._resource
+    );
+
+    // TODO: update after verifying if default rule will always match,
+    // this means that this method will always return return { decision: matchedRule.sample(attributes) }
+    // as long as the rule cache has not expired.
+    if (matchedRule) {
+      return { decision: matchedRule.sample(attributes) };
+    }
+
+    return this._fallBackSampler.shouldSample(
+      context,
+      traceId,
+      spanName,
+      spanKind,
+      attributes,
+      links
+    );
   }
 
   toString(): string {
@@ -88,11 +137,10 @@ export class AWSXRayRemoteSampler implements Sampler {
 
       samplingRules =
         responseJson?.SamplingRuleRecords.map(
-          (record: SamplingRuleRecord) => record.SamplingRule
+          (record: SamplingRuleRecord) => new SamplingRule(record.SamplingRule)
         ).filter(Boolean) ?? [];
 
-      // TODO: pass samplingRules to rule cache, temporarily logging the samplingRules array
-      this._samplerDiag.debug('sampling rules: ', samplingRules);
+      this._ruleCache.updateRules(samplingRules);
     } catch (error) {
       // Log error
       this._samplerDiag.warn('Error fetching sampling rules: ', error);

packages/opentelemetry-sampler-aws-xray/src/reservoir.ts

@@ -0,0 +1,24 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * 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
+ *
+ *      https://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.
+ */
+
+export class Reservoir {
+  // This class will be implemented in the next PR
+  private _quota: number;
+
+  constructor(quota: number) {
+    this._quota = quota;
+  }
+}

packages/opentelemetry-sampler-aws-xray/src/rule-cache.ts

@@ -0,0 +1,91 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * 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
+ *
+ *      https://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 { SamplingRule } from './sampling-rule';
+import { Attributes } from '@opentelemetry/api';
+import { Resource } from '@opentelemetry/resources';
+
+const RuleCacheTTL_SECONDS = 60 * 60; // The cache expires 1 hour after the last refresh time. (in seconds)
+
+export class RuleCache {
+  public rules: SamplingRule[];
+  private lastUpdated: number | undefined = undefined;
+
+  constructor() {
+    this.rules = [];
+  }
+
+  public isExpired = (): boolean => {
+    if (this.lastUpdated === undefined) {
+      return true;
+    }
+
+    const now = Math.floor(new Date().getTime() / 1000); // getTime returns milliseconds -> /1000 gives seconds
+    return now > this.lastUpdated + RuleCacheTTL_SECONDS;
+  };
+
+  public getMatchedRule = (
+    attributes: Attributes,
+    resource: Resource
+  ): SamplingRule | undefined => {
+    if (this.isExpired()) {
+      return undefined;
+    }
+
+    return this.rules.find(
+      rule => rule.matches(attributes, resource) || rule.RuleName === 'Default'
+    );
+  };
+
+  private _sortRulesByPriority = () => {
+    this.rules.sort((rule1: SamplingRule, rule2: SamplingRule): number => {
+      const value = rule1.Priority - rule2.Priority;
+
+      if (value !== 0) return value;
+
+      if (rule1.RuleName > rule2.RuleName) {
+        return 1;
+      } else return -1;
+    });
+  };
+
+  public updateRules = (newRules: SamplingRule[]): void => {
+    // store previous rules
+    const oldRules: { [key: string]: SamplingRule } = {};
+
+    this.rules.forEach((rule: SamplingRule) => {
+      // use Object.assign to create a new copy of the rules object to store by value
+      oldRules[rule.RuleName] = Object.assign({}, rule);
+    });
+
+    // update rules in the cache
+    this.rules = newRules;
+
+    this.rules.forEach((rule: SamplingRule) => {
+      const oldRule: SamplingRule = oldRules[rule.RuleName];
+
+      if (oldRule) {
+        rule.reservoir = oldRule.reservoir;
+        rule.statistics = oldRule.statistics;
+      }
+    });
+
+    // sort rules by priority and make lastUpdated = now
+    this._sortRulesByPriority();
+    this.lastUpdated = Math.floor(new Date().getTime() / 1000);
+
+    return;
+  };
+}

packages/opentelemetry-sampler-aws-xray/src/sampling-rule.ts

@@ -0,0 +1,92 @@
+/*
+ * Copyright The OpenTelemetry Authors
+ *
+ * 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
+ *
+ *      https://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 { ISamplingRule } from './types';
+import { wildcardMatch, attributeMatch } from './utils';
+import { SamplingDecision } from '@opentelemetry/sdk-trace-base';
+import { Attributes, AttributeValue } from '@opentelemetry/api';
+import {
+  SemanticAttributes,
+  SemanticResourceAttributes,
+} from '@opentelemetry/semantic-conventions';
+import { Resource } from '@opentelemetry/resources';
+import { Reservoir } from './reservoir';
+import { Statistics } from './statistics';
+
+export class SamplingRule implements ISamplingRule {
+  public RuleName: string;
+  public Priority: number;
+  public ReservoirSize: number;
+  public FixedRate: number;
+  public ServiceName: string;
+  public ServiceType: string;
+  public Host: string;
+  public HTTPMethod: string;
+  public URLPath: string;
+  public ResourceARN: string;
+  public Attributes?: { [key: string]: string };
+  public Version: number;
+  public reservoir: Reservoir;
+  public statistics: Statistics;
+
+  constructor(samplingRule: any) {
+    this.RuleName = samplingRule.ruleName;
+    this.Priority = samplingRule.priority;
+    this.ReservoirSize = samplingRule.reservoirSize;
+    this.FixedRate = samplingRule.fixedRate;
+    this.ServiceName = samplingRule.serviceName;
+    this.ServiceType = samplingRule.serviceType;
+    this.Host = samplingRule.host;
+    this.HTTPMethod = samplingRule.httpMethod;
+    this.URLPath = samplingRule.urlPath;
+    this.ResourceARN = samplingRule.resourceARN;
+    this.Version = samplingRule.version;
+    this.Attributes = samplingRule.attributes;
+
+    this.reservoir = new Reservoir(this.ReservoirSize);
+    this.statistics = new Statistics();
+  }
+
+  public matches = (attributes: Attributes, resource: Resource): boolean => {
+    const host: AttributeValue | undefined =
+      attributes[SemanticAttributes.HTTP_HOST];
+    const httpMethod = attributes[SemanticAttributes.HTTP_METHOD];
+    const serviceName = attributes[SemanticResourceAttributes.SERVICE_NAME];
+    const urlPath = attributes[SemanticAttributes.HTTP_URL];
+    const serviceType =
+      resource.attributes[SemanticResourceAttributes.CLOUD_PLATFORM];
+
+    // "Default" is a reserved keyword from X-Ray back-end.
+    if (this.RuleName === 'Default') {
+      return true;
+    }
+
+    return (
+      attributeMatch(attributes, this.Attributes) &&
+      (!host || wildcardMatch(this.Host, host)) &&
+      (!httpMethod || wildcardMatch(this.HTTPMethod, httpMethod)) &&
+      (!serviceName || wildcardMatch(this.ServiceName, serviceName)) &&
+      (!urlPath || wildcardMatch(this.URLPath, urlPath)) &&
+      (!serviceType || wildcardMatch(this.ServiceType, serviceType))
+    );
+  };
+
+  public sample = (attributes: Attributes): SamplingDecision => {
+    // Returning a drop sample decision
+    // TODO: use reservoir in next PR to make this decision
+    return SamplingDecision.NOT_RECORD;
+  };
+}