sqlreplay: do not write traffic files if meta exists (#721)

Author: djshow832

pkg/sqlreplay/capture/capture.go

@@ -72,6 +72,7 @@ func (cfg *CaptureConfig) Validate() error {
 		if !st.IsDir() {
 			return errors.New("output should be a directory")
 		}
+		err = store.PreCheckMeta(cfg.Output)
 	} else if os.IsNotExist(err) {
 		err = os.MkdirAll(cfg.Output, 0755)
 	}

pkg/sqlreplay/capture/capture_test.go

@@ -21,14 +21,15 @@ import (
 )
 
 func TestStartAndStop(t *testing.T) {
+	dir := t.TempDir()
 	cpt := NewCapture(zap.NewNop())
 	defer cpt.Close()
 
 	packet := append([]byte{pnet.ComQuery.Byte()}, []byte("select 1")...)
 	cpt.Capture(packet, time.Now(), 100, mockInitSession)
 	writer := newMockWriter(store.WriterCfg{})
 	cfg := CaptureConfig{
-		Output:    t.TempDir(),
+		Output:    dir,
 		Duration:  10 * time.Second,
 		cmdLogger: writer,
 	}
@@ -48,25 +49,28 @@ func TestStartAndStop(t *testing.T) {
 	require.Equal(t, len(data), len(writer.getData()))
 
 	// start capture again
+	removeMeta(dir)
 	require.NoError(t, cpt.Start(cfg))
 	cpt.Capture(packet, time.Now(), 100, mockInitSession)
 	cpt.Stop(nil)
 	require.Greater(t, len(writer.getData()), len(data))
 
 	// duplicated start and stop
+	removeMeta(dir)
 	require.NoError(t, cpt.Start(cfg))
 	require.Error(t, cpt.Start(cfg))
 	cpt.Stop(nil)
 	cpt.Stop(nil)
 }
 
 func TestConcurrency(t *testing.T) {
+	dir := t.TempDir()
 	cpt := NewCapture(zap.NewNop())
 	defer cpt.Close()
 
 	writer := newMockWriter(store.WriterCfg{})
 	cfg := CaptureConfig{
-		Output:         t.TempDir(),
+		Output:         dir,
 		Duration:       10 * time.Second,
 		bufferCap:      12 * 1024,
 		flushThreshold: 8 * 1024,
@@ -107,6 +111,7 @@ func TestConcurrency(t *testing.T) {
 					cpt.Stop(nil)
 					started = false
 				} else {
+					removeMeta(dir)
 					require.NoError(t, cpt.Start(cfg))
 					started = true
 				}
@@ -292,14 +297,16 @@ func TestFilterCmds(t *testing.T) {
 		},
 	}
 
+	dir := t.TempDir()
 	cfg := CaptureConfig{
-		Output:   t.TempDir(),
+		Output:   dir,
 		Duration: 10 * time.Second,
 	}
 	for i, test := range tests {
 		cpt := NewCapture(zap.NewNop())
 		writer := newMockWriter(store.WriterCfg{})
 		cfg.cmdLogger = writer
+		removeMeta(dir)
 		require.NoError(t, cpt.Start(cfg))
 		cpt.Capture(test.packet, time.Now(), 100, func() (string, error) {
 			return "init session 100", nil
@@ -322,3 +329,7 @@ func TestFilterCmds(t *testing.T) {
 		cpt.Close()
 	}
 }
+
+func removeMeta(dir string) {
+	_ = os.Remove(filepath.Join(dir, "meta"))
+}

pkg/sqlreplay/store/meta.go

@@ -58,3 +58,12 @@ func (m *Meta) Read(path string) error {
 	}
 	return nil
 }
+
+func PreCheckMeta(path string) error {
+	filePath := filepath.Join(path, metaFile)
+	_, err := os.Stat(filePath)
+	if errors.Is(err, os.ErrNotExist) {
+		return nil
+	}
+	return errors.Errorf("file %s already exists, please remove it before capture", filePath)
+}

pkg/sqlreplay/store/meta_test.go

@@ -29,3 +29,12 @@ func TestMeta(t *testing.T) {
 	require.NoError(t, os.Remove(filepath.Join(dir, metaFile)))
 	require.Error(t, m3.Read(dir))
 }
+
+func TestPrecheckMeta(t *testing.T) {
+	dir := t.TempDir()
+	require.NoError(t, PreCheckMeta(dir))
+	f, err := os.Create(filepath.Join(dir, metaFile))
+	require.NoError(t, err)
+	require.NoError(t, f.Close())
+	require.Error(t, PreCheckMeta(dir))
+}