sqlreplay: fix replay quits immediately after graceful pause (#995)

Author: djshow832

pkg/sqlreplay/manager/manager.go

@@ -164,7 +164,6 @@ func (jm *jobManager) StartReplay(cfg replay.ReplayConfig) error {
 		jm.lg.Warn("start replay failed", zap.String("job", newJob.String()), zap.Error(err))
 		return errors.Wrapf(err, "start replay failed")
 	}
-	jm.lg.Info("start replay", zap.String("job", newJob.String()))
 	jm.addToHistory(newJob)
 	return nil
 }

pkg/sqlreplay/replay/replay.go

@@ -86,7 +86,7 @@ type ReplayConfig struct {
 	Format   string
 	Input    string
 	Username string
-	Password string
+	Password string `json:"-"`
 	KeyFile  string
 	// It's specified when executing with the statement `TRAFFIC REPLAY` so that all TiProxy instances
 	// use the same start time and the time acts as the job ID.
@@ -303,11 +303,13 @@ func (r *replay) Start(cfg ReplayConfig, backendTLSConfig *tls.Config, hsHandler
 		return err
 	}
 
+	r.lg.Info("start replay", zap.Any("config", cfg))
 	r.Lock()
 	defer r.Unlock()
 	r.cfg = cfg
 	r.storages = storages
 	r.meta = *r.readMeta()
+	r.gracefulStop.Store(false)
 	r.startTime = cfg.StartTime
 	r.endTime = time.Time{}
 	r.progress = 0
@@ -449,6 +451,7 @@ func (r *replay) readCommands(ctx context.Context) {
 				err = nil
 				continue
 			} else {
+				r.lg.Info("decode failed, stop", zap.Error(err))
 				break
 			}
 		}
@@ -509,7 +512,9 @@ func (r *replay) readCommands(ctx context.Context) {
 		zap.Duration("extra_wait_time", extraWaitTime),
 		zap.Int("alive_conns", connCount),
 		zap.Time("last_cmd_start_ts", time.Unix(0, r.replayStats.CurCmdTs.Load())),
-		zap.Time("last_cmd_end_ts", time.Unix(0, r.replayStats.CurCmdEndTs.Load())))
+		zap.Time("last_cmd_end_ts", time.Unix(0, r.replayStats.CurCmdEndTs.Load())),
+		zap.NamedError("ctx_err", ctx.Err()),
+		zap.Bool("graceful_stop", r.gracefulStop.Load()))
 
 	// Notify the connections that the commands are finished.
 	for _, conn := range conns {
@@ -928,15 +933,9 @@ func (r *replay) stop(err error) {
 	}
 	commonFields := r.commonFields()
 	fields := append(commonFields, []zap.Field{
-		zap.Time("start_time", r.startTime),
-		zap.Time("end_time", r.endTime),
-		zap.Time("command_start_time", r.cfg.CommandStartTime),
-		zap.Time("command_end_time", r.cfg.CommandEndTime),
-		zap.String("format", r.cfg.Format),
-		zap.String("username", r.cfg.Username),
-		zap.Bool("ignore_errs", r.cfg.IgnoreErrs),
-		zap.Float64("speed", r.cfg.Speed),
-		zap.Bool("read_only", r.cfg.ReadOnly),
+		zap.Time("replay_start_time", r.startTime),
+		zap.Time("replay_end_time", r.endTime),
+		zap.Any("config", r.cfg),
 	}...)
 	if r.meta.Cmds > 0 {
 		r.progress = float64(decodedCmds) / float64(r.meta.Cmds)

pkg/sqlreplay/replay/replay_test.go

@@ -431,43 +431,46 @@ func TestGracefulStop(t *testing.T) {
 	replay := NewReplay(zap.NewNop(), id.NewIDManager())
 	defer replay.Close()
 
-	i := 0
-	loader := &customizedReader{
-		getCmd: func() *cmd.Command {
-			j := rand.Uint64N(100) + 1
-			command := newMockCommand(j)
-			i++
-			command.StartTs = time.Unix(0, int64(i)*int64(time.Microsecond))
-			return command
-		},
-	}
+	// Test 2 rounds to test that the graceful flag will be reset before each round.
+	for n := 0; n < 2; n++ {
+		i := 0
+		loader := &customizedReader{
+			getCmd: func() *cmd.Command {
+				j := rand.Uint64N(100) + 1
+				command := newMockCommand(j)
+				i++
+				command.StartTs = time.Unix(0, int64(i)*int64(time.Microsecond))
+				return command
+			},
+		}
 
-	cfg := ReplayConfig{
-		Input:     t.TempDir(),
-		Username:  "u1",
-		StartTime: time.Now(),
-		readers:   []cmd.LineReader{loader},
-		connCreator: func(connID uint64, _ uint64) conn.Conn {
-			return &mockDelayConn{
-				stats:   &replay.replayStats,
-				closeCh: replay.closeConnCh,
-				connID:  connID,
-			}
-		},
-		report:          newMockReport(replay.exceptionCh),
-		PSCloseStrategy: cmd.PSCloseStrategyDirected,
-	}
-	require.NoError(t, replay.Start(cfg, nil, nil, &backend.BCConfig{}))
+		cfg := ReplayConfig{
+			Input:     t.TempDir(),
+			Username:  "u1",
+			StartTime: time.Now(),
+			readers:   []cmd.LineReader{loader},
+			connCreator: func(connID uint64, _ uint64) conn.Conn {
+				return &mockDelayConn{
+					stats:   &replay.replayStats,
+					closeCh: replay.closeConnCh,
+					connID:  connID,
+				}
+			},
+			report:          newMockReport(replay.exceptionCh),
+			PSCloseStrategy: cmd.PSCloseStrategyDirected,
+		}
+		require.NoError(t, replay.Start(cfg, nil, nil, &backend.BCConfig{}))
 
-	time.Sleep(2 * time.Second)
-	replay.Stop(errors.New("graceful stop"), true)
-	// check that all the pending commands are replayed
-	curCmdTs := replay.replayStats.CurCmdTs.Load()
-	require.EqualValues(t, 0, replay.replayStats.PendingCmds.Load())
-	require.EqualValues(t, curCmdTs, int64(replay.replayStats.ReplayedCmds.Load())*int64(time.Microsecond))
-	_, _, lastTs, _, _, err := replay.Progress()
-	require.ErrorContains(t, err, "graceful stop")
-	require.Equal(t, curCmdTs, lastTs.UnixNano())
+		time.Sleep(2 * time.Second)
+		replay.Stop(errors.New("graceful stop"), true)
+		// check that all the pending commands are replayed
+		curCmdTs := replay.replayStats.CurCmdTs.Load()
+		require.EqualValues(t, 0, replay.replayStats.PendingCmds.Load())
+		require.EqualValues(t, curCmdTs, int64(replay.replayStats.ReplayedCmds.Load())*int64(time.Microsecond))
+		_, _, lastTs, _, _, err := replay.Progress()
+		require.ErrorContains(t, err, "graceful stop")
+		require.Equal(t, curCmdTs, lastTs.UnixNano())
+	}
 }
 
 func BenchmarkMultiBufferedDecoder(b *testing.B) {