github.com/pingcap/tiflow@v0.0.0-20240520035814-5bf52d54e205/pkg/compression/compress.go (about)

     1  // Copyright 2023 PingCAP, Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package compression
    15  
    16  import (
    17  	"bytes"
    18  	"sync"
    19  
    20  	"github.com/klauspost/compress/snappy"
    21  	"github.com/pierrec/lz4/v4"
    22  	cerror "github.com/pingcap/tiflow/pkg/errors"
    23  )
    24  
    25  const (
    26  	// None no compression
    27  	None string = "none"
    28  
    29  	// Snappy compression
    30  	Snappy string = "snappy"
    31  
    32  	// LZ4 compression
    33  	LZ4 string = "lz4"
    34  )
    35  
    36  var (
    37  	lz4ReaderPool = sync.Pool{
    38  		New: func() interface{} {
    39  			return lz4.NewReader(nil)
    40  		},
    41  	}
    42  
    43  	bufferPool = sync.Pool{
    44  		New: func() interface{} {
    45  			return new(bytes.Buffer)
    46  		},
    47  	}
    48  )
    49  
    50  // Supported return true if the given compression is supported.
    51  func Supported(cc string) bool {
    52  	switch cc {
    53  	case None, Snappy, LZ4:
    54  		return true
    55  	}
    56  	return false
    57  }
    58  
    59  // Encode the given data by the given compression codec.
    60  func Encode(cc string, data []byte) ([]byte, error) {
    61  	switch cc {
    62  	case None:
    63  		return data, nil
    64  	case Snappy:
    65  		return snappy.Encode(nil, data), nil
    66  	case LZ4:
    67  		var buf bytes.Buffer
    68  		writer := lz4.NewWriter(&buf)
    69  		if _, err := writer.Write(data); err != nil {
    70  			return nil, cerror.WrapError(cerror.ErrCompressionFailed, err)
    71  		}
    72  		if err := writer.Close(); err != nil {
    73  			return nil, cerror.WrapError(cerror.ErrCompressionFailed, err)
    74  		}
    75  		return buf.Bytes(), nil
    76  	default:
    77  	}
    78  
    79  	return nil, cerror.ErrCompressionFailed.GenWithStack("Unsupported compression %s", cc)
    80  }
    81  
    82  // Decode the given data by the given compression codec.
    83  func Decode(cc string, data []byte) ([]byte, error) {
    84  	switch cc {
    85  	case None:
    86  		return data, nil
    87  	case Snappy:
    88  		return snappy.Decode(nil, data)
    89  	case LZ4:
    90  		reader, ok := lz4ReaderPool.Get().(*lz4.Reader)
    91  		if !ok {
    92  			reader = lz4.NewReader(bytes.NewReader(data))
    93  		} else {
    94  			reader.Reset(bytes.NewReader(data))
    95  		}
    96  		buffer := bufferPool.Get().(*bytes.Buffer)
    97  		_, err := buffer.ReadFrom(reader)
    98  		// copy the buffer to a new slice with the correct length
    99  		// reuse lz4Reader and buffer
   100  		lz4ReaderPool.Put(reader)
   101  		res := make([]byte, buffer.Len())
   102  		copy(res, buffer.Bytes())
   103  		buffer.Reset()
   104  		bufferPool.Put(buffer)
   105  
   106  		return res, err
   107  	default:
   108  	}
   109  
   110  	return nil, cerror.ErrCompressionFailed.GenWithStack("Unsupported compression %s", cc)
   111  }