sqlreplay, backend: fix bugs to make traffic replay runnable (#671)

Author: djshow832

pkg/proxy/backend/backend_conn_mgr.go

@@ -321,7 +321,9 @@ func (mgr *BackendConnManager) ExecuteCmd(ctx context.Context, request []byte) (
 	}
 	mgr.processLock.Lock()
 	defer func() {
-		mgr.setQuitSourceByErr(err)
+		if err != nil && !pnet.IsMySQLError(err) {
+			mgr.setQuitSourceByErr(err)
+		}
 		mgr.handshakeHandler.OnTraffic(mgr)
 		now := time.Now()
 		if err != nil && errors.Is(err, ErrBackendConn) {
@@ -403,12 +405,7 @@ func (mgr *BackendConnManager) ExecuteCmd(ctx context.Context, request []byte) (
 		_, err = mgr.cmdProcessor.executeCmd(request, mgr.clientIO, backendIO, false)
 		addCmdMetrics(cmd, backendIO.RemoteAddr().String(), startTime)
 		mgr.updateTraffic(backendIO)
-		if err != nil && !pnet.IsMySQLError(err) {
-			return
-		}
 	}
-	// Ignore MySQL errors, only return unexpected errors.
-	err = nil
 	return
 }
 

pkg/proxy/backend/backend_conn_mgr_test.go

@@ -199,6 +199,9 @@ func (ts *backendMgrTester) forwardCmd4Proxy(clientIO, backendIO pnet.PacketIO)
 	prevCounter, err := readCmdCounter(pnet.Command(request[0]), ts.tc.backendListener.Addr().String())
 	require.NoError(ts.t, err)
 	rsErr := ts.mp.ExecuteCmd(context.Background(), request)
+	if pnet.IsMySQLError(rsErr) {
+		rsErr = nil
+	}
 	curCounter, err := readCmdCounter(pnet.Command(request[0]), ts.tc.backendListener.Addr().String())
 	require.NoError(ts.t, err)
 	require.Equal(ts.t, prevCounter+1, curCounter)
@@ -575,6 +578,41 @@ func TestSpecialCmds(t *testing.T) {
 	ts.runTests(runners)
 }
 
+// Test that ExecuteCmd may return a mysql error, which is required by traffic replay.
+func TestReturnMySQLError(t *testing.T) {
+	ts := newBackendMgrTester(t)
+	runners := []runner{
+		// 1st handshake
+		{
+			client:  ts.mc.authenticate,
+			proxy:   ts.firstHandshake4Proxy,
+			backend: ts.handshake4Backend,
+		},
+		// mysql error
+		{
+			client: func(packetIO pnet.PacketIO) error {
+				ts.mc.cmd = pnet.ComQuery
+				ts.mc.sql = "select $$"
+				return ts.mc.request(packetIO)
+			},
+			proxy: func(clientIO, backendIO pnet.PacketIO) error {
+				clientIO.ResetSequence()
+				request, err := clientIO.ReadPacket()
+				require.NoError(ts.t, err)
+				rsErr := ts.mp.ExecuteCmd(context.Background(), request)
+				require.True(ts.t, pnet.IsMySQLError(rsErr))
+				require.Equal(ts.t, SrcNone, ts.mp.QuitSource())
+				return nil
+			},
+			backend: func(packetIO pnet.PacketIO) error {
+				ts.mb.respondType = responseTypeErr
+				return ts.mb.respond(packetIO)
+			},
+		},
+	}
+	ts.runTests(runners)
+}
+
 // Test that closing the BackendConnMgr while it's receiving a redirection signal is OK.
 func TestCloseWhileRedirect(t *testing.T) {
 	ts := newBackendMgrTester(t)

pkg/proxy/client/client_conn.go

@@ -75,7 +75,7 @@ func (cc *ClientConnection) processMsg(ctx context.Context) error {
 			return err
 		}
 		err = cc.connMgr.ExecuteCmd(ctx, clientPkt)
-		if err != nil {
+		if err != nil && !pnet.IsMySQLError(err) {
 			return err
 		}
 		if pnet.Command(clientPkt[0]) == pnet.ComQuit {

pkg/proxy/net/packetio.go

@@ -366,25 +366,28 @@ func (p *packetIO) WritePacket(data []byte, flush bool) (err error) {
 func (p *packetIO) ForwardUntil(destIO PacketIO, isEnd func(firstByte byte, firstPktLen int) (end, needData bool),
 	process func(response []byte) error) error {
 	p.readWriter.BeginRW(rwRead)
-	dest := destIO.(*packetIO)
-	dest.readWriter.BeginRW(rwWrite)
+	dest, _ := destIO.(*packetIO)
+	// destIO is not packetIO in traffic replay.
+	if dest != nil {
+		dest.readWriter.BeginRW(rwWrite)
+	}
 	p.limitReader.R = p.readWriter
 	for {
 		header, err := p.readWriter.Peek(5)
 		if err != nil {
-			return p.wrapErr(errors.Wrap(err, ErrReadConn))
+			return p.wrapErr(errors.Wrap(errors.WithStack(err), ErrReadConn))
 		}
 		length := int(header[0]) | int(header[1])<<8 | int(header[2])<<16
 		end, needData := isEnd(header[4], length)
 		var data []byte
 		// Just call ReadFrom if the caller doesn't need the data, even if it's the last packet.
-		if end && needData {
+		if (end && needData) || dest == nil {
 			// TODO: allocate a buffer from pool and return the buffer after `process`.
 			data, err = p.ReadPacket()
 			if err != nil {
 				return p.wrapErr(errors.Wrap(err, ErrReadConn))
 			}
-			if err := dest.WritePacket(data, false); err != nil {
+			if err := destIO.WritePacket(data, false); err != nil {
 				return p.wrapErr(errors.Wrap(err, ErrWriteConn))
 			}
 		} else {

pkg/proxy/proxy.go

@@ -150,8 +150,8 @@ func (s *SQLServer) onConn(ctx context.Context, conn net.Conn, addr string) {
 			return false, nil, 0, nil
 		}
 
-		connID := s.mu.connID
 		s.mu.connID++
+		connID := s.mu.connID
 		logger := s.logger.With(zap.Uint64("connID", connID), zap.String("client_addr", conn.RemoteAddr().String()),
 			zap.String("addr", addr))
 		clientConn := client.NewClientConnection(logger.Named("conn"), conn, s.certMgr.ServerSQLTLS(), s.certMgr.SQLTLS(),

pkg/proxy/proxy_test.go

@@ -258,7 +258,7 @@ func TestRecoverPanic(t *testing.T) {
 	require.NoError(t, err)
 	server, err := NewSQLServer(lg, &config.Config{}, certManager, nil, &mockHsHandler{
 		handshakeResp: func(ctx backend.ConnContext, _ *pnet.HandshakeResp) error {
-			if ctx.Value(backend.ConnContextKeyConnID).(uint64) == 0 {
+			if ctx.Value(backend.ConnContextKeyConnID).(uint64) == 1 {
 				panic("HandleHandshakeResp panic")
 			}
 			return nil

pkg/sqlreplay/capture/capture.go

@@ -86,13 +86,17 @@ var _ Capture = (*capture)(nil)
 
 type capture struct {
 	sync.Mutex
-	cfg       CaptureConfig
-	wg        waitgroup.WaitGroup
-	cancel    context.CancelFunc
-	cmdCh     chan *cmd.Command
-	err       error
-	startTime time.Time
-	lg        *zap.Logger
+	cfg          CaptureConfig
+	wg           waitgroup.WaitGroup
+	cancel       context.CancelFunc
+	cmdCh        chan *cmd.Command
+	err          error
+	startTime    time.Time
+	endTime      time.Time
+	progress     float64
+	capturedCmds uint64
+	filteredCmds uint64
+	lg           *zap.Logger
 }
 
 func NewCapture(lg *zap.Logger) *capture {
@@ -115,6 +119,10 @@ func (c *capture) Start(cfg CaptureConfig) error {
 	c.stopNoLock(nil)
 	c.cfg = cfg
 	c.startTime = time.Now()
+	c.endTime = time.Time{}
+	c.progress = 0
+	c.capturedCmds = 0
+	c.filteredCmds = 0
 	c.err = nil
 	childCtx, cancel := context.WithTimeout(context.Background(), c.cfg.Duration)
 	c.cancel = cancel
@@ -143,6 +151,7 @@ func (c *capture) collectCmds(bufCh chan<- *bytes.Buffer) {
 			c.Stop(errors.Wrapf(err, "failed to encode command"))
 			return
 		}
+		c.capturedCmds++
 		if buf.Len() > c.cfg.flushThreshold {
 			select {
 			case bufCh <- buf:
@@ -208,7 +217,7 @@ func (c *capture) Progress() (float64, error) {
 	c.Lock()
 	defer c.Unlock()
 	if c.startTime.IsZero() || c.cfg.Duration == 0 {
-		return 0, c.err
+		return c.progress, c.err
 	}
 	return float64(time.Since(c.startTime)) / float64(c.cfg.Duration), c.err
 }
@@ -219,16 +228,33 @@ func (c *capture) stopNoLock(err error) {
 	if c.startTime.IsZero() {
 		return
 	}
-	if err != nil {
-		c.lg.Error("stop capture", zap.Error(err))
-	}
-	c.err = err
 	if c.cancel != nil {
 		c.cancel()
 		c.cancel = nil
 	}
-	c.startTime = time.Time{}
 	close(c.cmdCh)
+
+	c.endTime = time.Now()
+	fields := []zap.Field{
+		zap.Time("start_time", c.startTime),
+		zap.Time("end_time", c.endTime),
+		zap.Uint64("captured_cmds", c.capturedCmds),
+	}
+	if err != nil {
+		if c.cfg.Duration > 0 {
+			c.progress = float64(c.endTime.Sub(c.startTime)) / float64(c.cfg.Duration)
+			if c.progress > 1 {
+				c.progress = 1
+			}
+		}
+		c.err = err
+		fields = append(fields, zap.Error(err))
+		c.lg.Error("capture failed", fields...)
+	} else {
+		c.progress = 1
+		c.lg.Info("capture finished", fields...)
+	}
+	c.startTime = time.Time{}
 }
 
 func (c *capture) Stop(err error) {

pkg/sqlreplay/capture/capture_test.go

@@ -11,16 +11,15 @@ import (
 	"time"
 
 	"github.com/pingcap/tiproxy/lib/util/errors"
-	"github.com/pingcap/tiproxy/lib/util/logger"
 	"github.com/pingcap/tiproxy/lib/util/waitgroup"
 	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/store"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
 )
 
 func TestStartAndStop(t *testing.T) {
-	lg, _ := logger.CreateLoggerForTest(t)
-	cpt := NewCapture(lg)
+	cpt := NewCapture(zap.NewNop())
 	defer cpt.Close()
 
 	packet := append([]byte{pnet.ComQuery.Byte()}, []byte("select 1")...)
@@ -35,15 +34,12 @@ func TestStartAndStop(t *testing.T) {
 	// start capture and the traffic should be outputted
 	require.NoError(t, cpt.Start(cfg))
 	cpt.Capture(packet, time.Now(), 100)
-	_, err := cpt.Progress()
-	require.NoError(t, err)
 	cpt.Stop(errors.Errorf("mock error"))
-	_, err = cpt.Progress()
-	require.ErrorContains(t, err, "mock error")
 	cpt.wg.Wait()
 	data := writer.getData()
 	require.Greater(t, len(data), 0)
 	require.Contains(t, string(data), "select 1")
+	require.Equal(t, uint64(1), cpt.capturedCmds)
 
 	// stop capture and traffic should not be outputted
 	cpt.Capture(packet, time.Now(), 100)
@@ -65,8 +61,7 @@ func TestStartAndStop(t *testing.T) {
 }
 
 func TestConcurrency(t *testing.T) {
-	lg, _ := logger.CreateLoggerForTest(t)
-	cpt := NewCapture(lg)
+	cpt := NewCapture(zap.NewNop())
 	defer cpt.Close()
 
 	writer := newMockWriter(store.WriterCfg{})
@@ -145,3 +140,38 @@ func TestCaptureCfgError(t *testing.T) {
 	require.Equal(t, maxBuffers, cfg.maxBuffers)
 	require.Equal(t, maxPendingCommands, cfg.maxPendingCommands)
 }
+
+func TestProgress(t *testing.T) {
+	cpt := NewCapture(zap.NewNop())
+	defer cpt.Close()
+
+	writer := newMockWriter(store.WriterCfg{})
+	cfg := CaptureConfig{
+		Output:    t.TempDir(),
+		Duration:  10 * time.Second,
+		cmdLogger: writer,
+	}
+	setStartTime := func(t time.Time) {
+		cpt.Lock()
+		cpt.startTime = t
+		cpt.Unlock()
+	}
+
+	now := time.Now()
+	require.NoError(t, cpt.Start(cfg))
+	progress, err := cpt.Progress()
+	require.NoError(t, err)
+	require.Less(t, progress, 0.3)
+
+	setStartTime(now.Add(-5 * time.Second))
+	progress, err = cpt.Progress()
+	require.NoError(t, err)
+	require.GreaterOrEqual(t, progress, 0.5)
+
+	cpt.Stop(errors.Errorf("mock error"))
+	cpt.wg.Wait()
+	progress, err = cpt.Progress()
+	require.ErrorContains(t, err, "mock error")
+	require.GreaterOrEqual(t, progress, 0.5)
+	require.Less(t, progress, 1.0)
+}