github.com/observiq/carbon@v0.9.11-0.20200820160507-1b872e368a5e/operator/builtin/transformer/rate_limit.go (about)

     1  package transformer
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"time"
     7  
     8  	"github.com/observiq/carbon/entry"
     9  	"github.com/observiq/carbon/operator"
    10  	"github.com/observiq/carbon/operator/helper"
    11  )
    12  
    13  func init() {
    14  	operator.Register("rate_limit", func() operator.Builder { return NewRateLimitConfig("") })
    15  }
    16  
    17  func NewRateLimitConfig(operatorID string) *RateLimitConfig {
    18  	return &RateLimitConfig{
    19  		TransformerConfig: helper.NewTransformerConfig(operatorID, "rate_limit"),
    20  	}
    21  }
    22  
    23  // RateLimitConfig is the configuration of a rate filter operator.
    24  type RateLimitConfig struct {
    25  	helper.TransformerConfig `yaml:",inline"`
    26  
    27  	Rate     float64           `json:"rate,omitempty"     yaml:"rate,omitempty"`
    28  	Interval operator.Duration `json:"interval,omitempty" yaml:"interval,omitempty"`
    29  	Burst    uint              `json:"burst,omitempty"    yaml:"burst,omitempty"`
    30  }
    31  
    32  // Build will build a rate limit operator.
    33  func (c RateLimitConfig) Build(context operator.BuildContext) (operator.Operator, error) {
    34  	transformerOperator, err := c.TransformerConfig.Build(context)
    35  	if err != nil {
    36  		return nil, err
    37  	}
    38  
    39  	var interval time.Duration
    40  	switch {
    41  	case c.Rate != 0 && c.Interval.Raw() != 0:
    42  		return nil, fmt.Errorf("only one of 'rate' or 'interval' can be defined")
    43  	case c.Rate < 0 || c.Interval.Raw() < 0:
    44  		return nil, fmt.Errorf("rate and interval must be greater than zero")
    45  	case c.Rate > 0:
    46  		interval = time.Second / time.Duration(c.Rate)
    47  	default:
    48  		interval = c.Interval.Raw()
    49  	}
    50  
    51  	rateLimitOperator := &RateLimitOperator{
    52  		TransformerOperator: transformerOperator,
    53  		interval:            interval,
    54  		burst:               c.Burst,
    55  	}
    56  
    57  	return rateLimitOperator, nil
    58  }
    59  
    60  // RateLimitOperator is an operator that limits the rate of log consumption between operators.
    61  type RateLimitOperator struct {
    62  	helper.TransformerOperator
    63  
    64  	interval time.Duration
    65  	burst    uint
    66  	isReady  chan struct{}
    67  	cancel   context.CancelFunc
    68  }
    69  
    70  // Process will wait until a rate is met before sending an entry to the output.
    71  func (p *RateLimitOperator) Process(ctx context.Context, entry *entry.Entry) error {
    72  	<-p.isReady
    73  	p.Write(ctx, entry)
    74  	return nil
    75  }
    76  
    77  // Start will start the rate limit operator.
    78  func (p *RateLimitOperator) Start() error {
    79  	p.isReady = make(chan struct{}, p.burst)
    80  	ticker := time.NewTicker(p.interval)
    81  
    82  	ctx, cancel := context.WithCancel(context.Background())
    83  	p.cancel = cancel
    84  
    85  	// Buffer the ticker ticks in isReady to allow bursts
    86  	go func() {
    87  		defer ticker.Stop()
    88  		defer close(p.isReady)
    89  		for {
    90  			select {
    91  			case <-ticker.C:
    92  				p.isReady <- struct{}{}
    93  			case <-ctx.Done():
    94  				return
    95  			}
    96  		}
    97  	}()
    98  
    99  	return nil
   100  }
   101  
   102  // Stop will stop the rate limit operator.
   103  func (p *RateLimitOperator) Stop() error {
   104  	p.cancel()
   105  	return nil
   106  }