replay: slow down decoding when there are too many pending commands (#701)

Author: djshow832

pkg/metrics/metrics.go

@@ -32,6 +32,7 @@ const (
 	LabelMonitor = "monitor"
 	LabelBackend = "backend"
 	LabelTraffic = "traffic"
+	LabelReplay  = "replay"
 )
 
 // MetricsManager manages metrics.
@@ -120,6 +121,8 @@ func init() {
 		OutboundBytesCounter,
 		OutboundPacketsCounter,
 		CrossLocationBytesCounter,
+		ReplayPendingCmdsGauge,
+		ReplayWaitTime,
 	}
 }
 

pkg/metrics/replay.go

@@ -0,0 +1,24 @@
+// Copyright 2024 PingCAP, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package metrics
+
+import "github.com/prometheus/client_golang/prometheus"
+
+var (
+	ReplayPendingCmdsGauge = prometheus.NewGauge(
+		prometheus.GaugeOpts{
+			Namespace: ModuleProxy,
+			Subsystem: LabelReplay,
+			Name:      "pending_cmds",
+			Help:      "Counter of pending commands.",
+		})
+
+	ReplayWaitTime = prometheus.NewGauge(
+		prometheus.GaugeOpts{
+			Namespace: ModuleProxy,
+			Subsystem: LabelReplay,
+			Name:      "wait_time",
+			Help:      "Wait time of replaying commands.",
+		})
+)

pkg/proxy/backend/mock_proxy_test.go

@@ -145,8 +145,8 @@ func (mc *mockCapture) Capture(packet []byte, startTime time.Time, connID uint64
 	}
 }
 
