replay: use start-time config to start replay later, and add dryrun argument for api. (#1003)

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

Author: YangKeao

cmd/replayer/main.go

@@ -59,6 +59,7 @@ func main() {
 	outputPath := rootCmd.PersistentFlags().String("output-path", "", "the file path to store replayed sql. Empty indicates do not output replayed sql.")
 	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.")
 
 	rootCmd.RunE = func(cmd *cobra.Command, _ []string) error {
 		// set up general managers
@@ -134,7 +135,7 @@ func main() {
 				Password:           *password,
 				Format:             *format,
 				ReadOnly:           *readonly,
-				StartTime:          time.Now(),
+				StartTime:          *startTime,
 				CommandStartTime:   *cmdStartTime,
 				CommandEndTime:     *cmdEndTime,
 				IgnoreErrs:         *ignoreErrs,

pkg/server/api/traffic.go

@@ -155,6 +155,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")
 	h.lg.Info("request: traffic replay", zap.Any("cfg", cfg))
 
 	if err := h.mgr.ReplayJobMgr.StartReplay(cfg); err != nil {

pkg/sqlreplay/replay/replay.go

@@ -176,8 +176,6 @@ func (cfg *ReplayConfig) Validate() ([]storage.ExternalStorage, error) {
 	now := time.Now()
 	if cfg.StartTime.IsZero() {
 		return storages, errors.New("start time is not specified")
-	} else if now.Add(time.Minute).Before(cfg.StartTime) {
-		return storages, errors.New("start time should not be in the future")
 	} else if cfg.StartTime.Add(time.Minute).Before(now) {
 		return storages, errors.New("start time should not be in the past")
 	}
@@ -269,7 +267,7 @@ type replay struct {
 	exceptionCh      chan conn.Exception
 	closeConnCh      chan uint64
 	execInfoCh       chan conn.ExecInfo
-	gracefulStop     atomic.Bool
+	gracefulStopCh   chan struct{}
 	wg               waitgroup.WaitGroup
 	cancel           context.CancelFunc
 	connCreator      conn.ConnCreator
@@ -309,7 +307,7 @@ func (r *replay) Start(cfg ReplayConfig, backendTLSConfig *tls.Config, hsHandler
 	r.cfg = cfg
 	r.storages = storages
 	r.meta = *r.readMeta()
-	r.gracefulStop.Store(false)
+	r.gracefulStopCh = make(chan struct{})
 	r.startTime = cfg.StartTime
 	r.endTime = time.Time{}
 	r.progress = 0
@@ -395,7 +393,45 @@ func (r *replay) Start(cfg ReplayConfig, backendTLSConfig *tls.Config, hsHandler
 	return nil
 }
 
+// waitUntilStartTime waits until the configured start time and returns whether the replay has stopped
+func (r *replay) waitUntilStartTime(ctx context.Context) bool {
+	waitDuration := time.Until(r.cfg.StartTime)
+	if waitDuration > 0 {
+		r.lg.Info("wait until the specified time to start replaying", zap.Time("until", r.cfg.StartTime),
+			zap.Duration("duration", waitDuration))
+		select {
+		case <-ctx.Done():
+			return true
+		case <-time.After(waitDuration):
+			return false
+		case <-r.gracefulStopCh:
+			return true
+		}
+	}
+
+	return false
+}
+
+// isGracefulStopped returns whether the replay is gracefully stopped.
+func (r *replay) isGracefulStopped() bool {
+	select {
+	case _, ok := <-r.gracefulStopCh:
+		if !ok {
+			return true
+		} else {
+			return false
+		}
+	default:
+		return false
+	}
+}
+
 func (r *replay) readCommands(ctx context.Context) {
+	if r.waitUntilStartTime(ctx) {
+		r.stop(nil)
+		return
+	}
+
 	var decoder decoder
 	var err error
 
@@ -429,7 +465,7 @@ func (r *replay) readCommands(ctx context.Context) {
 	connCount := 0                      // alive connection count
 	maxPendingCmds := int64(0)
 	extraWaitTime := time.Duration(0)
-	for ctx.Err() == nil && !r.gracefulStop.Load() {
+	for ctx.Err() == nil && !r.isGracefulStopped() {
 		for hasCloseEvent := true; hasCloseEvent; {
 			select {
 			case id := <-r.closeConnCh:
@@ -501,6 +537,7 @@ func (r *replay) readCommands(ctx context.Context) {
 				select {
 				case <-ctx.Done():
 				case <-time.After(expectedInterval):
+				case <-r.gracefulStopCh:
 				}
 			}
 		}
@@ -513,7 +550,7 @@ func (r *replay) readCommands(ctx context.Context) {
 		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.Bool("graceful_stop", r.gracefulStop.Load()),
+		zap.Bool("graceful_stop", r.isGracefulStopped()),
 		zap.Error(err),
 		zap.NamedError("ctx_err", ctx.Err()))
 
@@ -980,6 +1017,10 @@ func (r *replay) Wait() {
 	r.wg.Wait()
 }
 
+func (r *replay) gracefulStop() {
+	close(r.gracefulStopCh)
+}
+
 func (r *replay) Stop(err error, graceful bool) {
 	r.Lock()
 	// already stopped
@@ -989,7 +1030,7 @@ func (r *replay) Stop(err error, graceful bool) {
 	}
 	r.err = err
 	if graceful {
-		r.gracefulStop.Store(true)
+		r.gracefulStop()
 	} else if r.cancel != nil {
 		r.cancel()
 		r.cancel = nil

pkg/sqlreplay/replay/replay_test.go

@@ -1102,3 +1102,31 @@ func TestDynamicInputCoverAllDirectories(t *testing.T) {
 	}
 	require.Equal(t, expectCommandCount, actualCommandCount)
 }
+
+func TestWaitUntil(t *testing.T) {
+	replay := NewReplay(zap.NewNop(), id.NewIDManager())
+	defer replay.Close()
+	loader := newMockNormalLoader()
+	startReplayTime := time.Now()
+	waitUntil := startReplayTime.Add(500 * time.Millisecond)
+	cfg := ReplayConfig{
+		DryRun:          true,
+		Input:           t.TempDir(),
+		StartTime:       waitUntil,
+		readers:         []cmd.LineReader{loader},
+		PSCloseStrategy: cmd.PSCloseStrategyDirected,
+	}
+
+	now := time.Now()
+	command := newMockCommand(1)
+	command.StartTs = now
+	loader.writeCommand(command, cmd.FormatAuditLogPlugin)
+	require.NoError(t, replay.Start(cfg, nil, nil, &backend.BCConfig{}))
+	loader.Close()
+
+	require.Eventually(t, func() bool {
+		return replay.replayStats.CurCmdTs.Load() != 0
+	}, time.Second*5, time.Millisecond)
+	actualStartTime := time.Now()
+	require.GreaterOrEqual(t, actualStartTime.Sub(startReplayTime), 500*time.Millisecond)
+}