Line data    Source code

       1              : // Copyright The OpenTelemetry Authors
       2              : // SPDX-License-Identifier: Apache-2.0
       3              : 
       4              : import 'dart:async';
       5              : import 'package:dartastic_opentelemetry_api/dartastic_opentelemetry_api.dart';
       6              : import 'sampler.dart';
       7              : 
       8              : /// A sampler that limits the number of sampled traces per time window.
       9              : class RateLimitingSampler implements Sampler {
      10              :   final double _maxTracesPerSecond;
      11              :   final Duration _timeWindow;
      12              :   double _tokenBalance;
      13              :   DateTime _lastTokenUpdate;
      14              :   late final Timer _tokenReplenishTimer;
      15              : 
      16            1 :   @override
      17              :   String get description =>
      18            2 :       'RateLimitingSampler{$_maxTracesPerSecond per second}';
      19              : 
      20              :   /// Creates a rate limiting sampler.
      21              :   /// [maxTracesPerSecond] specifies how many traces can be sampled per second.
      22              :   /// [timeWindow] specifies how often the token balance is updated (defaults to 100ms).
      23            2 :   RateLimitingSampler(
      24              :     double maxTracesPerSecond, {
      25              :     Duration timeWindow = const Duration(milliseconds: 100),
      26              :   })  : _maxTracesPerSecond = maxTracesPerSecond,
      27              :         _timeWindow = timeWindow,
      28              :         _tokenBalance =
      29              :             maxTracesPerSecond, // Start with tokens already available
      30            2 :         _lastTokenUpdate = DateTime.now() {
      31            2 :     if (maxTracesPerSecond <= 0) {
      32            1 :       throw ArgumentError('maxTracesPerSecond must be positive');
      33              :     }
      34            2 :     _updateTokens();
      35            6 :     _tokenReplenishTimer = Timer.periodic(timeWindow, (_) => _updateTokens());
      36              :   }
      37              : 
      38            2 :   void _updateTokens() {
      39            2 :     final now = DateTime.now();
      40              :     final elapsedSeconds =
      41            8 :         now.difference(_lastTokenUpdate).inMilliseconds / 1000;
      42            2 :     _lastTokenUpdate = now;
      43              : 
      44              :     // Calculate how many tokens to add based on elapsed time and rate
      45              :     // Don't use floor() to ensure even small time periods add tokens
      46            4 :     final tokensToAdd = _maxTracesPerSecond * elapsedSeconds;
      47              : 
      48              :     // Calculate max tokens based on rate and time window
      49           10 :     final maxTokens = _maxTracesPerSecond * _timeWindow.inMilliseconds / 1000;
      50              : 
      51              :     // Update balance, ensuring we don't exceed max
      52            8 :     _tokenBalance = (_tokenBalance + tokensToAdd).clamp(0.0, maxTokens);
      53              :   }
      54              : 
      55            1 :   @override
      56              :   SamplingResult shouldSample({
      57              :     required Context parentContext,
      58              :     required String traceId,
      59              :     required String name,
      60              :     required SpanKind spanKind,
      61              :     required Attributes? attributes,
      62              :     required List<SpanLink>? links,
      63              :   }) {
      64              :     // Update tokens first
      65            1 :     _updateTokens();
      66              : 
      67            1 :     final parentTraceState = parentContext.spanContext?.traceState;
      68              : 
      69              :     // If we have tokens available, sample the trace
      70            2 :     if (_tokenBalance >= 1.0) {
      71            2 :       _tokenBalance -= 1.0;
      72            1 :       return SamplingResult(
      73              :         decision: SamplingDecision.recordAndSample,
      74              :         source: SamplingDecisionSource.tracerConfig,
      75              :         traceState: parentTraceState,
      76              :       );
      77              :     }
      78              : 
      79            1 :     return SamplingResult(
      80              :       decision: SamplingDecision.drop,
      81              :       source: SamplingDecisionSource.tracerConfig,
      82              :       traceState: parentTraceState,
      83              :     );
      84              :   }
      85              : 
      86              :   /// Clean up timer resources
      87            2 :   void dispose() {
      88            4 :     _tokenReplenishTimer.cancel();
      89              :   }
      90              : }