feat(opentelemetry-sampler-aws-xray): Add Rate Limiter and Sampling Targets Poller Logic

Author: jj22ee

incubator/opentelemetry-sampler-aws-xray/src/aws-xray-sampling-client.ts

@@ -18,7 +18,7 @@
 // Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
 // SPDX-License-Identifier: Apache-2.0
 
-import { DiagLogFunction, DiagLogger, context } from '@opentelemetry/api';
+import { DiagLogFunction, context, diag } from '@opentelemetry/api';
 import { suppressTracing } from '@opentelemetry/core';
 import * as http from 'http';
 import {
@@ -30,12 +30,10 @@ import {
 export class AWSXRaySamplingClient {
   private getSamplingRulesEndpoint: string;
   private samplingTargetsEndpoint: string;
-  private samplerDiag: DiagLogger;
 
-  constructor(endpoint: string, samplerDiag: DiagLogger) {
+  constructor(endpoint: string) {
     this.getSamplingRulesEndpoint = endpoint + '/GetSamplingRules';
     this.samplingTargetsEndpoint = endpoint + '/SamplingTargets';
-    this.samplerDiag = samplerDiag;
   }
 
   public fetchSamplingTargets(
@@ -45,7 +43,7 @@ export class AWSXRaySamplingClient {
     this.makeSamplingRequest<GetSamplingTargetsResponse>(
       this.samplingTargetsEndpoint,
       callback,
-      this.samplerDiag.debug,
+      diag.debug,
       JSON.stringify(requestBody)
     );
   }
@@ -56,7 +54,7 @@ export class AWSXRaySamplingClient {
     this.makeSamplingRequest<GetSamplingRulesResponse>(
       this.getSamplingRulesEndpoint,
       callback,
-      this.samplerDiag.error
+      diag.error
     );
   }
 
@@ -98,10 +96,8 @@ export class AWSXRaySamplingClient {
                 callback(responseObject);
               }
             } else {
-              this.samplerDiag.debug(
-                `${url} Response Code is: ${response.statusCode}`
-              );
-              this.samplerDiag.debug(`${url} responseData is: ${responseData}`);
+              diag.debug(`${url} Response Code is: ${response.statusCode}`);
+              diag.debug(`${url} responseData is: ${responseData}`);
             }
           });
         })

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

@@ -21,16 +21,20 @@
 import { Attributes, Context, Link, SpanKind } from '@opentelemetry/api';
 import {
   Sampler,
+  SamplingDecision,
   SamplingResult,
   TraceIdRatioBasedSampler,
 } from '@opentelemetry/sdk-trace-base';
+import { RateLimitingSampler } from './rate-limiting-sampler';
 
 // FallbackSampler samples 1 req/sec and additional 5% of requests using TraceIdRatioBasedSampler.
 export class FallbackSampler implements Sampler {
   private fixedRateSampler: TraceIdRatioBasedSampler;
+  private rateLimitingSampler: RateLimitingSampler;
 
   constructor() {
     this.fixedRateSampler = new TraceIdRatioBasedSampler(0.05);
+    this.rateLimitingSampler = new RateLimitingSampler(1);
   }
 
   shouldSample(
@@ -41,7 +45,19 @@ export class FallbackSampler implements Sampler {
     attributes: Attributes,
     links: Link[]
   ): SamplingResult {
-    // TODO: implement and use Rate Limiting Sampler
+    const samplingResult: SamplingResult =
+      this.rateLimitingSampler.shouldSample(
+        context,
+        traceId,
+        spanName,
+        spanKind,
+        attributes,
+        links
+      );
+
+    if (samplingResult.decision !== SamplingDecision.NOT_RECORD) {
+      return samplingResult;
+    }
 
     return this.fixedRateSampler.shouldSample(context, traceId);
   }

incubator/opentelemetry-sampler-aws-xray/src/rate-limiter.ts

@@ -0,0 +1,67 @@
+/*
+ * 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
+
+/*
+ * The RateLimiter keeps track of the current reservoir quota balance available (measured via available time)
+ * If enough time has elapsed, the RateLimiter will allow quota balance to be consumed/taken (decrease available time)
+ * A RateLimitingSampler uses this RateLimiter to determine if it should sample or not based on the quota balance available.
+ */
+export class RateLimiter {
+  // Quota assigned to client to dictate maximum quota balance that can be consumed per second.
+  private quota: number;
+  private MAX_BALANCE_MILLIS: number;
+  // Used to measure current quota balance.
+  private walletFloorMillis: number;
+
+  constructor(quota: number, maxBalanceInSeconds = 1) {
+    this.MAX_BALANCE_MILLIS = maxBalanceInSeconds * 1000.0;
+    this.quota = quota;
+    this.walletFloorMillis = Date.now();
+    // current "balance" would be `ceiling - floor`
+  }
+
+  public take(cost = 1): boolean {
+    if (this.quota === 0) {
+      return false;
+    }
+
+    const quotaPerMillis: number = this.quota / 1000.0;
+
+    // assume divide by zero not possible
+    const costInMillis: number = cost / quotaPerMillis;
+
+    const walletCeilingMillis: number = Date.now();
+    let currentBalanceMillis: number =
+      walletCeilingMillis - this.walletFloorMillis;
+    currentBalanceMillis = Math.min(
+      currentBalanceMillis,
+      this.MAX_BALANCE_MILLIS
+    );
+    const pendingRemainingBalanceMillis: number =
+      currentBalanceMillis - costInMillis;
+    if (pendingRemainingBalanceMillis >= 0) {
+      this.walletFloorMillis =
+        walletCeilingMillis - pendingRemainingBalanceMillis;
+      return true;
+    }
+    // No changes to the wallet state
+    return false;
+  }
+}

incubator/opentelemetry-sampler-aws-xray/src/rate-limiting-sampler.ts

@@ -0,0 +1,58 @@
+/*
+ * 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,
+  SamplingDecision,
+  SamplingResult,
+} from '@opentelemetry/sdk-trace-base';
+import { RateLimiter } from './rate-limiter';
+
+export class RateLimitingSampler implements Sampler {
+  private quota: number;
+  private reservoir: RateLimiter;
+
+  constructor(quota: number) {
+    this.quota = quota;
+    this.reservoir = new RateLimiter(quota);
+  }
+
+  shouldSample(
+    context: Context,
+    traceId: string,
+    spanName: string,
+    spanKind: SpanKind,
+    attributes: Attributes,
+    links: Link[]
+  ): SamplingResult {
+    if (this.reservoir.take(1)) {
+      return {
+        decision: SamplingDecision.RECORD_AND_SAMPLED,
+        attributes: attributes,
+      };
+    }
+    return { decision: SamplingDecision.NOT_RECORD, attributes: attributes };
+  }
+
+  public toString(): string {
+    return `RateLimitingSampler{rate limiting sampling with sampling config of ${this.quota} req/sec and 0% of additional requests}`;
+  }
+}