github.com/MetalBlockchain/metalgo@v1.11.9/network/throttling/dial_throttler.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package throttling
     5  
     6  import (
     7  	"context"
     8  
     9  	"golang.org/x/time/rate"
    10  )
    11  
    12  var (
    13  	_ DialThrottler = (*dialThrottler)(nil)
    14  	_ DialThrottler = (*noDialThrottler)(nil)
    15  )
    16  
    17  type DialThrottler interface {
    18  	// Block until the event associated with this Acquire can happen.
    19  	// If [ctx] is canceled, gives up and returns an error.
    20  	Acquire(ctx context.Context) error
    21  }
    22  
    23  type dialThrottler struct {
    24  	limiter *rate.Limiter
    25  }
    26  
    27  type noDialThrottler struct{}
    28  
    29  func (t dialThrottler) Acquire(ctx context.Context) error {
    30  	return t.limiter.Wait(ctx)
    31  }
    32  
    33  func NewDialThrottler(throttleLimit int) DialThrottler {
    34  	return dialThrottler{
    35  		limiter: rate.NewLimiter(rate.Limit(throttleLimit), throttleLimit),
    36  	}
    37  }
    38  
    39  func NewNoDialThrottler() DialThrottler {
    40  	return noDialThrottler{}
    41  }
    42  
    43  func (noDialThrottler) Acquire(context.Context) error {
    44  	return nil
    45  }