github.com/Jeffail/benthos/v3@v3.65.0/public/service/rate_limit.go (about)

     1  package service
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/Jeffail/benthos/v3/internal/component/ratelimit"
     8  	"github.com/Jeffail/benthos/v3/lib/metrics"
     9  	"github.com/Jeffail/benthos/v3/lib/types"
    10  )
    11  
    12  // RateLimit is an interface implemented by Benthos rate limits.
    13  type RateLimit 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 may be accessed) or a reasonable length of time to wait before
    17  	// requesting again.
    18  	Access(context.Context) (time.Duration, error)
    19  
    20  	Closer
    21  }
    22  
    23  //------------------------------------------------------------------------------
    24  
    25  func newAirGapRateLimit(c RateLimit, stats metrics.Type) types.RateLimit {
    26  	return ratelimit.NewV2ToV1RateLimit(c, stats)
    27  }
    28  
    29  //------------------------------------------------------------------------------
    30  
    31  // Implements RateLimit around a types.RateLimit
    32  type reverseAirGapRateLimit struct {
    33  	r types.RateLimit
    34  }
    35  
    36  func newReverseAirGapRateLimit(r types.RateLimit) *reverseAirGapRateLimit {
    37  	return &reverseAirGapRateLimit{r}
    38  }
    39  
    40  func (a *reverseAirGapRateLimit) Access(context.Context) (time.Duration, error) {
    41  	return a.r.Access()
    42  }
    43  
    44  func (a *reverseAirGapRateLimit) Close(ctx context.Context) error {
    45  	a.r.CloseAsync()
    46  	for {
    47  		// Gross but will do for now until we replace these with context params.
    48  		if err := a.r.WaitForClose(time.Millisecond * 100); err == nil {
    49  			return nil
    50  		}
    51  		select {
    52  		case <-ctx.Done():
    53  			return ctx.Err()
    54  		default:
    55  		}
    56  	}
    57  }