github.com/Jeffail/benthos/v3@v3.65.0/lib/processor/rate_limit.go (about)

     1  package processor
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  
     8  	"github.com/Jeffail/benthos/v3/internal/docs"
     9  	"github.com/Jeffail/benthos/v3/internal/interop"
    10  	"github.com/Jeffail/benthos/v3/lib/log"
    11  	"github.com/Jeffail/benthos/v3/lib/metrics"
    12  	"github.com/Jeffail/benthos/v3/lib/types"
    13  )
    14  
    15  //------------------------------------------------------------------------------
    16  
    17  func init() {
    18  	Constructors[TypeRateLimit] = TypeSpec{
    19  		constructor: NewRateLimit,
    20  		Categories: []Category{
    21  			CategoryUtility,
    22  		},
    23  		Summary: `
    24  Throttles the throughput of a pipeline according to a specified
    25  ` + "[`rate_limit`](/docs/components/rate_limits/about)" + ` resource. Rate limits are
    26  shared across components and therefore apply globally to all processing
    27  pipelines.`,
    28  		FieldSpecs: docs.FieldSpecs{
    29  			docs.FieldCommon("resource", "The target [`rate_limit` resource](/docs/components/rate_limits/about)."),
    30  		},
    31  	}
    32  }
    33  
    34  //------------------------------------------------------------------------------
    35  
    36  // RateLimitConfig contains configuration fields for the RateLimit processor.
    37  type RateLimitConfig struct {
    38  	Resource string `json:"resource" yaml:"resource"`
    39  }
    40  
    41  // NewRateLimitConfig returns a RateLimitConfig with default values.
    42  func NewRateLimitConfig() RateLimitConfig {
    43  	return RateLimitConfig{
    44  		Resource: "",
    45  	}
    46  }
    47  
    48  //------------------------------------------------------------------------------
    49  
    50  // RateLimit is a processor that performs an RateLimit request using the message as the
    51  // request body, and returns the response.
    52  type RateLimit struct {
    53  	rlName string
    54  	mgr    types.Manager
    55  
    56  	log log.Modular
    57  
    58  	mCount       metrics.StatCounter
    59  	mRateLimited metrics.StatCounter
    60  	mErr         metrics.StatCounter
    61  	mSent        metrics.StatCounter
    62  	mBatchSent   metrics.StatCounter
    63  
    64  	closeChan chan struct{}
    65  	closeOnce sync.Once
    66  }
    67  
    68  // NewRateLimit returns a RateLimit processor.
    69  func NewRateLimit(
    70  	conf Config, mgr types.Manager, log log.Modular, stats metrics.Type,
    71  ) (Type, error) {
    72  	if err := interop.ProbeRateLimit(context.Background(), mgr, conf.RateLimit.Resource); err != nil {
    73  		return nil, err
    74  	}
    75  	r := &RateLimit{
    76  		rlName:       conf.RateLimit.Resource,
    77  		mgr:          mgr,
    78  		log:          log,
    79  		mCount:       stats.GetCounter("count"),
    80  		mRateLimited: stats.GetCounter("rate.limited"),
    81  		mErr:         stats.GetCounter("error"),
    82  		mSent:        stats.GetCounter("sent"),
    83  		mBatchSent:   stats.GetCounter("batch.sent"),
    84  		closeChan:    make(chan struct{}),
    85  	}
    86  	return r, nil
    87  }
    88  
    89  //------------------------------------------------------------------------------
    90  
    91  // ProcessMessage applies the processor to a message, either creating >0
    92  // resulting messages or a response to be sent back to the message source.
    93  func (r *RateLimit) ProcessMessage(msg types.Message) ([]types.Message, types.Response) {
    94  	r.mCount.Incr(1)
    95  
    96  	msg.Iter(func(i int, p types.Part) error {
    97  		var waitFor time.Duration
    98  		var err error
    99  		if rerr := interop.AccessRateLimit(context.Background(), r.mgr, r.rlName, func(rl types.RateLimit) {
   100  			waitFor, err = rl.Access()
   101  		}); rerr != nil {
   102  			err = rerr
   103  		}
   104  		for err != nil || waitFor > 0 {
   105  			if err == types.ErrTypeClosed {
   106  				return err
   107  			}
   108  			if err != nil {
   109  				r.mErr.Incr(1)
   110  				r.log.Errorf("Failed to access rate limit: %v\n", err)
   111  				waitFor = time.Second
   112  			} else {
   113  				r.mRateLimited.Incr(1)
   114  			}
   115  			select {
   116  			case <-time.After(waitFor):
   117  			case <-r.closeChan:
   118  				return types.ErrTypeClosed
   119  			}
   120  			if rerr := interop.AccessRateLimit(context.Background(), r.mgr, r.rlName, func(rl types.RateLimit) {
   121  				waitFor, err = rl.Access()
   122  			}); rerr != nil {
   123  				err = rerr
   124  			}
   125  		}
   126  		return err
   127  	})
   128  
   129  	r.mBatchSent.Incr(1)
   130  	r.mSent.Incr(int64(msg.Len()))
   131  	return []types.Message{msg}, nil
   132  }
   133  
   134  // CloseAsync shuts down the processor and stops processing requests.
   135  func (r *RateLimit) CloseAsync() {
   136  	r.closeOnce.Do(func() {
   137  		close(r.closeChan)
   138  	})
   139  }
   140  
   141  // WaitForClose blocks until the processor has closed down.
   142  func (r *RateLimit) WaitForClose(timeout time.Duration) error {
   143  	return nil
   144  }
   145  
   146  //------------------------------------------------------------------------------