diff --git a/common/rate/conn.go b/common/rate/conn.go index c94b041..5b3e68a 100644 --- a/common/rate/conn.go +++ b/common/rate/conn.go @@ -19,13 +19,19 @@ type Conn struct { } func (c *Conn) Read(b []byte) (n int, err error) { - c.limiter.Wait(int64(len(b))) - return c.Conn.Read(b) + n, err = c.Conn.Read(b) + if n > 0 { + c.limiter.Wait(int64(n)) + } + return n, err } func (c *Conn) Write(b []byte) (n int, err error) { - c.limiter.Wait(int64(len(b))) - return c.Conn.Write(b) + n, err = c.Conn.Write(b) + if n > 0 { + c.limiter.Wait(int64(n)) + } + return n, err } /*