func (conn *ThrottledConn) Write(buffer []byte) (int, error) {
// See comments in Read.
conn.writeLock.Lock()
defer conn.writeLock.Unlock()
if atomic.LoadInt64(&conn.writeUnthrottledBytes) > 0 {
n, err := conn.Conn.Write(buffer)
atomic.AddInt64(&conn.writeUnthrottledBytes, -int64(n))
return n, err
}
if atomic.LoadInt32(&conn.closeAfterExhausted) == 1 {
conn.Conn.Close()
return 0, errors.New("throttled conn exhausted")
}
rate := atomic.SwapInt64(&conn.writeBytesPerSecond, -1)
if rate != -1 {
if rate == 0 {
conn.throttledWriter = conn.Conn
} else {
conn.throttledWriter = ratelimit.Writer(
conn.Conn,
ratelimit.NewBucketWithRate(float64(rate), rate))
}
}
return conn.throttledWriter.Write(buffer)
}
func (conn *ThrottledConn) Read(buffer []byte) (int, error) {
// A mutex is used to ensure conformance with net.Conn
// concurrency semantics. The atomic.SwapInt64 and
// subsequent assignment of throttledReader could be
// a race condition with concurrent reads.
conn.readLock.Lock()
defer conn.readLock.Unlock()
// Use the base conn until the unthrottled count is
// exhausted. This is only an approximate enforcement
// since this read, or concurrent reads, could exceed
// the remaining count.
if atomic.LoadInt64(&conn.readUnthrottledBytes) > 0 {
n, err := conn.Conn.Read(buffer)
atomic.AddInt64(&conn.readUnthrottledBytes, -int64(n))
return n, err
}
if atomic.LoadInt32(&conn.closeAfterExhausted) == 1 {
conn.Conn.Close()
return 0, errors.New("throttled conn exhausted")
}
rate := atomic.SwapInt64(&conn.readBytesPerSecond, -1)
if rate != -1 {
// SetLimits has been called and a new rate limiter
// must be initialized. When no limit is specified,
// the reader/writer is simply the base conn.
// No state is retained from the previous rate limiter,
// so a pending I/O throttle sleep may be skipped when
// the old and new rate are similar.
if rate == 0 {
conn.throttledReader = conn.Conn
} else {
conn.throttledReader = ratelimit.Reader(
conn.Conn,
ratelimit.NewBucketWithRate(float64(rate), rate))
}
}
return conn.throttledReader.Read(buffer)
}