Merge branch 'main' into bedrock-invoke-model-stream

Author: yuliia-fryshko

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

@@ -0,0 +1,52 @@
+/*
+ * 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.
+ */
+
+// Includes work from:
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { Attributes, Context, Link, SpanKind } from '@opentelemetry/api';
+import {
+  Sampler,
+  SamplingResult,
+  TraceIdRatioBasedSampler,
+} from '@opentelemetry/sdk-trace-base';
+
+// FallbackSampler samples 1 req/sec and additional 5% of requests using TraceIdRatioBasedSampler.
+export class FallbackSampler implements Sampler {
+  private fixedRateSampler: TraceIdRatioBasedSampler;
+
+  constructor() {
+    this.fixedRateSampler = new TraceIdRatioBasedSampler(0.05);
+  }
+
+  shouldSample(
+    context: Context,
+    traceId: string,
+    spanName: string,
+    spanKind: SpanKind,
+    attributes: Attributes,
+    links: Link[]
+  ): SamplingResult {
+    // TODO: implement and use Rate Limiting Sampler
+
+    return this.fixedRateSampler.shouldSample(context, traceId);
+  }
+
+  public toString(): string {
+    return 'FallbackSampler{fallback sampling with sampling config of 1 req/sec and 5% of additional requests}';
+  }
+}

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

