github.com/Uhtred009/v2ray-core-1@v4.31.2+incompatible/app/dispatcher/rate.go (about)

     1  package dispatcher
     2  
     3  import (
     4  	"time"
     5  
     6  	"v2ray.com/core/common/buf"
     7  )
     8  
     9  type Writer struct {
    10  	writer  buf.Writer
    11  	limiter *RateLimiter
    12  }
    13  
    14  type RateLimiter struct {
    15  	rate  time.Duration
    16  	count int64
    17  	t     time.Time
    18  }
    19  
    20  func NewRateLimiter(rate int64) *RateLimiter {
    21  	return &RateLimiter{
    22  		rate:  time.Duration(rate),
    23  		count: 0,
    24  		t:     time.Now(),
    25  	}
    26  }
    27  
    28  func RateWriter(writer buf.Writer, limiter *RateLimiter) buf.Writer {
    29  	return &Writer{
    30  		writer:  writer,
    31  		limiter: limiter,
    32  	}
    33  }
    34  
    35  func (l *RateLimiter) RateWait(count int64) {
    36  	l.count += count
    37  	t := time.Duration(l.count)*time.Second/l.rate - time.Since(l.t)
    38  	if t > 0 {
    39  		time.Sleep(t)
    40  	}
    41  }
    42  
    43  func (w *Writer) WriteMultiBuffer(mb buf.MultiBuffer) error {
    44  	w.limiter.RateWait(int64(mb.Len()))
    45  	return w.writer.WriteMultiBuffer(mb)
    46  }