Last active
July 25, 2022 00:18
-
-
Save harshavardhana/547ab92014ba5b1ccd5f8eea9e3a2e97 to your computer and use it in GitHub Desktop.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package deadlineconn | |
// DeadlineConn - is a generic stream-oriented network connection supporting buffered reader and read/write timeout. | |
type DeadlineConn struct { | |
net.Conn | |
readTimeout time.Duration // sets the read timeout in the connection. | |
writeTimeout time.Duration // sets the write timeout in the connection. | |
} | |
// Sets read timeout | |
func (c *DeadlineConn) setReadTimeout() { | |
if c.readTimeout != 0 { | |
c.SetReadDeadline(time.Now().UTC().Add(c.readTimeout)) | |
} | |
} | |
func (c *DeadlineConn) setWriteTimeout() { | |
if c.writeTimeout != 0 { | |
c.SetWriteDeadline(time.Now().UTC().Add(c.writeTimeout)) | |
} | |
} | |
// Read - reads data from the connection using wrapped buffered reader. | |
func (c *DeadlineConn) Read(b []byte) (n int, err error) { | |
c.setReadTimeout() | |
n, err = c.Conn.Read(b) | |
return n, err | |
} | |
// Write - writes data to the connection. | |
func (c *DeadlineConn) Write(b []byte) (n int, err error) { | |
c.setWriteTimeout() | |
n, err = c.Conn.Write(b) | |
return n, err | |
} | |
// NewDeadlineConn - creates a new connection object wrapping net.Conn with deadlines. | |
func NewDeadlineConn(c net.Conn, readTimeout, writeTimeout time.Duration) net.Conn { | |
return &DeadlineConn{ | |
Conn: c, | |
readTimeout: readTimeout, | |
writeTimeout: writeTimeout, | |
} | |
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
package deadlineconn | |
// Test deadlineconn handles read timeout properly by reading two messages beyond deadline. | |
func TestBuffConnReadTimeout(t *testing.T) { | |
l, err := net.Listen("tcp", "localhost:0") | |
if err != nil { | |
t.Fatalf("unable to create listener. %v", err) | |
} | |
defer l.Close() | |
serverAddr := l.Addr().String() | |
tcpListener, ok := l.(*net.TCPListener) | |
if !ok { | |
t.Fatalf("failed to assert to net.TCPListener") | |
} | |
var wg sync.WaitGroup | |
wg.Add(1) | |
go func() { | |
defer wg.Done() | |
tcpConn, terr := tcpListener.AcceptTCP() | |
if terr != nil { | |
t.Errorf("failed to accept new connection. %v", terr) | |
return | |
} | |
deadlineconn := NewDeadlineConn(tcpConn, 1*time.Second, 1*time.Second) | |
defer deadlineconn.Close() | |
// Read a line | |
var b = make([]byte, 12) | |
_, terr = deadlineconn.Read(b) | |
if terr != nil { | |
t.Errorf("failed to read from client. %v", terr) | |
return | |
} | |
received := string(b) | |
if received != "message one\n" { | |
t.Errorf(`server: expected: "message one\n", got: %v`, received) | |
return | |
} | |
// Wait for more than read timeout to simulate processing. | |
time.Sleep(3 * time.Second) | |
_, terr = deadlineconn.Read(b) | |
if terr != nil { | |
t.Errorf("failed to read from client. %v", terr) | |
return | |
} | |
received = string(b) | |
if received != "message two\n" { | |
t.Errorf(`server: expected: "message two\n", got: %v`, received) | |
return | |
} | |
// Send a response. | |
_, terr = io.WriteString(deadlineconn, "messages received\n") | |
if terr != nil { | |
t.Errorf("failed to write to client. %v", terr) | |
return | |
} | |
}() | |
c, err := net.Dial("tcp", serverAddr) | |
if err != nil { | |
t.Fatalf("unable to connect to server. %v", err) | |
} | |
defer c.Close() | |
_, err = io.WriteString(c, "message one\n") | |
if err != nil { | |
t.Fatalf("failed to write to server. %v", err) | |
} | |
_, err = io.WriteString(c, "message two\n") | |
if err != nil { | |
t.Fatalf("failed to write to server. %v", err) | |
} | |
received, err := bufio.NewReader(c).ReadString('\n') | |
if err != nil { | |
t.Fatalf("failed to read from server. %v", err) | |
} | |
if received != "messages received\n" { | |
t.Fatalf(`client: expected: "messages received\n", got: %v`, received) | |
} | |
wg.Wait() | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment