sqlreplay, api: support compress option to traffic capture (#720)

Author: djshow832

lib/cli/traffic.go

@@ -30,8 +30,14 @@ func GetTrafficCaptureCmd(ctx *Context) *cobra.Command {
 	output := captureCmd.PersistentFlags().String("output", "", "output directory for traffic files")
 	duration := captureCmd.PersistentFlags().String("duration", "", "the duration of traffic capture")
 	encrypt := captureCmd.PersistentFlags().String("encrypt-method", "", "the encryption method used for encrypting traffic files")
+	compress := captureCmd.PersistentFlags().Bool("compress", true, "whether compress the traffic files")
 	captureCmd.RunE = func(cmd *cobra.Command, args []string) error {
-		reader := GetFormReader(map[string]string{"output": *output, "duration": *duration, "encrypt-method": *encrypt})
+		reader := GetFormReader(map[string]string{
+			"output":         *output,
+			"duration":       *duration,
+			"encrypt-method": *encrypt,
+			"compress":       strconv.FormatBool(*compress),
+		})
 		resp, err := doRequest(cmd.Context(), ctx, http.MethodPost, "/api/traffic/capture", reader)
 		if err != nil {
 			return err

pkg/server/api/traffic.go

@@ -12,6 +12,7 @@ import (
 	"github.com/gin-gonic/gin"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/capture"
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/replay"
+	"go.uber.org/zap"
 )
 
 func (h *Server) registerTraffic(group *gin.RouterGroup) {
@@ -33,6 +34,16 @@ func (h *Server) TrafficCapture(c *gin.Context) {
 		cfg.Duration = duration
 	}
 	cfg.EncryptMethod = c.PostForm("encrypt-method")
+
+	compress := true
+	if compressStr := c.PostForm("compress"); compressStr != "" {
+		var err error
+		if compress, err = strconv.ParseBool(compressStr); err != nil {
+			h.lg.Warn("parsing argument 'compress' error, using true", zap.String("compress", c.PostForm("compress")), zap.Error(err))
+			compress = true
+		}
+	}
+	cfg.Compress = compress
 	cfg.KeyFile = h.mgr.CfgMgr.GetConfig().Security.Encryption.KeyPath
 
 	if err := h.mgr.ReplayJobMgr.StartCapture(cfg); err != nil {

pkg/server/api/traffic_test.go

@@ -39,7 +39,25 @@ func TestTraffic(t *testing.T) {
 		require.NoError(t, err)
 		require.Equal(t, "capture started", string(all))
 		require.Equal(t, "capture", mgr.curJob)
-		require.Equal(t, capture.CaptureConfig{Duration: time.Hour, Output: "/tmp"}, mgr.captureCfg)
+		require.Equal(t, capture.CaptureConfig{Duration: time.Hour, Output: "/tmp", Compress: true}, mgr.captureCfg)
+	})
+	doHTTP(t, http.MethodPost, "/api/traffic/cancel", httpOpts{}, func(t *testing.T, r *http.Response) {
+		require.Equal(t, http.StatusOK, r.StatusCode)
+		all, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		require.Equal(t, "stopped", string(all))
+		require.Equal(t, "", mgr.curJob)
+	})
+	doHTTP(t, http.MethodPost, "/api/traffic/capture", httpOpts{
+		reader: cli.GetFormReader(map[string]string{"output": "/tmp", "duration": "1h", "encrypt-method": "aes256-ctr", "compress": "false"}),
+		header: map[string]string{"Content-Type": "application/x-www-form-urlencoded"},
+	}, func(t *testing.T, r *http.Response) {
+		require.Equal(t, http.StatusOK, r.StatusCode)
+		all, err := io.ReadAll(r.Body)
+		require.NoError(t, err)
+		require.Equal(t, "capture started", string(all))
+		require.Equal(t, "capture", mgr.curJob)
+		require.Equal(t, capture.CaptureConfig{Duration: time.Hour, Output: "/tmp", EncryptMethod: "aes256-ctr", Compress: false}, mgr.captureCfg)
 	})
 	doHTTP(t, http.MethodPost, "/api/traffic/replay", httpOpts{
 		reader: cli.GetFormReader(map[string]string{"input": "/tmp"}),
@@ -52,10 +70,6 @@ func TestTraffic(t *testing.T) {
 	})
 	doHTTP(t, http.MethodPost, "/api/traffic/cancel", httpOpts{}, func(t *testing.T, r *http.Response) {
 		require.Equal(t, http.StatusOK, r.StatusCode)
-		all, err := io.ReadAll(r.Body)
-		require.NoError(t, err)
-		require.Equal(t, "stopped", string(all))
-		require.Equal(t, "", mgr.curJob)
 	})
 	doHTTP(t, http.MethodPost, "/api/traffic/replay", httpOpts{
 		reader: cli.GetFormReader(map[string]string{"input": "/tmp", "speed": "abc"}),

pkg/sqlreplay/capture/capture.go

@@ -55,6 +55,7 @@ type CaptureConfig struct {
 	EncryptMethod      string
 	KeyFile            string
 	Duration           time.Duration
+	Compress           bool
 	cmdLogger          store.Writer
 	bufferCap          int
 	flushThreshold     int
@@ -226,6 +227,7 @@ func (c *capture) flushBuffer(bufCh <-chan *bytes.Buffer) {
 			Dir:           c.cfg.Output,
 			EncryptMethod: c.cfg.EncryptMethod,
 			KeyFile:       c.cfg.KeyFile,
+			Compress:      c.cfg.Compress,
 		})
 		if err != nil {
 			c.lg.Error("failed to create capture writer", zap.Error(err))

pkg/sqlreplay/store/line.go

@@ -22,6 +22,7 @@ type WriterCfg struct {
 	EncryptMethod string
 	KeyFile       string
 	FileSize      int
+	Compress      bool
 }
 
 func NewWriter(cfg WriterCfg) (Writer, error) {

pkg/sqlreplay/store/rotate.go

@@ -38,7 +38,7 @@ func newRotateWriter(cfg WriterCfg) *rotateWriter {
 			Filename:  filepath.Join(cfg.Dir, fileName),
 			MaxSize:   cfg.FileSize,
 			LocalTime: true,
-			Compress:  true,
+			Compress:  cfg.Compress,
 		},
 	}
 }

pkg/sqlreplay/store/rotate_test.go

@@ -22,6 +22,7 @@ func TestFileRotation(t *testing.T) {
 	writer := newRotateWriter(WriterCfg{
 		Dir:      tmpDir,
 		FileSize: 1,
+		Compress: true,
 	})
 	defer writer.Close()
 
@@ -186,39 +187,55 @@ func TestIterateFiles(t *testing.T) {
 
 func TestReadGZip(t *testing.T) {
 	tmpDir := t.TempDir()
-	writer := newRotateWriter(WriterCfg{
-		Dir:      tmpDir,
-		FileSize: 1,
-	})
-	defer writer.Close()
+	for _, compress := range []bool{true, false} {
+		require.NoError(t, os.RemoveAll(tmpDir))
+		require.NoError(t, os.MkdirAll(tmpDir, 0777))
 
-	data := make([]byte, 100*1024)
-	for i := 0; i < 11; i++ {
-		require.NoError(t, writer.Write(data))
-	}
-	// files are rotated and compressed at backendground asynchronously
-	require.Eventually(t, func() bool {
-		files := listFiles(t, tmpDir)
-		for _, f := range files {
-			if strings.HasPrefix(f, "traffic") && strings.HasSuffix(f, ".gz") {
-				return true
+		writer := newRotateWriter(WriterCfg{
+			Dir:      tmpDir,
+			FileSize: 1,
+			Compress: compress,
+		})
+		data := make([]byte, 100*1024)
+		for i := 0; i < 11; i++ {
+			require.NoError(t, writer.Write(data))
+		}
+		// files are rotated and compressed at backendground asynchronously
+		if compress {
+			require.Eventually(t, func() bool {
+				files := listFiles(t, tmpDir)
+				for _, f := range files {
+					if strings.HasPrefix(f, "traffic") && strings.HasSuffix(f, ".gz") {
+						return true
+					}
+				}
+				t.Logf("traffic files: %v", files)
+				return false
+			}, 5*time.Second, 10*time.Millisecond)
+		} else {
+			time.Sleep(100 * time.Millisecond)
+			files := listFiles(t, tmpDir)
+			for _, f := range files {
+				if strings.HasPrefix(f, "traffic") {
+					require.False(t, strings.HasSuffix(f, ".gz"))
+				}
 			}
 		}
-		t.Logf("traffic files: %v", files)
-		return false
-	}, 5*time.Second, 10*time.Millisecond)
+		require.NoError(t, writer.Close())
 
-	lg, _ := logger.CreateLoggerForTest(t)
-	l := newRotateReader(lg, tmpDir)
-	for i := 0; i < 11; i++ {
-		data = make([]byte, 100*1024)
-		_, err := io.ReadFull(l, data)
-		require.NoError(t, err)
-		for j := 0; j < 100*1024; j++ {
-			require.Equal(t, byte(0), data[j])
+		lg, _ := logger.CreateLoggerForTest(t)
+		l := newRotateReader(lg, tmpDir)
+		for i := 0; i < 11; i++ {
+			data = make([]byte, 100*1024)
+			_, err := io.ReadFull(l, data)
+			require.NoError(t, err)
+			for j := 0; j < 100*1024; j++ {
+				require.Equal(t, byte(0), data[j])
+			}
 		}
+		data = make([]byte, 1)
+		_, err := l.Read(data)
+		require.True(t, errors.Is(err, io.EOF))
+		l.Close()
 	}
-	data = make([]byte, 1)
-	_, err := l.Read(data)
-	require.True(t, errors.Is(err, io.EOF))
 }