internal: return an error on setting deadline error

internal/pool/conn.go

@@ -68,25 +68,31 @@ func (cn *Conn) NextID() string {
 }
 
 func (cn *Conn) WithReader(
-	c context.Context, timeout time.Duration, fn func(rd *internal.BufReader) error,
+	ctx context.Context, timeout time.Duration, fn func(rd *internal.BufReader) error,
 ) error {
-	_ = cn.netConn.SetReadDeadline(cn.deadline(c, timeout))
+	err := cn.netConn.SetReadDeadline(cn.deadline(ctx, timeout))
+	if err != nil {
+		return err
+	}
 	return fn(cn.rd)
 }
 
 func (cn *Conn) WithWriter(
-	c context.Context, timeout time.Duration, fn func(wb *WriteBuffer) error,
+	ctx context.Context, timeout time.Duration, fn func(wb *WriteBuffer) error,
 ) error {
-	_ = cn.netConn.SetWriteDeadline(cn.deadline(c, timeout))
-	firstErr := fn(cn.wb)
+	err := cn.netConn.SetWriteDeadline(cn.deadline(ctx, timeout))
+	if err != nil {
+		return err
+	}
 
-	buf := cn.wb.Flush()
-	_, err := cn.netConn.Write(buf)
-	if err != nil && firstErr == nil {
-		firstErr = err
+	cn.wb.Reset()
+	err = fn(cn.wb)
+	if err != nil {
+		return err
 	}
 
-	return firstErr
+	_, err = cn.netConn.Write(cn.wb.Bytes)
+	return err
 }
 
 func (cn *Conn) Close() error {

internal/pool/write_buffer.go

@@ -25,12 +25,6 @@ func (buf *WriteBuffer) ResetBuffer(b []byte) {
 	buf.Bytes = b[:0]
 }
 
-func (buf *WriteBuffer) Flush() []byte {
-	b := buf.Bytes
-	buf.Reset()
-	return b
-}
-
 func (buf *WriteBuffer) StartMessage(c byte) {
 	if c == 0 {
 		buf.msgStart = len(buf.Bytes)

main_test.go

@@ -158,6 +158,14 @@ type badConn struct {
 
 var _ net.Conn = &badConn{}
 
+func (cn *badConn) SetReadDeadline(t time.Time) error {
+	return nil
+}
+
+func (cn *badConn) SetWriteDeadline(t time.Time) error {
+	return nil
+}
+
 func (cn *badConn) Read([]byte) (int, error) {
 	if cn.readDelay != 0 {
 		time.Sleep(cn.readDelay)