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 }