vitess.io/vitess@v0.16.2/go/netutil/conn.go (about) 1 /* 2 Copyright 2019 The Vitess Authors. 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 http://www.apache.org/licenses/LICENSE-2.0 7 Unless required by applicable law or agreed to in writing, software 8 distributed under the License is distributed on an "AS IS" BASIS, 9 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 10 See the License for the specific language governing permissions and 11 limitations under the License. 12 */ 13 14 package netutil 15 16 import ( 17 "net" 18 "time" 19 ) 20 21 var _ net.Conn = (*ConnWithTimeouts)(nil) 22 23 // A ConnWithTimeouts is a wrapper to net.Comm that allows to set a read and write timeouts. 24 type ConnWithTimeouts struct { 25 net.Conn 26 readTimeout time.Duration 27 writeTimeout time.Duration 28 } 29 30 // NewConnWithTimeouts wraps a net.Conn with read and write deadilnes. 31 func NewConnWithTimeouts(conn net.Conn, readTimeout time.Duration, writeTimeout time.Duration) ConnWithTimeouts { 32 return ConnWithTimeouts{Conn: conn, readTimeout: readTimeout, writeTimeout: writeTimeout} 33 } 34 35 // Implementation of the Conn interface. 36 37 // Read sets a read deadilne and delegates to conn.Read. 38 func (c ConnWithTimeouts) Read(b []byte) (int, error) { 39 if c.readTimeout == 0 { 40 return c.Conn.Read(b) 41 } 42 if err := c.Conn.SetReadDeadline(time.Now().Add(c.readTimeout)); err != nil { 43 return 0, err 44 } 45 return c.Conn.Read(b) 46 } 47 48 // Write sets a write deadline and delegates to conn.Write 49 func (c ConnWithTimeouts) Write(b []byte) (int, error) { 50 if c.writeTimeout == 0 { 51 return c.Conn.Write(b) 52 } 53 if err := c.Conn.SetWriteDeadline(time.Now().Add(c.writeTimeout)); err != nil { 54 return 0, err 55 } 56 return c.Conn.Write(b) 57 } 58 59 // SetDeadline implements the Conn SetDeadline method. 60 func (c ConnWithTimeouts) SetDeadline(t time.Time) error { 61 panic("can't call SetDeadline for ConnWithTimeouts") 62 } 63 64 // SetReadDeadline implements the Conn SetReadDeadline method. 65 func (c ConnWithTimeouts) SetReadDeadline(t time.Time) error { 66 panic("can't call SetReadDeadline for ConnWithTimeouts") 67 } 68 69 // SetWriteDeadline implements the Conn SetWriteDeadline method. 70 func (c ConnWithTimeouts) SetWriteDeadline(t time.Time) error { 71 panic("can't call SetWriteDeadline for ConnWithTimeouts") 72 }