replay: add lock to avoid starting replay task multiple times (#992)

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

Author: YangKeao

cmd/replayer/main.go

@@ -58,13 +58,14 @@ func main() {
 	replayerIndex := rootCmd.PersistentFlags().Uint64("replayer-index", 0, "the index of this replayer instance. Used only when dynamic-input is enabled.")
 	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")
 
 	rootCmd.RunE = func(cmd *cobra.Command, _ []string) error {
 		// set up general managers
 		cfg := &config.Config{
 			Log: config.Log{
 				LogOnline: config.LogOnline{
-					Level:   "info",
+					Level:   *logLevel,
 					LogFile: config.LogFile{Filename: *logFile},
 				},
 			},

pkg/sqlreplay/manager/manager.go

@@ -6,6 +6,7 @@ package manager
 import (
 	"crypto/tls"
 	"encoding/json"
+	"sync"
 	"time"
 
 	"github.com/pingcap/tiproxy/lib/config"
@@ -46,6 +47,8 @@ type JobManager interface {
 var _ JobManager = (*jobManager)(nil)
 
 type jobManager struct {
+	mu sync.Mutex
+
 	jobHistory         []Job
 	capture            capture.Capture
 	replay             replay.Replay
@@ -100,6 +103,9 @@ func (jm *jobManager) runningJob() Job {
 }
 
 func (jm *jobManager) StartCapture(cfg capture.CaptureConfig) error {
+	jm.mu.Lock()
+	defer jm.mu.Unlock()
+
 	running := jm.runningJob()
 	if running != nil {
 		return errors.Errorf("a job is running: %s", running.String())
@@ -121,6 +127,9 @@ func (jm *jobManager) StartCapture(cfg capture.CaptureConfig) error {
 }
 
 func (jm *jobManager) StartReplay(cfg replay.ReplayConfig) error {
+	jm.mu.Lock()
+	defer jm.mu.Unlock()
+
 	running := jm.runningJob()
 	if running != nil {
 		return errors.Errorf("a job is running: %s", running.String())
@@ -174,6 +183,9 @@ func (jm *jobManager) GetCapture() capture.Capture {
 }
 
 func (jm *jobManager) Jobs() string {
+	jm.mu.Lock()
+	defer jm.mu.Unlock()
+
 	jm.updateProgress()
 	b, err := json.MarshalIndent(jm.jobHistory, "", "  ")
 	if err != nil {
@@ -182,6 +194,8 @@ func (jm *jobManager) Jobs() string {
 	return hack.String(b)
 }
 
+// Wait waits for the running job to finish.
+// As `Wait` is a blocking call, it'll not acquire the jobManager lock. For now it's only used in standalone player mode.
 func (jm *jobManager) Wait() {
 	job := jm.runningJob()
 	if job == nil {
@@ -196,6 +210,9 @@ func (jm *jobManager) Wait() {
 }
 
 func (jm *jobManager) Stop(cfg CancelConfig) string {
+	jm.mu.Lock()
+	defer jm.mu.Unlock()
+
 	job := jm.runningJob()
 	if job == nil {
 		return "no job running"
@@ -214,6 +231,9 @@ func (jm *jobManager) Stop(cfg CancelConfig) string {
 }
 
 func (jm *jobManager) Close() {
+	jm.mu.Lock()
+	defer jm.mu.Unlock()
+
 	if jm.capture != nil {
 		jm.capture.Close()
 	}

pkg/sqlreplay/manager/manager_test.go

@@ -4,13 +4,15 @@
 package manager
 
 import (
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/pingcap/tiproxy/lib/config"
 	"github.com/pingcap/tiproxy/lib/util/errors"
 	"github.com/pingcap/tiproxy/pkg/manager/id"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/capture"
+	"github.com/pingcap/tiproxy/pkg/sqlreplay/cmd"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/replay"
 	"github.com/stretchr/testify/require"
 	"go.uber.org/zap"
@@ -198,3 +200,35 @@ func TestAllowAddrOnlyForStandaloneService(t *testing.T) {
 	})
 	require.NoError(t, err)
 }
+
+func TestAvoidConcurrentReplay(t *testing.T) {
+	mgr := NewJobManager(zap.NewNop(), &config.Config{}, &mockCertMgr{}, id.NewIDManager(), nil, true)
+	defer mgr.Close()
+	mgr.replay = &mockReplay{}
+
+	successCount := 0
+	attempts := 1000
+
+	wg := &sync.WaitGroup{}
+	for i := 0; i < attempts; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+
+			err := mgr.StartReplay(replay.ReplayConfig{
+				Addr:            "127.0.0.1:10000",
+				Input:           "/tmp/traffic",
+				Username:        "root",
+				StartTime:       time.Now(),
+				PSCloseStrategy: cmd.PSCloseStrategyAlways,
+			})
+			if err == nil {
+				successCount++
+			} else {
+				require.ErrorContains(t, err, "a job is running")
+			}
+		}()
+	}
+	wg.Wait()
+	require.Equal(t, 1, successCount)
+}

pkg/sqlreplay/replay/replay.go

@@ -238,7 +238,14 @@ func (cfg *ReplayConfig) LoadFromCheckpoint() error {
 	decoder := json.NewDecoder(file)
 	var state replayCheckpoint
 	if err := decoder.Decode(&state); err != nil {
-		return errors.Wrapf(err, "failed to decode checkpoint file %s", cfg.CheckPointFilePath)
+		// Allow empty file.
+		stat, err2 := file.Stat()
+		if err2 != nil {
+			return errors.Wrapf(err, "failed to stat checkpoint file %s", cfg.CheckPointFilePath)
+		}
+		if stat.Size() != 0 {
+			return errors.Wrapf(err, "failed to decode checkpoint file %s", cfg.CheckPointFilePath)
+		}
 	}
 
 	if state.CurCmdTs > 0 {

pkg/sqlreplay/replay/replay_test.go

@@ -570,7 +570,7 @@ func TestLoadFromCheckpoint(t *testing.T) {
 		{
 			checkpointData: "",
 			setupFile:      true,
-			expectedError:  true,
+			expectedError:  false,
 		},
 	}