github.com/Jeffail/benthos/v3@v3.65.0/internal/bundle/ratelimits.go (about)

     1  package bundle
     2  
     3  import (
     4  	"sort"
     5  
     6  	"github.com/Jeffail/benthos/v3/internal/docs"
     7  	"github.com/Jeffail/benthos/v3/lib/ratelimit"
     8  	"github.com/Jeffail/benthos/v3/lib/types"
     9  )
    10  
    11  // AllRateLimits is a set containing every single ratelimit that has been imported.
    12  var AllRateLimits = &RateLimitSet{
    13  	specs: map[string]rateLimitSpec{},
    14  }
    15  
    16  //------------------------------------------------------------------------------
    17  
    18  // RateLimitAdd adds a new ratelimit to this environment by providing a
    19  // constructor and documentation.
    20  func (e *Environment) RateLimitAdd(constructor RateLimitConstructor, spec docs.ComponentSpec) error {
    21  	return e.rateLimits.Add(constructor, spec)
    22  }
    23  
    24  // RateLimitInit attempts to initialise a ratelimit from a config.
    25  func (e *Environment) RateLimitInit(conf ratelimit.Config, mgr NewManagement) (types.RateLimit, error) {
    26  	return e.rateLimits.Init(conf, mgr)
    27  }
    28  
    29  // RateLimitDocs returns a slice of ratelimit specs, which document each method.
    30  func (e *Environment) RateLimitDocs() []docs.ComponentSpec {
    31  	return e.rateLimits.Docs()
    32  }
    33  
    34  //------------------------------------------------------------------------------
    35  
    36  // RateLimitConstructor constructs an ratelimit component.
    37  type RateLimitConstructor func(ratelimit.Config, NewManagement) (types.RateLimit, error)
    38  
    39  type rateLimitSpec struct {
    40  	constructor RateLimitConstructor
    41  	spec        docs.ComponentSpec
    42  }
    43  
    44  // RateLimitSet contains an explicit set of ratelimits available to a Benthos service.
    45  type RateLimitSet struct {
    46  	specs map[string]rateLimitSpec
    47  }
    48  
    49  // Add a new ratelimit to this set by providing a spec (name, documentation, and
    50  // constructor).
    51  func (s *RateLimitSet) Add(constructor RateLimitConstructor, spec docs.ComponentSpec) error {
    52  	if s.specs == nil {
    53  		s.specs = map[string]rateLimitSpec{}
    54  	}
    55  	s.specs[spec.Name] = rateLimitSpec{
    56  		constructor: constructor,
    57  		spec:        spec,
    58  	}
    59  	docs.RegisterDocs(spec)
    60  	return nil
    61  }
    62  
    63  // Init attempts to initialise an ratelimit from a config.
    64  func (s *RateLimitSet) Init(conf ratelimit.Config, mgr NewManagement) (types.RateLimit, error) {
    65  	spec, exists := s.specs[conf.Type]
    66  	if !exists {
    67  		// TODO: V4 Remove this
    68  		if ctor, exists := ratelimit.GetDeprecatedPlugin(conf.Type); exists {
    69  			return ctor(conf, mgr, mgr.Logger(), mgr.Metrics())
    70  		}
    71  		return nil, types.ErrInvalidRateLimitType
    72  	}
    73  	return spec.constructor(conf, mgr)
    74  }
    75  
    76  // Docs returns a slice of ratelimit specs, which document each method.
    77  func (s *RateLimitSet) Docs() []docs.ComponentSpec {
    78  	var docs []docs.ComponentSpec
    79  	for _, v := range s.specs {
    80  		docs = append(docs, v.spec)
    81  	}
    82  	sort.Slice(docs, func(i, j int) bool {
    83  		return docs[i].Name < docs[j].Name
    84  	})
    85  	return docs
    86  }
    87  
    88  // DocsFor returns the documentation for a given component name, returns a
    89  // boolean indicating whether the component name exists.
    90  func (s *RateLimitSet) DocsFor(name string) (docs.ComponentSpec, bool) {
    91  	c, ok := s.specs[name]
    92  	if !ok {
    93  		return docs.ComponentSpec{}, false
    94  	}
    95  	return c.spec, true
    96  }