proxy, bufio: fix client network error is summarized as backend network error (#1051)

Author: djshow832

pkg/proxy/backend/backend_conn_mgr.go

@@ -356,7 +356,7 @@ func (mgr *BackendConnManager) ExecuteCmd(ctx context.Context, request []byte) (
 	mgr.processLock.Lock()
 	defer func() {
 		if err != nil && !pnet.IsMySQLError(err) {
-			mgr.setQuitSourceByErr(err)
+			mgr.SetQuitSourceByErr(err)
 		}
 		mgr.handshakeHandler.OnTraffic(mgr)
 		now := time.Now()
@@ -864,7 +864,7 @@ func (mgr *BackendConnManager) setKeepAlive() {
 	}
 }
 
-func (mgr *BackendConnManager) setQuitSourceByErr(err error) {
+func (mgr *BackendConnManager) SetQuitSourceByErr(err error) {
 	if err == nil {
 		return
 	}

pkg/proxy/client/client_conn.go

@@ -72,6 +72,7 @@ func (cc *ClientConnection) processMsg(ctx context.Context) error {
 		cc.pkt.ResetSequence()
 		clientPkt, err := cc.pkt.ReadPacket()
 		if err != nil {
+			cc.connMgr.SetQuitSourceByErr(err)
 			return err
 		}
 		err = cc.connMgr.ExecuteCmd(ctx, clientPkt)

pkg/proxy/net/error.go

@@ -16,7 +16,6 @@ import (
 var (
 	ErrReadConn     = errors.New("failed to read the connection")
 	ErrWriteConn    = errors.New("failed to write the connection")
-	ErrRelayConn    = errors.New("failed to relay the connection")
 	ErrFlushConn    = errors.New("failed to flush the connection")
 	ErrCloseConn    = errors.New("failed to close the connection")
 	ErrHandshakeTLS = errors.New("failed to complete tls handshake")

pkg/proxy/net/packetio.go

@@ -396,7 +396,11 @@ func (p *packetIO) ForwardUntil(destIO PacketIO, isEnd func(firstByte byte, firs
 				return p.wrapErr(errors.Wrap(err, ErrReadConn))
 			}
 			if err := destIO.WritePacket(data, false); err != nil {
-				return p.wrapErr(errors.Wrap(err, ErrWriteConn))
+				err = errors.Wrap(err, ErrWriteConn)
+				if dest != nil {
+					err = dest.wrapErr(err)
+				}
+				return err
 			}
 		} else {
 			for {
@@ -409,7 +413,10 @@ func (p *packetIO) ForwardUntil(destIO PacketIO, isEnd func(firstByte byte, firs
 				dest.readWriter.SetSequence(dest.readWriter.Sequence() + 1)
 				p.limitReader.N = int64(length + 4)
 				if _, err := dest.readWriter.ReadFrom(&p.limitReader); err != nil {
-					return p.wrapErr(errors.Wrap(err, ErrRelayConn))
+					if errors.Is(err, bufio.ErrWriteFail) {
+						return dest.wrapErr(errors.Wrap(err, ErrWriteConn))
+					}
+					return p.wrapErr(errors.Wrap(err, ErrReadConn))
 				}
 				p.inPackets++
 				dest.outPackets++

pkg/proxy/net/packetio_test.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/go-mysql-org/go-mysql/mysql"
 	"github.com/pingcap/tiproxy/lib/config"
+	"github.com/pingcap/tiproxy/lib/util/errors"
 	"github.com/pingcap/tiproxy/lib/util/logger"
 	"github.com/pingcap/tiproxy/lib/util/security"
 	"github.com/pingcap/tiproxy/lib/util/waitgroup"
@@ -558,6 +559,48 @@ func TestForwardUntilLongData(t *testing.T) {
 	wg.Wait()
 }
 
+func TestForwardUntilError(t *testing.T) {
+	srvCh := make(chan *packetIO)
+	var wg waitgroup.WaitGroup
+	selfErr, peerErr := errors.New("self"), errors.New("peer")
+	// client1 writes to server1
+	// server1 forwards to server2
+	// server2 writes to client2 while client2 closes
+	wg.Run(func() {
+		testTCPConn(t,
+			func(t *testing.T, cli *packetIO) {
+				data := make([]byte, DefaultConnBufferSize*2)
+				data[0] = byte(0)
+				require.NoError(t, cli.WritePacket(data, true))
+			},
+			func(t *testing.T, srv1 *packetIO) {
+				srv1.ApplyOpts(WithWrapError(selfErr))
+				srv2 := <-srvCh
+				err := srv1.ForwardUntil(srv2, func(firstByte byte, firstPktLen int) (bool, bool) {
+					return true, true
+				}, func(response []byte) error {
+					return srv2.Flush()
+				})
+				require.ErrorIs(t, err, peerErr)
+			},
+			1,
+		)
+	})
+	wg.Run(func() {
+		testTCPConn(t,
+			func(t *testing.T, cli *packetIO) {
+				require.NoError(t, cli.Close())
+			},
+			func(t *testing.T, srv2 *packetIO) {
+				srv2.ApplyOpts(WithWrapError(peerErr))
+				srvCh <- srv2
+			},
+			1,
+		)
+	})
+	wg.Wait()
+}
+
 func BenchmarkWritePacket(b *testing.B) {
 	b.ReportAllocs()
 	cli, srv := net.Pipe()

pkg/util/bufio/bufio.go

@@ -5,14 +5,15 @@
 // Use of this source code is governed by a BSD-style
 // license that can be found in the LICENSE file.
 
-// Package bufio is a simplified Golang bufio.
-// It mainly avoids calling TCPConn.ReadFrom in Writer.ReadFrom because
-// TCPConn.ReadFrom calls sys calls.
+// Package bufio is a simplified Golang bufio. It mainly changes:
+// - Avoid calling TCPConn.ReadFrom in Writer.ReadFrom because TCPConn.ReadFrom calls sys calls.
+// - Wrap the error with ErrWriteFail/ErrReadFail so that the caller knows which peer fails.
 package bufio
 
 import (
-	"errors"
 	"io"
+
+	"github.com/pingcap/tiproxy/lib/util/errors"
 )
 
 const (
@@ -22,6 +23,9 @@ const (
 var (
 	ErrBufferFull    = errors.New("bufio: buffer full")
 	ErrNegativeCount = errors.New("bufio: negative count")
+
+	ErrReadFail  = errors.New("read failed")
+	ErrWriteFail = errors.New("write failed")
 )
 
 // Buffered input.
@@ -250,9 +254,12 @@ func (b *Reader) Buffered() int { return b.w - b.r }
 // This may make multiple calls to the Read method of the underlying Reader.
 // If the underlying reader supports the WriteTo method,
 // this calls the underlying WriteTo without buffering.
+//
+// Wrap the error with ErrWriteFail/ErrReadFail so that the caller knows which peer fails.
 func (b *Reader) WriteTo(w io.Writer) (n int64, err error) {
 	n, err = b.writeBuf(w)
 	if err != nil {
+		err = errors.Wrap(err, ErrWriteFail)
 		return
 	}
 
@@ -280,7 +287,7 @@ func (b *Reader) WriteTo(w io.Writer) (n int64, err error) {
 		b.err = nil
 	}
 
-	return n, b.readErr()
+	return n, errors.Wrap(b.readErr(), ErrReadFail)
 }
 
 var errNegativeWrite = errors.New("bufio: writer returned negative count from Write")
@@ -430,15 +437,17 @@ func (b *Writer) Write(p []byte) (nn int, err error) {
 // supports the ReadFrom method, this calls the underlying ReadFrom.
 // If there is buffered data and an underlying ReadFrom, this fills
 // the buffer and writes it before calling ReadFrom.
+//
+// Wrap the error with ErrWriteFail/ErrReadFail so that the caller knows which peer fails.
 func (b *Writer) ReadFrom(r io.Reader) (n int64, err error) {
 	if b.err != nil {
-		return 0, b.err
+		return 0, errors.Wrap(b.err, ErrWriteFail)
 	}
 	var m int
 	for {
 		if b.Available() == 0 {
 			if err1 := b.Flush(); err1 != nil {
-				return n, err1
+				return n, errors.Wrap(err1, ErrWriteFail)
 			}
 		}
 		nr := 0
@@ -450,7 +459,7 @@ func (b *Writer) ReadFrom(r io.Reader) (n int64, err error) {
 			nr++
 		}
 		if nr == maxConsecutiveEmptyReads {
-			return n, io.ErrNoProgress
+			return n, errors.Wrap(io.ErrNoProgress, ErrReadFail)
 		}
 		b.n += m
 		n += int64(m)
@@ -461,12 +470,14 @@ func (b *Writer) ReadFrom(r io.Reader) (n int64, err error) {
 	if err == io.EOF {
 		// If we filled the buffer exactly, flush preemptively.
 		if b.Available() == 0 {
-			err = b.Flush()
+			if err = b.Flush(); err != nil {
+				err = errors.Wrap(err, ErrWriteFail)
+			}
 		} else {
 			err = nil
 		}
 	}
-	return n, err
+	return n, errors.Wrap(err, ErrReadFail)
 }
 
 // buffered input and output

pkg/util/bufio/bufio_test.go

@@ -0,0 +1,47 @@
+// Copyright 2025 PingCAP, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package bufio
+
+import (
+	"bytes"
+	"io"
+	"net"
+	"testing"
+
+	"github.com/pingcap/tiproxy/pkg/util/waitgroup"
+	"github.com/stretchr/testify/require"
+)
+
+func TestReadFrom(t *testing.T) {
+	// The reader breaks and ReadFrom returns ErrReadFail.
+	var wg waitgroup.WaitGroup
+	p1, p2 := net.Pipe()
+	writer := bytes.NewBuffer(make([]byte, 10))
+	bufWriter := NewWriter(writer)
+	wg.Run(func() {
+		_, err := p1.Write([]byte("h"))
+		require.NoError(t, err)
+		require.NoError(t, p2.Close())
+		require.NoError(t, p1.Close())
+	}, nil)
+	wg.Run(func() {
+		n, err := bufWriter.ReadFrom(io.LimitReader(p2, 1))
+		require.NoError(t, err)
+		require.EqualValues(t, 1, n)
+		_, err = bufWriter.ReadFrom(p2)
+		require.ErrorIs(t, err, ErrReadFail)
+	}, nil)
+	wg.Wait()
+
+	// The writer breaks and ReadFrom returns ErrWriteFail.
+	p1, p2 = net.Pipe()
+	bufWriter = NewWriterSize(p1, 2)
+	reader := bytes.NewBuffer([]byte("he"))
+	n, err := bufWriter.ReadFrom(io.LimitReader(reader, 1))
+	require.NoError(t, err)
+	require.EqualValues(t, 1, n)
+	require.NoError(t, p2.Close())
+	_, err = bufWriter.ReadFrom(io.LimitReader(reader, 1))
+	require.ErrorIs(t, err, ErrWriteFail)
+}