trpc.group/trpc-go/trpc-go@v1.0.3/codec/compress_snappy.go (about)

     1  //
     2  //
     3  // Tencent is pleased to support the open source community by making tRPC available.
     4  //
     5  // Copyright (C) 2023 THL A29 Limited, a Tencent company.
     6  // All rights reserved.
     7  //
     8  // If you have downloaded a copy of the tRPC source code from Tencent,
     9  // please note that tRPC source code is licensed under the  Apache 2.0 License,
    10  // A copy of the Apache 2.0 License is included in this file.
    11  //
    12  //
    13  
    14  package codec
    15  
    16  import (
    17  	"bytes"
    18  	"io"
    19  	"sync"
    20  
    21  	"github.com/golang/snappy"
    22  )
    23  
    24  func init() {
    25  	RegisterCompressor(CompressTypeSnappy, NewSnappyCompressor())
    26  	RegisterCompressor(CompressTypeStreamSnappy, NewSnappyCompressor())
    27  	RegisterCompressor(CompressTypeBlockSnappy, NewSnappyBlockCompressor())
    28  }
    29  
    30  // SnappyCompress is snappy compressor using stream snappy format.
    31  //
    32  // There are actually two Snappy formats: block and stream. They are related,
    33  // but different: trying to decompress block-compressed data as a Snappy stream
    34  // will fail, and vice versa.
    35  type SnappyCompress struct {
    36  	writerPool *sync.Pool
    37  	readerPool *sync.Pool
    38  }
    39  
    40  // NewSnappyCompressor returns a stream format snappy compressor instance.
    41  func NewSnappyCompressor() *SnappyCompress {
    42  	s := &SnappyCompress{}
    43  	s.writerPool = &sync.Pool{
    44  		New: func() interface{} {
    45  			return snappy.NewBufferedWriter(&bytes.Buffer{})
    46  		},
    47  	}
    48  	s.readerPool = &sync.Pool{
    49  		New: func() interface{} {
    50  			return snappy.NewReader(&bytes.Buffer{})
    51  		},
    52  	}
    53  	return s
    54  }
    55  
    56  // Compress returns binary data compressed by snappy stream format.
    57  func (c *SnappyCompress) Compress(in []byte) ([]byte, error) {
    58  	if len(in) == 0 {
    59  		return in, nil
    60  	}
    61  
    62  	buf := &bytes.Buffer{}
    63  	writer := c.getSnappyWriter(buf)
    64  	defer func() {
    65  		if c.writerPool != nil {
    66  			c.writerPool.Put(writer)
    67  		}
    68  	}()
    69  
    70  	if _, err := writer.Write(in); err != nil {
    71  		writer.Close()
    72  		return nil, err
    73  	}
    74  	if err := writer.Close(); err != nil {
    75  		return nil, err
    76  	}
    77  	return buf.Bytes(), nil
    78  }
    79  
    80  // Decompress returns binary data decompressed by snappy stream format.
    81  func (c *SnappyCompress) Decompress(in []byte) ([]byte, error) {
    82  	if len(in) == 0 {
    83  		return in, nil
    84  	}
    85  
    86  	inReader := bytes.NewReader(in)
    87  	reader := c.getSnappyReader(inReader)
    88  	defer func() {
    89  		if c.readerPool != nil {
    90  			c.readerPool.Put(reader)
    91  		}
    92  	}()
    93  
    94  	out, err := io.ReadAll(reader)
    95  	if err != nil {
    96  		return nil, err
    97  	}
    98  	return out, err
    99  }
   100  
   101  // SnappyBlockCompressor is snappy compressor using snappy block format.
   102  type SnappyBlockCompressor struct{}
   103  
   104  // NewSnappyBlockCompressor returns a block format snappy compressor instance.
   105  func NewSnappyBlockCompressor() *SnappyBlockCompressor {
   106  	return &SnappyBlockCompressor{}
   107  }
   108  
   109  // Compress returns binary data compressed by snappy block formats.
   110  func (c *SnappyBlockCompressor) Compress(in []byte) ([]byte, error) {
   111  	if len(in) == 0 {
   112  		return in, nil
   113  	}
   114  	return snappy.Encode(nil, in), nil
   115  }
   116  
   117  // Decompress returns binary data decompressed by snappy block formats.
   118  func (c *SnappyBlockCompressor) Decompress(in []byte) ([]byte, error) {
   119  	if len(in) == 0 {
   120  		return in, nil
   121  	}
   122  	return snappy.Decode(nil, in)
   123  }
   124  
   125  func (c *SnappyCompress) getSnappyWriter(buf *bytes.Buffer) *snappy.Writer {
   126  	if c.writerPool == nil {
   127  		return snappy.NewBufferedWriter(buf)
   128  	}
   129  
   130  	// get from pool
   131  	writer, ok := c.writerPool.Get().(*snappy.Writer)
   132  	if !ok || writer == nil {
   133  		return snappy.NewBufferedWriter(buf)
   134  	}
   135  	writer.Reset(buf)
   136  	return writer
   137  }
   138  
   139  func (c *SnappyCompress) getSnappyReader(inReader *bytes.Reader) *snappy.Reader {
   140  	if c.readerPool == nil {
   141  		return snappy.NewReader(inReader)
   142  	}
   143  
   144  	// get from pool
   145  	reader, ok := c.readerPool.Get().(*snappy.Reader)
   146  	if !ok || reader == nil {
   147  		return snappy.NewReader(inReader)
   148  	}
   149  	reader.Reset(inReader)
   150  	return reader
   151  }