sqlreplay: remove lumberjack to support large commands (#727)

Author: djshow832

go.mod

@@ -34,7 +34,6 @@ require (
 	go.uber.org/ratelimit v0.2.0
 	go.uber.org/zap v1.26.0
 	google.golang.org/grpc v1.54.0
-	gopkg.in/natefinch/lumberjack.v2 v2.2.1
 )
 
 require (
@@ -117,6 +116,7 @@ require (
 	golang.org/x/time v0.3.0 // indirect
 	google.golang.org/genproto v0.0.0-20230110181048-76db0878b65f // indirect
 	google.golang.org/protobuf v1.31.0 // indirect
+	gopkg.in/natefinch/lumberjack.v2 v2.2.1 // indirect
 	gopkg.in/yaml.v2 v2.4.0 // indirect
 	gopkg.in/yaml.v3 v3.0.1 // indirect
 	sigs.k8s.io/yaml v1.3.0 // indirect

pkg/sqlreplay/capture/capture.go

@@ -6,6 +6,7 @@ package capture
 import (
 	"bytes"
 	"context"
+	"io"
 	"os"
 	"sync"
 	"time"
@@ -59,7 +60,7 @@ type CaptureConfig struct {
 	StartTime          time.Time
 	Duration           time.Duration
 	Compress           bool
-	cmdLogger          store.Writer
+	cmdLogger          io.WriteCloser
 	bufferCap          int
 	flushThreshold     int
 	maxBuffers         int
@@ -231,12 +232,14 @@ func (c *capture) collectCmds(bufCh chan<- *bytes.Buffer) {
 	}
 }
 
+// Writing commands requires a bytes buffer instead of a simple bufio.Writer,
+// so the buffer can not be pushed down to the store package.
 func (c *capture) flushBuffer(bufCh <-chan *bytes.Buffer) {
 	// cfg.cmdLogger is set in tests
 	cmdLogger := c.cfg.cmdLogger
 	if cmdLogger == nil {
 		var err error
-		cmdLogger, err = store.NewWriter(store.WriterCfg{
+		cmdLogger, err = store.NewWriter(c.lg.Named("writer"), store.WriterCfg{
 			Dir:           c.cfg.Output,
 			EncryptMethod: c.cfg.EncryptMethod,
 			KeyFile:       c.cfg.KeyFile,
@@ -249,8 +252,7 @@ func (c *capture) flushBuffer(bufCh <-chan *bytes.Buffer) {
 	}
 	// Flush all buffers even if the context is timeout.
 	for buf := range bufCh {
-		// TODO: each write size should be less than MaxSize.
-		if err := cmdLogger.Write(buf.Bytes()); err != nil {
+		if _, err := cmdLogger.Write(buf.Bytes()); err != nil {
 			c.stop(errors.Wrapf(err, "failed to flush traffic to disk"))
 			break
 		}

pkg/sqlreplay/capture/mock_test.go

@@ -5,12 +5,13 @@ package capture
 
 import (
 	"bytes"
+	"io"
 	"sync"
 
 	"github.com/pingcap/tiproxy/pkg/sqlreplay/store"
 )
 
-var _ store.Writer = (*mockWriter)(nil)
+var _ io.WriteCloser = (*mockWriter)(nil)
 
 type mockWriter struct {
 	sync.Mutex
@@ -21,11 +22,10 @@ func newMockWriter(store.WriterCfg) *mockWriter {
 	return &mockWriter{}
 }
 
-func (w *mockWriter) Write(p []byte) error {
+func (w *mockWriter) Write(p []byte) (int, error) {
 	w.Lock()
 	defer w.Unlock()
-	_, err := w.buf.Write(p)
-	return err
+	return w.buf.Write(p)
 }
 
 func (w *mockWriter) getData() []byte {

pkg/sqlreplay/store/compress.go

@@ -0,0 +1,50 @@
+// Copyright 2024 PingCAP, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package store
+
+import (
+	"compress/gzip"
+	"io"
+
+	"go.uber.org/zap"
+)
+
+var _ io.WriteCloser = (*compressWriter)(nil)
+
+type compressWriter struct {
+	io.WriteCloser
+	internalWriter io.WriteCloser
+	lg             *zap.Logger
+}
+
+func newCompressWriter(lg *zap.Logger, writer io.WriteCloser) *compressWriter {
+	return &compressWriter{
+		WriteCloser:    gzip.NewWriter(writer),
+		internalWriter: writer,
+		lg:             lg,
+	}
+}
+
+func (w *compressWriter) Close() error {
+	if err := w.WriteCloser.Close(); err != nil {
+		w.lg.Warn("failed to close writer", zap.Error(err))
+	}
+	return w.internalWriter.Close()
+}
+
+var _ io.Reader = (*compressReader)(nil)
+
+type compressReader struct {
+	*gzip.Reader
+}
+
+func newCompressReader(reader io.Reader) (*compressReader, error) {
+	gr, err := gzip.NewReader(reader)
+	if err != nil {
+		return nil, err
+	}
+	return &compressReader{
+		Reader: gr,
+	}, nil
+}

pkg/sqlreplay/store/compress_test.go

@@ -0,0 +1,38 @@
+// Copyright 2024 PingCAP, Inc.
+// SPDX-License-Identifier: Apache-2.0
+
+package store
+
+import (
+	"io"
+	"os"
+	"path/filepath"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.uber.org/zap"
+)
+
+func TestCompressReadWrite(t *testing.T) {
+	dir := t.TempDir()
+	path := filepath.Join(dir, "test")
+
+	file, err := os.OpenFile(path, os.O_CREATE|os.O_TRUNC|os.O_WRONLY, 0600)
+	require.NoError(t, err)
+	writer := newCompressWriter(zap.NewNop(), file)
+	n, err := writer.Write([]byte("test"))
+	require.NoError(t, err)
+	require.Equal(t, 4, n)
+	require.NoError(t, writer.Close())
+	// file is already closed
+	require.Error(t, file.Close())
+
+	file, err = os.OpenFile(path, os.O_RDONLY, 0600)
+	require.NoError(t, err)
+	reader, err := newCompressReader(file)
+	require.NoError(t, err)
+	data := make([]byte, 100)
+	n, err = io.ReadFull(reader, data)
+	require.Equal(t, 4, n)
+	require.ErrorContains(t, err, "EOF")
+}

pkg/sqlreplay/store/const.go

@@ -4,10 +4,9 @@
 package store
 
 const (
-	fileNamePrefix     = "traffic"
+	fileNamePrefix     = "traffic-"
 	fileNameSuffix     = ".log"
-	fileName           = fileNamePrefix + fileNameSuffix
-	fileTsLayout       = "2006-01-02T15-04-05.000"
 	fileCompressFormat = ".gz"
-	fileSize           = 300 // 300MB
+	fileSize           = 300 << 20
+	bufferSize         = 1 << 20
 )

pkg/sqlreplay/store/encrypt.go

@@ -7,27 +7,43 @@ import (
 	"crypto/aes"
 	"crypto/cipher"
 	"crypto/rand"
+	"fmt"
 	"io"
 	"os"
-	"reflect"
+	"strings"
 
 	"github.com/pingcap/tiproxy/lib/util/errors"
 )
 
-var _ Writer = (*aesCTRWriter)(nil)
+const (
+	EncryptPlain = "plaintext"
+	EncryptAes   = "aes256-ctr"
+)
+
+var _ io.WriteCloser = (*aesCTRWriter)(nil)
 
 type aesCTRWriter struct {
-	Writer
+	io.WriteCloser
 	stream cipher.Stream
-	iv     []byte
-	inited bool
 }
 
-func newAESCTRWriter(writer Writer, keyFile string) (*aesCTRWriter, error) {
+func newWriterWithEncryptOpts(writer io.WriteCloser, encryptMethod string, keyFile string) (io.WriteCloser, error) {
+	switch strings.ToLower(encryptMethod) {
+	case "", EncryptPlain:
+		return writer, nil
+	case EncryptAes:
+	default:
+		return nil, fmt.Errorf("unsupported encrypt method: %s", encryptMethod)
+	}
+
 	key, err := readAesKey(keyFile)
 	if err != nil {
 		return nil, err
 	}
+	return newAESCTRWriter(writer, key)
+}
+
+func newAESCTRWriter(writer io.WriteCloser, key []byte) (*aesCTRWriter, error) {
 	block, err := aes.NewCipher(key)
 	if err != nil {
 		return nil, errors.WithStack(err)
@@ -36,93 +52,76 @@ func newAESCTRWriter(writer Writer, keyFile string) (*aesCTRWriter, error) {
 	if _, err := io.ReadFull(rand.Reader, iv); err != nil {
 		return nil, errors.WithStack(err)
 	}
-	return &aesCTRWriter{
-		Writer: writer,
-		stream: cipher.NewCTR(block, iv),
-		iv:     iv,
-	}, nil
-}
-
-func (ctr *aesCTRWriter) Write(data []byte) error {
-	if !ctr.inited {
-		if err := ctr.writeIV(); err != nil {
-			return err
-		}
-		ctr.inited = true
+	ctr := &aesCTRWriter{
+		WriteCloser: writer,
+		stream:      cipher.NewCTR(block, iv),
 	}
-	ctr.stream.XORKeyStream(data, data)
-	return ctr.Writer.Write(data)
+	_, err = ctr.WriteCloser.Write(iv)
+	return ctr, err
 }
 
-func (ctr *aesCTRWriter) writeIV() error {
-	return ctr.Writer.Write(ctr.iv)
-}
-
-func (ctr *aesCTRWriter) Close() error {
-	return ctr.Writer.Close()
+func (ctr *aesCTRWriter) Write(data []byte) (int, error) {
+	ctr.stream.XORKeyStream(data, data)
+	return ctr.WriteCloser.Write(data)
 }
 
-var _ Reader = (*aesCTRReader)(nil)
+var _ io.Reader = (*aesCTRReader)(nil)
 
 type aesCTRReader struct {
-	Reader
+	io.Reader
 	stream cipher.Stream
-	key    []byte
 }
 
-func newAESCTRReader(reader Reader, keyFile string) (*aesCTRReader, error) {
+func newReaderWithEncryptOpts(reader io.Reader, encryptMethod string, keyFile string) (io.Reader, error) {
+	switch strings.ToLower(encryptMethod) {
+	case "", EncryptPlain:
+		return reader, nil
+	case EncryptAes:
+	default:
+		return nil, fmt.Errorf("unsupported encrypt method: %s", encryptMethod)
+	}
+
 	key, err := readAesKey(keyFile)
 	if err != nil {
 		return nil, err
 	}
+	return newAESCTRReader(reader, key)
+}
+
+func newAESCTRReader(reader io.Reader, key []byte) (*aesCTRReader, error) {
+	block, err := aes.NewCipher(key)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	iv := make([]byte, aes.BlockSize)
+	for readLen := 0; readLen < len(iv); {
+		m, err := reader.Read(iv[readLen:])
+		if err != nil {
+			return nil, err
+		}
+		readLen += m
+	}
 	return &aesCTRReader{
 		Reader: reader,
-		key:    key,
+		stream: cipher.NewCTR(block, iv),
 	}, nil
 }
 
 func (ctr *aesCTRReader) Read(data []byte) (int, error) {
-	if ctr.stream == nil || reflect.ValueOf(ctr.stream).IsNil() {
-		if err := ctr.init(); err != nil {
-			return 0, err
-		}
-	}
 	n, err := ctr.Reader.Read(data)
 	if n > 0 {
 		ctr.stream.XORKeyStream(data[:n], data[:n])
 	}
 	if err != nil {
-		return n, err
+		return n, errors.WithStack(err)
 	}
 	return n, nil
 }
 
-func (ctr *aesCTRReader) init() error {
-	block, err := aes.NewCipher(ctr.key)
-	if err != nil {
-		return errors.WithStack(err)
-	}
-	iv := make([]byte, aes.BlockSize)
-	for readLen := 0; readLen < len(iv); {
-		m, err := ctr.Reader.Read(iv[readLen:])
-		if err != nil {
-			return err
-		}
-		readLen += m
-	}
-	ctr.stream = cipher.NewCTR(block, iv)
-	return nil
-}
-
-func (ctr *aesCTRReader) CurFile() string {
-	return ctr.Reader.CurFile()
-}
-
-func (ctr *aesCTRReader) Close() {
-	ctr.Reader.Close()
-}
-
 func readAesKey(filename string) ([]byte, error) {
+	if len(filename) == 0 {
+		return nil, errors.New("encryption key file name is not set")
+	}
 	key, err := os.ReadFile(filename)
 	if err != nil {
 		return nil, errors.WithStack(err)