@@ -29,15 +29,17 @@ import {
 import {
   ParentBasedSampler,
   Sampler,
-  SamplingDecision,
   SamplingResult,
 } from '@opentelemetry/sdk-trace-base';
 import { AWSXRaySamplingClient } from './aws-xray-sampling-client';
+import { FallbackSampler } from './fallback-sampler';
 import {
   AWSXRayRemoteSamplerConfig,
   GetSamplingRulesResponse,
   SamplingRuleRecord,
 } from './types';
+import { RuleCache } from './rule-cache';
+
 import { SamplingRuleApplier } from './sampling-rule-applier';
 
 // 5 minute default sampling rules polling interval
@@ -50,12 +52,14 @@ const DEFAULT_AWS_PROXY_ENDPOINT = 'http://localhost:2000';
 export class AWSXRayRemoteSampler implements Sampler {
   private _root: ParentBasedSampler;
   private internalXraySampler: _AWSXRayRemoteSampler;
+
   constructor(samplerConfig: AWSXRayRemoteSamplerConfig) {
     this.internalXraySampler = new _AWSXRayRemoteSampler(samplerConfig);
     this._root = new ParentBasedSampler({
       root: this.internalXraySampler,
     });
   }
+
   public shouldSample(
     context: Context,
     traceId: string,
@@ -91,8 +95,11 @@ export class AWSXRayRemoteSampler implements Sampler {
 export class _AWSXRayRemoteSampler implements Sampler {
   private rulePollingIntervalMillis: number;
   private awsProxyEndpoint: string;
+  private ruleCache: RuleCache;
+  private fallbackSampler: FallbackSampler;
   private samplerDiag: DiagLogger;
   private rulePoller: NodeJS.Timeout | undefined;
+  private clientId: string;
   private rulePollingJitterMillis: number;
   private samplingClient: AWSXRaySamplingClient;
 
@@ -117,6 +124,10 @@ export class _AWSXRayRemoteSampler implements Sampler {
     this.awsProxyEndpoint = samplerConfig.endpoint
       ? samplerConfig.endpoint
       : DEFAULT_AWS_PROXY_ENDPOINT;
+    this.fallbackSampler = new FallbackSampler();
+    // TODO: Use clientId for retrieving Sampling Targets
+    this.clientId = _AWSXRayRemoteSampler.generateClientId();
+    this.ruleCache = new RuleCache(samplerConfig.resource);
 
     this.samplingClient = new AWSXRaySamplingClient(
       this.awsProxyEndpoint,
@@ -137,8 +148,51 @@ export class _AWSXRayRemoteSampler implements Sampler {
     attributes: Attributes,
     links: Link[]
   ): SamplingResult {
-    // Implementation to be added
-    return { decision: SamplingDecision.NOT_RECORD };
+    if (this.ruleCache.isExpired()) {
+      this.samplerDiag.debug(
+        'Rule cache is expired, so using fallback sampling strategy'
+      );
+      return this.fallbackSampler.shouldSample(
+        context,
+        traceId,
+        spanName,
+        spanKind,
+        attributes,
+        links
+      );
+    }
+
+    try {
+      const matchedRule: SamplingRuleApplier | undefined =
+        this.ruleCache.getMatchedRule(attributes);
+      if (matchedRule) {
+        return matchedRule.shouldSample(
+          context,
+          traceId,
+          spanName,
+          spanKind,
+          attributes,
+          links
+        );
+      }
+    } catch (e: unknown) {
+      this.samplerDiag.debug(
+        'Unexpected error occurred when trying to match or applying a sampling rule',
+        e
+      );
+    }
+
+    this.samplerDiag.debug(
+      'Using fallback sampler as no rule match was found. This is likely due to a bug, since default rule should always match'
+    );
+    return this.fallbackSampler.shouldSample(
+      context,
+      traceId,
+      spanName,
+      spanKind,
+      attributes,
+      links
+    );
   }
 
   public toString(): string {
@@ -180,13 +234,37 @@ export class _AWSXRayRemoteSampler implements Sampler {
           }
         }
       );
-
-      // TODO: pass samplingRules to rule cache, temporarily logging the samplingRules array
-      this.samplerDiag.debug('sampling rules: ', samplingRules);
+      this.ruleCache.updateRules(samplingRules);
     } else {
       this.samplerDiag.error(
         'SamplingRuleRecords from GetSamplingRules request is not defined'
       );
     }
   }
+
+  private static generateClientId(): string {
+    const hexChars: string[] = [
+      '0',
+      '1',
+      '2',
+      '3',
+      '4',
+      '5',
+      '6',
+      '7',
+      '8',
+      '9',
+      'a',
+      'b',
+      'c',
+      'd',
+      'e',
+      'f',
+    ];
+    const clientIdArray: string[] = [];
+    for (let _ = 0; _ < 24; _ += 1) {
+      clientIdArray.push(hexChars[Math.floor(Math.random() * hexChars.length)]);
+    }
+    return clientIdArray.join('');
+  }
 }

incubator/opentelemetry-sampler-aws-xray/src/rule-cache.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.
+ */
+
+// Includes work from:
+// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+// SPDX-License-Identifier: Apache-2.0
+
+import { Attributes } from '@opentelemetry/api';
+import { Resource } from '@opentelemetry/resources';
+import { SamplingRuleApplier } from './sampling-rule-applier';
+
+// The cache expires 1 hour after the last refresh time.
+const RULE_CACHE_TTL_MILLIS: number = 60 * 60 * 1000;
+
+export class RuleCache {
+  private ruleAppliers: SamplingRuleApplier[];
+  private lastUpdatedEpochMillis: number;
+  private samplerResource: Resource;
+
+  constructor(samplerResource: Resource) {
+    this.ruleAppliers = [];
+    this.samplerResource = samplerResource;
+    this.lastUpdatedEpochMillis = Date.now();
+  }
+
+  public isExpired(): boolean {
+    const nowInMillis: number = Date.now();
+    return nowInMillis > this.lastUpdatedEpochMillis + RULE_CACHE_TTL_MILLIS;
+  }
+
+  public getMatchedRule(
+    attributes: Attributes
+  ): SamplingRuleApplier | undefined {
+    // `this.ruleAppliers` should be sorted by priority, so `find()` is used here
+    // to determine the first highest priority rule that is matched. The last rule
+    // in the list should be the 'Default' rule with hardcoded priority of 10000.
+    return this.ruleAppliers.find(
+      rule =>
+        rule.matches(attributes, this.samplerResource) ||
+        rule.samplingRule.RuleName === 'Default'
+    );
+  }
+
+  private sortRulesByPriority(): void {
+    this.ruleAppliers.sort(
+      (rule1: SamplingRuleApplier, rule2: SamplingRuleApplier): number => {
+        if (rule1.samplingRule.Priority === rule2.samplingRule.Priority) {
+          return rule1.samplingRule.RuleName < rule2.samplingRule.RuleName
+            ? -1
+            : 1;
+        }
+        return rule1.samplingRule.Priority - rule2.samplingRule.Priority;
+      }
+    );
+  }
+
+  public updateRules(newRuleAppliers: SamplingRuleApplier[]): void {
+    const oldRuleAppliersMap = new Map<string, SamplingRuleApplier>();
+
+    this.ruleAppliers.forEach((rule: SamplingRuleApplier) => {
+      oldRuleAppliersMap.set(rule.samplingRule.RuleName, rule);
+    });
+
+    newRuleAppliers.forEach((newRule: SamplingRuleApplier, index: number) => {
+      const ruleNameToCheck: string = newRule.samplingRule.RuleName;
+      const oldRule = oldRuleAppliersMap.get(ruleNameToCheck);
+      if (oldRule) {
+        if (newRule.samplingRule.equals(oldRule.samplingRule)) {
+          newRuleAppliers[index] = oldRule;
+        }
+      }
+    });
+    this.ruleAppliers = newRuleAppliers;
+
+    // sort ruleAppliers by priority and update lastUpdatedEpochMillis
+    this.sortRulesByPriority();
+    this.lastUpdatedEpochMillis = Date.now();
+  }
+}