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 //------------------------------------------------------------------------------