backend, net: add more logs (#482)

Author: djshow832

pkg/proxy/backend/backend_conn_mgr.go

@@ -9,7 +9,6 @@ import (
 	"encoding/binary"
 	"encoding/json"
 	"fmt"
-
 	"net"
 	"os"
 	"strings"
@@ -20,6 +19,7 @@ import (
 
 	"github.com/cenkalti/backoff/v4"
 	"github.com/go-mysql-org/go-mysql/mysql"
+	"github.com/pingcap/tidb/parser"
 	"github.com/pingcap/tiproxy/lib/config"
 	"github.com/pingcap/tiproxy/lib/util/errors"
 	"github.com/pingcap/tiproxy/lib/util/waitgroup"
@@ -277,19 +277,34 @@ func (mgr *BackendConnManager) getBackendIO(ctx context.Context, cctx ConnContex
 // ExecuteCmd forwards messages between the client and the backend.
 // If it finds that the session is ready for redirection, it migrates the session.
 func (mgr *BackendConnManager) ExecuteCmd(ctx context.Context, request []byte) (err error) {
+	startTime := monotime.Now()
 	mgr.processLock.Lock()
 	defer func() {
 		mgr.setQuitSourceByErr(err)
 		mgr.handshakeHandler.OnTraffic(mgr)
-		mgr.lastActiveTime = monotime.Now()
+		now := monotime.Now()
+		if err != nil && errors.Is(err, ErrBackendConn) {
+			cmd, data := pnet.Command(request[0]), request[1:]
+			var query string
+			if cmd == pnet.ComQuery {
+				query = parser.Normalize(pnet.ParseQueryPacket(data))
+				if len(query) > 256 {
+					query = query[:256]
+				}
+			}
+			// idle_time: maybe the idle time exceeds wait_timeout?
+			// execute_time and query: maybe this query causes TiDB OOM?
+			mgr.logger.Info("backend disconnects", zap.Duration("idle_time", time.Duration(now-mgr.lastActiveTime)),
+				zap.Duration("execute_time", time.Duration(now-startTime)), zap.Stringer("cmd", cmd), zap.String("query", query))
+		}
+		mgr.lastActiveTime = now
 		mgr.processLock.Unlock()
 	}()
 	if len(request) < 1 {
 		err = mysql.ErrMalformPacket
 		return
 	}
 	cmd := pnet.Command(request[0])
-	startTime := monotime.Now()
 
 	// Once the request is accepted, it's treated in the transaction, so we don't check graceful shutdown here.
 	if mgr.closeStatus.Load() >= statusClosing {

pkg/proxy/backend/backend_conn_mgr_test.go

@@ -1297,6 +1297,49 @@ func TestTrafficMetrics(t *testing.T) {
 	ts.runTests(runners)
 }
 
+func TestDisconnectLog(t *testing.T) {
+	ts := newBackendMgrTester(t)
+	tests := []struct {
+		runner  runner
+		checker checker
+	}{
+		{
+			// 1st handshake
+			runner: runner{
+				client:  ts.mc.authenticate,
+				proxy:   ts.firstHandshake4Proxy,
+				backend: ts.handshake4Backend,
+			},
+		},
+		{
+			// proxy logs SQL when the backend disconnects
+			runner: runner{
+				client: func(packetIO *pnet.PacketIO) error {
+					ts.mc.sql = "select 1"
+					return ts.mc.request(packetIO)
+				},
+				proxy: func(clientIO, backendIO *pnet.PacketIO) error {
+					err := ts.forwardCmd4Proxy(clientIO, backendIO)
+					_ = clientIO.Close()
+					return err
+				},
+				backend: func(packetIO *pnet.PacketIO) error {
+					return packetIO.Close()
+				},
+			},
+			checker: func(t *testing.T, ts *testSuite) {
+				require.True(t, pnet.IsDisconnectError(ts.mc.err))
+				require.ErrorIs(t, ts.mp.err, ErrBackendConn)
+				require.True(t, strings.Contains(ts.mp.text.String(), "select ?"))
+			},
+		},
+	}
+	// Do not run ts.runTests(runners) to skip the general checker.
+	for _, test := range tests {
+		ts.runAndCheck(ts.t, test.checker, test.runner.client, test.runner.backend, test.runner.proxy)
+	}
+}
+
 func BenchmarkSyncMap(b *testing.B) {
 	for i := 0; i < b.N; i++ {
 		var m sync.Map

pkg/proxy/net/mysql.go

@@ -467,3 +467,12 @@ func Attr2ZapFields(attrs map[string]string) []zap.Field {
 	}
 	return fields
 }
+
+// ParseQueryPacket returns the statement in the CMD_QUERY packet.
+// data is the payload after byte CMD_QUERY.
+func ParseQueryPacket(data []byte) string {
+	if len(data) > 0 && data[len(data)-1] == 0 {
+		data = data[:len(data)-1]
+	}
+	return hack.String(data)
+}

pkg/proxy/net/packetio.go

@@ -247,7 +247,11 @@ func (p *PacketIO) readOnePacket() ([]byte, bool, error) {
 	if err := ReadFull(p.readWriter, p.header[:]); err != nil {
 		return nil, false, errors.Wrap(ErrReadConn, err)
 	}
-	p.readWriter.SetSequence(p.header[3] + 1)
+	sequence, pktSequence := p.header[3], p.readWriter.Sequence()
+	if sequence != pktSequence {
+		p.logger.Warn("sequence mismatch", zap.Uint8("expected", pktSequence), zap.Uint8("actual", sequence))
+	}
+	p.readWriter.SetSequence(sequence + 1)
 
 	length := int(p.header[0]) | int(p.header[1])<<8 | int(p.header[2])<<16
 	data := make([]byte, length)
@@ -349,7 +353,11 @@ func (p *PacketIO) ForwardUntil(dest *PacketIO, isEnd func(firstByte byte, first
 			}
 		} else {
 			for {
-				p.readWriter.SetSequence(header[3] + 1)
+				sequence, pktSequence := header[3], p.readWriter.Sequence()
+				if sequence != pktSequence {
+					p.logger.Warn("sequence mismatch", zap.Uint8("expected", pktSequence), zap.Uint8("actual", sequence))
+				}
+				p.readWriter.SetSequence(sequence + 1)
 				// Sequence may be different (e.g. with compression) so we can't just copy the data to the destination.
 				dest.readWriter.SetSequence(dest.readWriter.Sequence() + 1)
 				p.limitReader.N = int64(length + 4)