github.com/webx-top/com@v1.2.12/ratelimit/ratelimit.go (about)

     1  package ratelimit
     2  
     3  import (
     4  	"io"
     5  	"time"
     6  )
     7  
     8  // New 创建Limiter
     9  // Usage:
    10  // limiter := New(100*1024)
    11  // limiter.NewReader(reader)
    12  // limiter.NewWriter(writer)
    13  // limiter.NewReadSeeker(readSeeker)
    14  func New(bytesPerSecond int64) *Limiter {
    15  	return &Limiter{
    16  		rate:  time.Duration(bytesPerSecond),
    17  		count: 0,
    18  		t:     time.Now(),
    19  	}
    20  }
    21  
    22  // Reader 返回一个带有Limiter的io.Reader
    23  func Reader(r io.Reader, l *Limiter) io.Reader {
    24  	return &reader{
    25  		r: r,
    26  		l: l,
    27  	}
    28  }
    29  
    30  // ReadSeeker 返回一个带有Limiter的io.ReadSeeker
    31  func ReadSeeker(rs io.ReadSeeker, l *Limiter) io.ReadSeeker {
    32  	return &readSeeker{
    33  		reader: reader{
    34  			r: rs,
    35  			l: l,
    36  		},
    37  		s: rs,
    38  	}
    39  }
    40  
    41  // Writer 返回一个带有Limiter的io.Writer
    42  func Writer(w io.Writer, l *Limiter) io.Writer {
    43  	return &writer{
    44  		w: w,
    45  		l: l,
    46  	}
    47  }
    48  
    49  // Limiter 速度限制器
    50  type Limiter struct {
    51  	rate  time.Duration
    52  	count int64 // 最大8G
    53  	t     time.Time
    54  }
    55  
    56  // Wait 传入需要处理的数量，计算并等待需要经过的时间
    57  func (l *Limiter) Wait(count int) {
    58  	l.count += int64(count)
    59  	t := time.Duration(l.count)*time.Second/l.rate - time.Since(l.t)
    60  	if t > 0 {
    61  		time.Sleep(t)
    62  	}
    63  }
    64  
    65  func (l *Limiter) NewReader(r io.Reader) io.Reader {
    66  	return Reader(r, l)
    67  }
    68  
    69  func (l *Limiter) NewReadSeeker(r io.ReadSeeker) io.ReadSeeker {
    70  	return ReadSeeker(r, l)
    71  }
    72  
    73  func (l *Limiter) NewWriter(r io.Writer) io.Writer {
    74  	return Writer(r, l)
    75  }
    76  
    77  type reader struct {
    78  	r io.Reader
    79  	l *Limiter
    80  }
    81  
    82  // Read Read
    83  func (r *reader) Read(buf []byte) (int, error) {
    84  	n, err := r.r.Read(buf)
    85  	r.l.Wait(n)
    86  	return n, err
    87  }
    88  
    89  type readSeeker struct {
    90  	reader
    91  	s io.Seeker
    92  }
    93  
    94  func (rs *readSeeker) Seek(offset int64, whence int) (int64, error) {
    95  	return rs.s.Seek(offset, whence)
    96  }
    97  
    98  type writer struct {
    99  	w io.Writer
   100  	l *Limiter
   101  }
   102  
   103  // Write Write
   104  func (w *writer) Write(buf []byte) (int, error) {
   105  	w.l.Wait(len(buf))
   106  	return w.w.Write(buf)
   107  }