github.com/Jeffail/benthos/v3@v3.65.0/internal/component/ratelimit/rate_limit_v2.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/Jeffail/benthos/v3/internal/shutdown"
     8  	"github.com/Jeffail/benthos/v3/lib/metrics"
     9  	"github.com/Jeffail/benthos/v3/lib/types"
    10  )
    11  
    12  // V2 is a simpler interface to implement than types.RateLimit.
    13  type V2 interface {
    14  	// Access the rate limited resource. Returns a duration or an error if the
    15  	// rate limit check fails. The returned duration is either zero (meaning the
    16  	// resource can be accessed) or a reasonable length of time to wait before
    17  	// requesting again.
    18  	Access(ctx context.Context) (time.Duration, error)
    19  
    20  	// Close the component, blocks until either the underlying resources are
    21  	// cleaned up or the context is cancelled. Returns an error if the context
    22  	// is cancelled.
    23  	Close(ctx context.Context) error
    24  }
    25  
    26  //------------------------------------------------------------------------------
    27  
    28  // Implements types.RateLimit
    29  type v2ToV1RateLimit struct {
    30  	r   V2
    31  	sig *shutdown.Signaller
    32  
    33  	mChecked metrics.StatCounter
    34  	mLimited metrics.StatCounter
    35  	mErr     metrics.StatCounter
    36  }
    37  
    38  // NewV2ToV1RateLimit wraps a ratelimit.V2 with a struct that implements
    39  // types.RateLimit.
    40  func NewV2ToV1RateLimit(r V2, stats metrics.Type) types.RateLimit {
    41  	return &v2ToV1RateLimit{
    42  		r: r, sig: shutdown.NewSignaller(),
    43  
    44  		mChecked: stats.GetCounter("checked"),
    45  		mLimited: stats.GetCounter("limited"),
    46  		mErr:     stats.GetCounter("error"),
    47  	}
    48  }
    49  
    50  func (r *v2ToV1RateLimit) Access() (time.Duration, error) {
    51  	r.mChecked.Incr(1)
    52  	tout, err := r.r.Access(context.Background())
    53  	if err != nil {
    54  		r.mErr.Incr(1)
    55  	} else if tout > 0 {
    56  		r.mLimited.Incr(1)
    57  	}
    58  	return tout, err
    59  }
    60  
    61  func (r *v2ToV1RateLimit) CloseAsync() {
    62  	go func() {
    63  		if err := r.r.Close(context.Background()); err == nil {
    64  			r.sig.ShutdownComplete()
    65  		}
    66  	}()
    67  }
    68  
    69  func (r *v2ToV1RateLimit) WaitForClose(tout time.Duration) error {
    70  	select {
    71  	case <-r.sig.HasClosedChan():
    72  	case <-time.After(tout):
    73  		return types.ErrTimeout
    74  	}
    75  	return nil
    76  }