support filtering commands with RETRY field (#1011)

Signed-off-by: Yang Keao <yangkeao@chunibyo.icu>

Author: YangKeao

cmd/replayer/main.go

@@ -60,6 +60,7 @@ func main() {
 	serviceMode := rootCmd.PersistentFlags().Bool("service-mode", false, "run replayer in service mode")
 	logLevel := rootCmd.PersistentFlags().String("log-level", "info", "the log level: debug, info, warn, error, dpanic, panic, fatal")
 	startTime := rootCmd.PersistentFlags().Time("start-time", time.Now(), []string{time.RFC3339, time.RFC3339Nano}, "the time to start the replay. Format is RFC3339. Default is the current time.")
+	filterCommandWithRetry := rootCmd.PersistentFlags().Bool("filter-command-with-retry", false, "filter out commands that are retries according to the audit log.")
 
 	rootCmd.RunE = func(cmd *cobra.Command, _ []string) error {
 		// set up general managers
@@ -129,24 +130,25 @@ func main() {
 		} else {
 			// start replay
 			replayCfg := replay.ReplayConfig{
-				Input:              *input,
-				Speed:              *speed,
-				Username:           *username,
-				Password:           *password,
-				Format:             *format,
-				ReadOnly:           *readonly,
-				StartTime:          *startTime,
-				CommandStartTime:   *cmdStartTime,
-				CommandEndTime:     *cmdEndTime,
-				IgnoreErrs:         *ignoreErrs,
-				BufSize:            *bufSize,
-				PSCloseStrategy:    replaycmd.PSCloseStrategy(*psCloseStrategy),
-				DryRun:             *dryRun,
-				CheckPointFilePath: *checkPointFilePath,
-				DynamicInput:       *dynamicInput,
-				ReplayerCount:      *replayerCount,
-				ReplayerIndex:      *replayerIndex,
-				OutputPath:         *outputPath,
+				Input:                  *input,
+				Speed:                  *speed,
+				Username:               *username,
+				Password:               *password,
+				Format:                 *format,
+				ReadOnly:               *readonly,
+				StartTime:              *startTime,
+				CommandStartTime:       *cmdStartTime,
+				CommandEndTime:         *cmdEndTime,
+				IgnoreErrs:             *ignoreErrs,
+				BufSize:                *bufSize,
+				PSCloseStrategy:        replaycmd.PSCloseStrategy(*psCloseStrategy),
+				DryRun:                 *dryRun,
+				CheckPointFilePath:     *checkPointFilePath,
+				DynamicInput:           *dynamicInput,
+				ReplayerCount:          *replayerCount,
+				ReplayerIndex:          *replayerIndex,
+				OutputPath:             *outputPath,
+				FilterCommandWithRetry: *filterCommandWithRetry,
 			}
 			if err := r.StartReplay(replayCfg); err != nil {
 				cancel()

pkg/server/api/traffic.go

@@ -156,6 +156,7 @@ func (h *Server) TrafficReplay(c *gin.Context) {
 	cfg.OutputPath = c.PostForm("outputpath")
 	cfg.Addr = c.PostForm("addr")
 	cfg.DryRun = strings.EqualFold(c.PostForm("dryrun"), "true")
+	cfg.FilterCommandWithRetry = strings.EqualFold(c.PostForm("filtercommandwithretry"), "true")
 	h.lg.Info("request: traffic replay", zap.Any("cfg", cfg))
 
 	if err := h.mgr.ReplayJobMgr.StartReplay(cfg); err != nil {

pkg/sqlreplay/cmd/audit_log_plugin.go

@@ -29,6 +29,7 @@ const (
 	auditPluginKeyEvent          = "EVENT"
 	auditPluginKeyCostTime       = "COST_TIME"
 	auditPluginKeyPreparedStmtID = "PREPARED_STMT_ID"
+	auditPluginKeyRetry          = "RETRY"
 
 	auditPluginClassGeneral     = "GENERAL"
 	auditPluginClassTableAccess = "TABLE_ACCESS"
@@ -101,11 +102,12 @@ type AuditLogPluginDecoder struct {
 	commandStartTime time.Time
 	commandEndTime   time.Time
 	// pendingCmds contains the commands that has not been returned yet.
-	pendingCmds     []*Command
-	psCloseStrategy PSCloseStrategy
-	idAllocator     *ConnIDAllocator
-	dedup           *DeDup
-	lg              *zap.Logger
+	pendingCmds            []*Command
+	psCloseStrategy        PSCloseStrategy
+	filterCommandWithRetry bool
+	idAllocator            *ConnIDAllocator
+	dedup                  *DeDup
+	lg                     *zap.Logger
 }
 
 // ConnIDAllocator allocates connection IDs for new connections.
@@ -446,9 +448,18 @@ func (decoder *AuditLogPluginDecoder) parseGeneralEvent(kvs map[string]string, s
 		if err != nil {
 			return nil, errors.Wrapf(err, "unquote sql failed: %s", kvs[auditPluginKeySQL])
 		}
-		// deduplicate DML and SELECT FOR UPDATE
-		if decoder.isDuplicatedWrite(connInfo.lastCmd, kvs, cmdStr, sql, startTs, endTs) {
-			return nil, nil
+		if decoder.filterCommandWithRetry {
+			if retryStr, ok := kvs[auditPluginKeyRetry]; ok {
+				if retryStr == "true" {
+					// skip the retried command
+					return nil, nil
+				}
+			}
+		} else {
+			// deduplicate DML and SELECT FOR UPDATE
+			if decoder.isDuplicatedWrite(connInfo.lastCmd, kvs, cmdStr, sql, startTs, endTs) {
+				return nil, nil
+			}
 		}
 	}
 
@@ -641,3 +652,8 @@ func (decoder *AuditLogPluginDecoder) isDuplicatedWrite(lastCmd *Command, kvs ma
 	decoder.dedup.Unlock()
 	return true
 }
+
+// EnableFilterCommandWithRetry enables filtering out commands that are retries according to the audit log.
+func (decoder *AuditLogPluginDecoder) EnableFilterCommandWithRetry() {
+	decoder.filterCommandWithRetry = true
+}

pkg/sqlreplay/cmd/audit_log_plugin_test.go

@@ -1657,3 +1657,65 @@ func decodeCmds(decoder CmdDecoder, reader LineReader) ([]*Command, error) {
 	}
 	return cmds, err
 }
+
+func TestFilterOutCommandsWithRetry(t *testing.T) {
+	tests := []struct {
+		name     string
+		lines    string
+		enable   bool
+		cmdCount int
+	}{
+		{
+			name:     "filter out retried query command",
+			lines:    `[2025/09/18 17:48:20.614 +08:10] [INFO] [logger.go:77] [ID=17581889006155] [TIMESTAMP=2025/09/18 17:48:20.614 +08:10] [EVENT_CLASS=GENERAL] [EVENT_SUBCLASS=] [STATUS_CODE=0] [COST_TIME=1000] [HOST=127.0.0.1] [CLIENT_IP=127.0.0.1] [USER=root] [DATABASES="[]"] [TABLES="[]"] [SQL_TEXT="SELECT * FROM test"] [ROWS=1] [CONNECTION_ID=1001] [CLIENT_PORT=50112] [PID=542193] [COMMAND=Query] [SQL_STATEMENTS=Select] [EXECUTE_PARAMS="[]"] [CURRENT_DB=test] [RETRY=true] [EVENT=COMPLETED]`,
+			enable:   true,
+			cmdCount: 0,
+		},
+		{
+			name:     "do not filter when this function is disabled",
+			lines:    `[2025/09/18 17:48:20.614 +08:10] [INFO] [logger.go:77] [ID=17581889006155] [TIMESTAMP=2025/09/18 17:48:20.614 +08:10] [EVENT_CLASS=GENERAL] [EVENT_SUBCLASS=] [STATUS_CODE=0] [COST_TIME=1000] [HOST=127.0.0.1] [CLIENT_IP=127.0.0.1] [USER=root] [DATABASES="[]"] [TABLES="[]"] [SQL_TEXT="SELECT * FROM test"] [ROWS=1] [CONNECTION_ID=1001] [CLIENT_PORT=50112] [PID=542193] [COMMAND=Query] [SQL_STATEMENTS=Select] [EXECUTE_PARAMS="[]"] [CURRENT_DB=test] [RETRY=true] [EVENT=COMPLETED]`,
+			enable:   false,
+			cmdCount: 1,
+		},
+		{
+			name:     "do not filter non-retried commands",
+			lines:    `[2025/09/18 17:48:20.614 +08:10] [INFO] [logger.go:77] [ID=17581889006155] [TIMESTAMP=2025/09/18 17:48:20.614 +08:10] [EVENT_CLASS=GENERAL] [EVENT_SUBCLASS=] [STATUS_CODE=0] [COST_TIME=1000] [HOST=127.0.0.1] [CLIENT_IP=127.0.0.1] [USER=root] [DATABASES="[]"] [TABLES="[]"] [SQL_TEXT="SELECT * FROM test"] [ROWS=1] [CONNECTION_ID=1001] [CLIENT_PORT=50112] [PID=542193] [COMMAND=Query] [SQL_STATEMENTS=Select] [EXECUTE_PARAMS="[]"] [CURRENT_DB=test] [EVENT=COMPLETED]`,
+			enable:   true,
+			cmdCount: 1,
+		},
+		{
+			name: "mixed retried and non-retried commands",
+			lines: `[2025/09/18 17:48:20.614 +08:10] [INFO] [logger.go:77] [ID=17581889006155] [TIMESTAMP=2025/09/18 17:48:20.614 +08:10] [EVENT_CLASS=GENERAL] [EVENT_SUBCLASS=] [STATUS_CODE=0] [COST_TIME=1000] [HOST=127.0.0.1] [CLIENT_IP=127.0.0.1] [USER=root] [DATABASES="[]"] [TABLES="[]"] [SQL_TEXT="SELECT 1"] [ROWS=1] [CONNECTION_ID=1001] [CLIENT_PORT=50112] [PID=542193] [COMMAND=Query] [SQL_STATEMENTS=Select] [EXECUTE_PARAMS="[]"] [CURRENT_DB=test] [EVENT=COMPLETED]
+[2025/09/18 17:48:20.714 +08:10] [INFO] [logger.go:77] [ID=17581889006156] [TIMESTAMP=2025/09/18 17:48:20.714 +08:10] [EVENT_CLASS=GENERAL] [EVENT_SUBCLASS=] [STATUS_CODE=0] [COST_TIME=1000] [HOST=127.0.0.1] [CLIENT_IP=127.0.0.1] [USER=root] [DATABASES="[]"] [TABLES="[]"] [SQL_TEXT="SELECT 2"] [ROWS=1] [CONNECTION_ID=1001] [CLIENT_PORT=50112] [PID=542193] [COMMAND=Query] [SQL_STATEMENTS=Select] [EXECUTE_PARAMS="[]"] [CURRENT_DB=test] [RETRY=true] [EVENT=COMPLETED]
+[2025/09/18 17:48:20.814 +08:10] [INFO] [logger.go:77] [ID=17581889006157] [TIMESTAMP=2025/09/18 17:48:20.814 +08:10] [EVENT_CLASS=GENERAL] [EVENT_SUBCLASS=] [STATUS_CODE=0] [COST_TIME=1000] [HOST=127.0.0.1] [CLIENT_IP=127.0.0.1] [USER=root] [DATABASES="[]"] [TABLES="[]"] [SQL_TEXT="SELECT 3"] [ROWS=1] [CONNECTION_ID=1001] [CLIENT_PORT=50112] [PID=542193] [COMMAND=Query] [SQL_STATEMENTS=Select] [EXECUTE_PARAMS="[]"] [CURRENT_DB=test] [EVENT=COMPLETED]`,
+			enable:   true,
+			cmdCount: 2,
+		},
+		{
+			name:     "retry flag with other value should not be filtered",
+			lines:    `[2025/09/18 17:48:20.614 +08:10] [INFO] [logger.go:77] [ID=17581889006155] [TIMESTAMP=2025/09/18 17:48:20.614 +08:10] [EVENT_CLASS=GENERAL] [EVENT_SUBCLASS=] [STATUS_CODE=0] [COST_TIME=1000] [HOST=127.0.0.1] [CLIENT_IP=127.0.0.1] [USER=root] [DATABASES="[]"] [TABLES="[]"] [SQL_TEXT="SELECT * FROM test"] [ROWS=1] [CONNECTION_ID=1001] [CLIENT_PORT=50112] [PID=542193] [COMMAND=Query] [SQL_STATEMENTS=Select] [EXECUTE_PARAMS="[]"] [CURRENT_DB=test] [RETRY=false] [EVENT=COMPLETED]`,
+			enable:   true,
+			cmdCount: 1,
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			decoder := NewAuditLogPluginDecoder(NewDeDup(), zap.NewNop())
+			if tt.enable {
+				decoder.EnableFilterCommandWithRetry()
+			}
+			mr := mockReader{data: append([]byte(tt.lines), '\n')}
+			cmds := make([]*Command, 0, 4)
+			for {
+				cmd, err := decoder.Decode(&mr)
+				if cmd == nil {
+					require.ErrorIs(t, err, io.EOF)
+					break
+				}
+				cmds = append(cmds, cmd)
+			}
+			require.Equal(t, tt.cmdCount, len(cmds))
+		})
+	}
+}

pkg/sqlreplay/replay/replay.go

@@ -122,6 +122,8 @@ type ReplayConfig struct {
 	OutputPath string
 	// Addr is the downstream address to connect to
 	Addr string
+	// FilterCommandWithRetry indicates whether to filter out commands that are retries according to the audit log.
+	FilterCommandWithRetry bool
 	// the following fields are for testing
 	readers           []cmd.LineReader
 	report            report.Report
@@ -616,6 +618,10 @@ func (r *replay) constructDecoderForReader(ctx context.Context, reader cmd.LineR
 		auditLogDecoder.SetPSCloseStrategy(r.cfg.PSCloseStrategy)
 		auditLogDecoder.SetIDAllocator(idAllocator)
 		auditLogDecoder.SetCommandEndTime(r.cfg.CommandEndTime)
+
+		if r.cfg.FilterCommandWithRetry {
+			auditLogDecoder.EnableFilterCommandWithRetry()
+		}
 	}
 
 	var decoder decoder

pkg/sqlreplay/replay/replay_test.go

@@ -744,7 +744,7 @@ func TestDynamicInput(t *testing.T) {
 		}
 	}()
 
-	dirWatcherInterval := 10 * time.Millisecond
+	dirWatcherInterval := 100 * time.Millisecond
 	store.SetDirWatcherPollIntervalForTest(dirWatcherInterval)
 
 	auditLog := `[2025/09/08 21:16:29.585 +08:00] [INFO] [logger.go:77] [ID=17573373891] [TIMESTAMP=2025/09/06 16:16:29.585 +08:10] [EVENT_CLASS=GENERAL] [EVENT_SUBCLASS=] [STATUS_CODE=0] [COST_TIME=1057.834] [HOST=127.0.0.1] [CLIENT_IP=127.0.0.1] [USER=root] [DATABASES="[]"] [TABLES="[]"] [SQL_TEXT="select 1"] [ROWS=0] [CONNECTION_ID=3695181836] [CLIENT_PORT=52611] [PID=89967] [COMMAND=Query] [SQL_STATEMENTS=Set] [EXECUTE_PARAMS="[]"] [CURRENT_DB=] [EVENT=COMPLETED]