-func (mc *mockCapture) Progress() (float64, time.Time, error) {
-	return 0, time.Time{}, nil
+func (mc *mockCapture) Progress() (float64, time.Time, bool, error) {
+	return 0, time.Time{}, false, nil
 }
 
 func (mc *mockCapture) Close() {

pkg/sqlreplay/capture/capture.go

@@ -44,7 +44,7 @@ type Capture interface {
 	// Capture captures traffic
 	Capture(packet []byte, startTime time.Time, connID uint64, initSession func() (string, error))
 	// Progress returns the progress of the capture job
-	Progress() (float64, time.Time, error)
+	Progress() (float64, time.Time, bool, error)
 	// Close closes the capture
 	Close()
 }
@@ -235,9 +235,10 @@ func (c *capture) flushBuffer(bufCh <-chan *bytes.Buffer) {
 	c.Lock()
 	startTime := c.startTime
 	capturedCmds := c.capturedCmds
+	filteredCmds := c.filteredCmds
 	c.Unlock()
 	// Write meta outside of the lock to avoid affecting QPS.
-	c.writeMeta(time.Since(startTime), capturedCmds)
+	c.writeMeta(time.Since(startTime), capturedCmds, filteredCmds)
 }
 
 func (c *capture) InitConn(startTime time.Time, connID uint64, db string) {
@@ -334,24 +335,24 @@ func (c *capture) putCommand(command *cmd.Command) bool {
 	}
 }
 
-func (c *capture) writeMeta(duration time.Duration, cmds uint64) {
-	meta := store.Meta{Duration: duration, Cmds: cmds}
+func (c *capture) writeMeta(duration time.Duration, cmds, filteredCmds uint64) {
+	meta := store.NewMeta(duration, cmds, filteredCmds)
 	if err := meta.Write(c.cfg.Output); err != nil {
 		c.lg.Error("failed to write meta", zap.Error(err))
 	}
 }
 
-func (c *capture) Progress() (float64, time.Time, error) {
+func (c *capture) Progress() (float64, time.Time, bool, error) {
 	c.Lock()
 	defer c.Unlock()
 	if c.status == statusIdle || c.cfg.Duration == 0 {
-		return c.progress, c.endTime, c.err
+		return c.progress, c.endTime, true, c.err
 	}
 	progress := float64(time.Since(c.startTime)) / float64(c.cfg.Duration)
 	if progress > 1 {
 		progress = 1
 	}
-	return progress, c.endTime, c.err
+	return progress, c.endTime, false, c.err
 }
 
 // stopNoLock must be called after holding a lock.

pkg/sqlreplay/capture/capture_test.go

@@ -170,23 +170,25 @@ func TestProgress(t *testing.T) {
 
 	now := time.Now()
 	require.NoError(t, cpt.Start(cfg))
-	progress, _, err := cpt.Progress()
+	progress, _, done, err := cpt.Progress()
 	require.NoError(t, err)
 	require.Less(t, progress, 0.3)
+	require.False(t, done)
 
 	setStartTime(now.Add(-5 * time.Second))
-	progress, _, err = cpt.Progress()
+	progress, _, _, err = cpt.Progress()
 	require.NoError(t, err)
 	require.GreaterOrEqual(t, progress, 0.5)
 
 	packet := append([]byte{pnet.ComQuery.Byte()}, []byte("select 1")...)
 	cpt.Capture(packet, time.Now(), 100, mockInitSession)
 	cpt.Stop(errors.Errorf("mock error"))
 	cpt.wg.Wait()
-	progress, _, err = cpt.Progress()
+	progress, _, done, err = cpt.Progress()
 	require.ErrorContains(t, err, "mock error")
 	require.GreaterOrEqual(t, progress, 0.5)
 	require.Less(t, progress, 1.0)
+	require.True(t, done)
 
 	m := store.Meta{}
 	require.NoError(t, m.Read(cfg.Output))

pkg/sqlreplay/cmd/cmd.go

@@ -28,6 +28,7 @@ const (
 )
 
 type LineReader interface {
+	String() string
 	ReadLine() ([]byte, string, int, error)
 	Read([]byte) (string, int, error)
 	Close()

pkg/sqlreplay/cmd/mock_test.go

@@ -39,3 +39,7 @@ func (mr *mockReader) Read(data []byte) (string, int, error) {
 
 func (mr *mockReader) Close() {
 }
+
+func (mr *mockReader) String() string {
+	return "mockReader"
+}

pkg/sqlreplay/conn/conn.go

@@ -7,49 +7,66 @@ import (
 	"context"
 	"crypto/tls"
 	"encoding/binary"
-	"time"
+	"sync"
+	"sync/atomic"
 
-	"github.com/pingcap/tiproxy/lib/util/errors"
+	glist "github.com/bahlo/generic-list-go"
 	"github.com/pingcap/tiproxy/pkg/manager/id"
 	"github.com/pingcap/tiproxy/pkg/proxy/backend"
 	pnet "github.com/pingcap/tiproxy/pkg/proxy/net"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/cmd"
 	"go.uber.org/zap"
 )
 
-const (
-	maxPendingCommands = 100 // pending commands for each connection
-)
+// ReplayStats record the statistics during replay. All connections share one ReplayStats and update it concurrently.
+type ReplayStats struct {
+	// ReplayedCmds is the number of executed commands.
+	ReplayedCmds atomic.Uint64
+	// PendingCmds is the number of decoded but not executed commands.
+	PendingCmds atomic.Int64
+}
+
+func (s *ReplayStats) Reset() {
+	s.ReplayedCmds.Store(0)
+	s.PendingCmds.Store(0)
+}
 
 type Conn interface {
 	Run(ctx context.Context)
 	ExecuteCmd(command *cmd.Command)
+	Stop()
 }
 
 type ConnCreator func(connID uint64) Conn
 
 var _ Conn = (*conn)(nil)
 
 type conn struct {
-	cmdCh       chan *cmd.Command
-	exceptionCh chan<- Exception
-	closeCh     chan<- uint64
-	lg          *zap.Logger
-	backendConn BackendConn
-	connID      uint64 // capture ID, not replay ID
+	cmdLock         sync.Mutex
+	cmdCh           chan struct{}
+	cmdList         *glist.List[*cmd.Command]
+	exceptionCh     chan<- Exception
+	closeCh         chan<- uint64
+	lg              *zap.Logger
+	backendConn     BackendConn
+	connID          uint64 // capture ID, not replay ID
+	replayStats     *ReplayStats
+	lastPendingCmds int // last pending cmds reported to the stats
 }
 
 func NewConn(lg *zap.Logger, username, password string, backendTLSConfig *tls.Config, hsHandler backend.HandshakeHandler,
-	idMgr *id.IDManager, connID uint64, bcConfig *backend.BCConfig, exceptionCh chan<- Exception, closeCh chan<- uint64) *conn {
+	idMgr *id.IDManager, connID uint64, bcConfig *backend.BCConfig, exceptionCh chan<- Exception, closeCh chan<- uint64, replayStats *ReplayStats) *conn {
 	backendConnID := idMgr.NewID()
 	lg = lg.With(zap.Uint64("captureID", connID), zap.Uint64("replayID", backendConnID))
 	return &conn{
 		lg:          lg,
 		connID:      connID,
-		cmdCh:       make(chan *cmd.Command, maxPendingCommands),
+		cmdList:     glist.New[*cmd.Command](),
+		cmdCh:       make(chan struct{}, 1),
 		exceptionCh: exceptionCh,
 		closeCh:     closeCh,
 		backendConn: NewBackendConn(lg.Named("be"), backendConnID, hsHandler, bcConfig, backendTLSConfig, username, password),
+		replayStats: replayStats,
 	}
 }
 
@@ -59,22 +76,42 @@ func (c *conn) Run(ctx context.Context) {
 		c.exceptionCh <- NewOtherException(err, c.connID)
 		return
 	}
-	for {
+	// context is canceled when the replay is interrupted.
+	// cmdCh is closed when the replay is finished.
+	finished := false
+	for !finished {
 		select {
 		case <-ctx.Done():
-			// ctx is canceled when the replay is finished
 			return
-		case command := <-c.cmdCh:
-			if err := c.backendConn.ExecuteCmd(ctx, command.Payload); err != nil {
+		case _, ok := <-c.cmdCh:
+			if !ok {
+				finished = true
+			}
+		}
+		for ctx.Err() == nil {
+			c.cmdLock.Lock()
+			pendingCmds := c.cmdList.Len()
+			command := c.cmdList.Back()
+			if command != nil {
+				c.cmdList.Remove(command)
+			}
+			c.updatePendingCmds(pendingCmds)
+			c.cmdLock.Unlock()
+			if command == nil {
+				break
+			}
+			if err := c.backendConn.ExecuteCmd(ctx, command.Value.Payload); err != nil {
 				if pnet.IsDisconnectError(err) {
 					c.exceptionCh <- NewOtherException(err, c.connID)
+					c.lg.Debug("backend connection disconnected", zap.Error(err))
 					return
 				}
-				if c.updateCmdForExecuteStmt(command) {
-					c.exceptionCh <- NewFailException(err, command)
+				if c.updateCmdForExecuteStmt(command.Value) {
+					c.exceptionCh <- NewFailException(err, command.Value)
 				}
 			}
-			if command.Type == pnet.ComQuit {
+			c.replayStats.ReplayedCmds.Add(1)
+			if command.Value.Type == pnet.ComQuit {
 				return
 			}
 		}
@@ -99,23 +136,41 @@ func (c *conn) updateCmdForExecuteStmt(command *cmd.Command) bool {
 	return true
 }
 
-// ExecuteCmd executes a command asynchronously.
+// ExecuteCmd executes a command asynchronously by adding it to the list.
+// Adding commands should never block because it may cause cycle wait, so we don't use channels.
+// Conn A: wait for the lock held by conn B, and then its list becomes full and blocks the replay
+// Conn B: wait for next command, but the replay is blocked, so the lock won't be released
 func (c *conn) ExecuteCmd(command *cmd.Command) {
+	c.cmdLock.Lock()
+	c.cmdList.PushFront(command)
+	pendingCmds := c.cmdList.Len()
+	c.updatePendingCmds(pendingCmds)
+	c.cmdLock.Unlock()
 	select {
-	case c.cmdCh <- command:
-	case <-time.After(3 * time.Second):
-		// If the replay is slower, wait until it catches up, otherwise too many transactions are broken.
-		// But if it's blocked due to a bug, discard the command to avoid block the whole replay.
-		// If the discarded command is a COMMIT, let the next COMMIT finish the transaction.
-		select {
-		case c.exceptionCh <- NewOtherException(errors.New("too many pending commands, discard command"), c.connID):
-		default:
-			c.lg.Warn("too many pending errors, discard error")
-		}
+	case c.cmdCh <- struct{}{}:
+	default:
+	}
+}
+
+func (c *conn) Stop() {
+	close(c.cmdCh)
+}
+
+func (c *conn) updatePendingCmds(pendingCmds int) {
+	diff := pendingCmds - c.lastPendingCmds
+	c.lastPendingCmds = pendingCmds
+	if diff != 0 {
+		c.replayStats.PendingCmds.Add(int64(diff))
 	}
 }
 
 func (c *conn) close() {
+	c.cmdLock.Lock()
+	if c.cmdList.Len() > 0 {
+		c.lg.Debug("backend connection closed while there are still pending commands", zap.Int("pending_cmds", c.cmdList.Len()))
+	}
+	c.updatePendingCmds(0)
+	c.cmdLock.Unlock()
 	c.backendConn.Close()
 	c.closeCh <- c.connID
 }

pkg/sqlreplay/conn/conn_test.go

@@ -39,7 +39,7 @@ func TestConnectError(t *testing.T) {
 	var wg waitgroup.WaitGroup
 	for i, test := range tests {
 		exceptionCh, closeCh := make(chan Exception, 1), make(chan uint64, 1)
-		conn := NewConn(lg, "u1", "", nil, nil, id.NewIDManager(), 1, &backend.BCConfig{}, exceptionCh, closeCh)
+		conn := NewConn(lg, "u1", "", nil, nil, id.NewIDManager(), 1, &backend.BCConfig{}, exceptionCh, closeCh, &ReplayStats{})
 		backendConn := &mockBackendConn{connErr: test.connErr, execErr: test.execErr}
 		conn.backendConn = backendConn
 		wg.RunWithRecover(func() {
@@ -60,23 +60,41 @@ func TestExecuteCmd(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	var wg waitgroup.WaitGroup
 	exceptionCh, closeCh := make(chan Exception, 1), make(chan uint64, 1)
-	conn := NewConn(lg, "u1", "", nil, nil, id.NewIDManager(), 1, &backend.BCConfig{}, exceptionCh, closeCh)
+	stats := &ReplayStats{}
+	conn := NewConn(lg, "u1", "", nil, nil, id.NewIDManager(), 1, &backend.BCConfig{}, exceptionCh, closeCh, stats)
 	backendConn := &mockBackendConn{}
 	conn.backendConn = backendConn
 	childCtx, cancel := context.WithCancel(context.Background())
 	wg.RunWithRecover(func() {
 		conn.Run(childCtx)
 	}, nil, lg)
-	for i := 0; i < 100; i++ {
+	cmds := 1000
+	for i := 0; i < cmds; i++ {
 		conn.ExecuteCmd(&cmd.Command{ConnID: 1, Type: pnet.ComFieldList})
 	}
 	require.Eventually(t, func() bool {
-		return backendConn.cmds.Load() == 100
+		return backendConn.cmds.Load() == int32(cmds)
 	}, 3*time.Second, time.Millisecond)
+	require.EqualValues(t, cmds, stats.ReplayedCmds.Load())
+	require.EqualValues(t, 0, stats.PendingCmds.Load())
 	cancel()
 	wg.Wait()
 }
 
+func TestStopExecution(t *testing.T) {
+	lg, _ := logger.CreateLoggerForTest(t)
+	var wg waitgroup.WaitGroup
+	exceptionCh, closeCh := make(chan Exception, 1), make(chan uint64, 1)
+	conn := NewConn(lg, "u1", "", nil, nil, id.NewIDManager(), 1, &backend.BCConfig{}, exceptionCh, closeCh, &ReplayStats{})
+	conn.backendConn = &mockBackendConn{}
+	wg.RunWithRecover(func() {
+		conn.Run(context.Background())
+	}, nil, lg)
+	conn.ExecuteCmd(&cmd.Command{ConnID: 1, Type: pnet.ComFieldList})
+	conn.Stop()
+	wg.Wait()
+}
+
 func TestExecuteError(t *testing.T) {
 	tests := []struct {
 		prepare   func(*mockBackendConn) []byte
@@ -106,7 +124,7 @@ func TestExecuteError(t *testing.T) {
 	lg, _ := logger.CreateLoggerForTest(t)
 	var wg waitgroup.WaitGroup
 	exceptionCh, closeCh := make(chan Exception, 1), make(chan uint64, 1)
-	conn := NewConn(lg, "u1", "", nil, nil, id.NewIDManager(), 1, &backend.BCConfig{}, exceptionCh, closeCh)
+	conn := NewConn(lg, "u1", "", nil, nil, id.NewIDManager(), 1, &backend.BCConfig{}, exceptionCh, closeCh, &ReplayStats{})
 	backendConn := &mockBackendConn{execErr: errors.New("mock error")}
 	conn.backendConn = backendConn
 	childCtx, cancel := context.WithCancel(context.Background())