@@ -23,10 +23,11 @@ const maxCachedBufSize = 256 * 1024
2323// The buffer is similar to bufio.Reader / Writer but zero-copy-ish
2424// Also highly optimized for this particular use case.
2525type buffer struct {
26- buf []byte // read buffer.
27- cachedBuf []byte // buffer that will be reused. len(cachedBuf) <= maxCachedBufSize.
28- nc net.Conn
29- timeout time.Duration
26+ buf []byte // read buffer.
27+ cachedBuf []byte // buffer that will be reused. len(cachedBuf) <= maxCachedBufSize.
28+ nc net.Conn
29+ readTimeout time.Duration
30+ writeTimeout time.Duration
3031}
3132
3233// newBuffer allocates and returns a new buffer.
@@ -64,8 +65,8 @@ func (b *buffer) fill(need int) error {
6465 copy (dest [:n ], b .buf )
6566
6667 for {
67- if b .timeout > 0 {
68- if err := b .nc .SetReadDeadline (time .Now ().Add (b .timeout )); err != nil {
68+ if b .readTimeout > 0 {
69+ if err := b .nc .SetReadDeadline (time .Now ().Add (b .readTimeout )); err != nil {
6970 return err
7071 }
7172 }
@@ -159,5 +160,10 @@ func (b *buffer) store(buf []byte) {
159160// writePackets is a proxy function to nc.Write.
160161// This is used to make the buffer type compatible with compressed I/O.
161162func (b * buffer ) writePackets (packets []byte ) (int , error ) {
163+ if b .writeTimeout > 0 {
164+ if err := b .nc .SetWriteDeadline (time .Now ().Add (b .writeTimeout )); err != nil {
165+ return 0 , err
166+ }
167+ }
162168 return b .nc .Write (packets )
163169}
0 commit comments