github.com/pingcap/br@v5.3.0-alpha.0.20220125034240-ec59c7b6ce30+incompatible/pkg/storage/compress.go (about)

     1  // Copyright 2020 PingCAP, Inc. Licensed under Apache-2.0.
     2  
     3  package storage
     4  
     5  import (
     6  	"bytes"
     7  	"context"
     8  	"io"
     9  
    10  	berrors "github.com/pingcap/br/pkg/errors"
    11  
    12  	"github.com/pingcap/errors"
    13  )
    14  
    15  type withCompression struct {
    16  	ExternalStorage
    17  	compressType CompressType
    18  }
    19  
    20  // WithCompression returns an ExternalStorage with compress option
    21  func WithCompression(inner ExternalStorage, compressionType CompressType) ExternalStorage {
    22  	if compressionType == NoCompression {
    23  		return inner
    24  	}
    25  	return &withCompression{ExternalStorage: inner, compressType: compressionType}
    26  }
    27  
    28  func (w *withCompression) Create(ctx context.Context, name string) (ExternalFileWriter, error) {
    29  	var (
    30  		writer ExternalFileWriter
    31  		err    error
    32  	)
    33  	if s3Storage, ok := w.ExternalStorage.(*S3Storage); ok {
    34  		writer, err = s3Storage.CreateUploader(ctx, name)
    35  	} else {
    36  		writer, err = w.ExternalStorage.Create(ctx, name)
    37  	}
    38  	if err != nil {
    39  		return nil, errors.Trace(err)
    40  	}
    41  	compressedWriter := newBufferedWriter(writer, hardcodedS3ChunkSize, w.compressType)
    42  	return compressedWriter, nil
    43  }
    44  
    45  func (w *withCompression) Open(ctx context.Context, path string) (ExternalFileReader, error) {
    46  	fileReader, err := w.ExternalStorage.Open(ctx, path)
    47  	if err != nil {
    48  		return nil, errors.Trace(err)
    49  	}
    50  	uncompressReader, err := newInterceptReader(fileReader, w.compressType)
    51  	if err != nil {
    52  		return nil, errors.Trace(err)
    53  	}
    54  	return uncompressReader, nil
    55  }
    56  
    57  func (w *withCompression) WriteFile(ctx context.Context, name string, data []byte) error {
    58  	bf := bytes.NewBuffer(make([]byte, 0, len(data)))
    59  	compressBf := newCompressWriter(w.compressType, bf)
    60  	_, err := compressBf.Write(data)
    61  	if err != nil {
    62  		return errors.Trace(err)
    63  	}
    64  	err = compressBf.Close()
    65  	if err != nil {
    66  		return errors.Trace(err)
    67  	}
    68  	return w.ExternalStorage.WriteFile(ctx, name, bf.Bytes())
    69  }
    70  
    71  func (w *withCompression) ReadFile(ctx context.Context, name string) ([]byte, error) {
    72  	data, err := w.ExternalStorage.ReadFile(ctx, name)
    73  	if err != nil {
    74  		return data, errors.Trace(err)
    75  	}
    76  	bf := bytes.NewBuffer(data)
    77  	compressBf, err := newCompressReader(w.compressType, bf)
    78  	if err != nil {
    79  		return nil, err
    80  	}
    81  	return io.ReadAll(compressBf)
    82  }
    83  
    84  type compressReader struct {
    85  	io.ReadCloser
    86  }
    87  
    88  // nolint:interfacer
    89  func newInterceptReader(fileReader ExternalFileReader, compressType CompressType) (ExternalFileReader, error) {
    90  	if compressType == NoCompression {
    91  		return fileReader, nil
    92  	}
    93  	r, err := newCompressReader(compressType, fileReader)
    94  	if err != nil {
    95  		return nil, errors.Trace(err)
    96  	}
    97  	return &compressReader{
    98  		ReadCloser: r,
    99  	}, nil
   100  }
   101  
   102  func (r *compressReader) Seek(_ int64, _ int) (int64, error) {
   103  	return int64(0), errors.Annotatef(berrors.ErrStorageInvalidConfig, "compressReader doesn't support Seek now")
   104  }
   105  
   106  type flushStorageWriter struct {
   107  	writer  io.Writer
   108  	flusher flusher
   109  	closer  io.Closer
   110  }
   111  
   112  func (w *flushStorageWriter) Write(ctx context.Context, data []byte) (int, error) {
   113  	n, err := w.writer.Write(data)
   114  	return n, errors.Trace(err)
   115  }
   116  
   117  func (w *flushStorageWriter) Close(ctx context.Context) error {
   118  	err := w.flusher.Flush()
   119  	if err != nil {
   120  		return errors.Trace(err)
   121  	}
   122  	return w.closer.Close()
   123  }
   124  
   125  func newFlushStorageWriter(writer io.Writer, flusher2 flusher, closer io.Closer) *flushStorageWriter {
   126  	return &flushStorageWriter{
   127  		writer:  writer,
   128  		flusher: flusher2,
   129  		closer:  closer,
   130  	}
   131  }