github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/engine/common/grpc/compressor/snappy/snappy.go (about)

     1  package snappy
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  
     7  	"github.com/golang/snappy"
     8  	"google.golang.org/grpc/encoding"
     9  )
    10  
    11  // Name is the name registered for the Snappy compressor.
    12  const Name = "snappy"
    13  
    14  func init() {
    15  	c := &compressor{}
    16  	c.poolCompressor.New = func() interface{} {
    17  		return &writer{Writer: snappy.NewBufferedWriter(nil), pool: &c.poolCompressor}
    18  	}
    19  	encoding.RegisterCompressor(c)
    20  }
    21  
    22  type compressor struct {
    23  	poolCompressor   sync.Pool
    24  	poolDecompressor sync.Pool
    25  }
    26  
    27  func (c *compressor) Name() string {
    28  	return Name
    29  }
    30  
    31  func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
    32  	sw := c.poolCompressor.Get().(*writer)
    33  	sw.Reset(w)
    34  	return sw, nil
    35  }
    36  
    37  func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
    38  	sr, inPool := c.poolDecompressor.Get().(*reader)
    39  	if !inPool {
    40  		return snappy.NewReader(r), nil
    41  	}
    42  	sr.Reset(r)
    43  	return sr, nil
    44  }
    45  
    46  type writer struct {
    47  	*snappy.Writer
    48  	pool *sync.Pool
    49  }
    50  
    51  func (w *writer) Close() error {
    52  	defer w.pool.Put(w)
    53  	return w.Writer.Close()
    54  }
    55  
    56  type reader struct {
    57  	*snappy.Reader
    58  	pool *sync.Pool
    59  }
    60  
    61  func (r *reader) Read(p []byte) (n int, err error) {
    62  	n, err = r.Reader.Read(p)
    63  	if err == io.EOF {
    64  		r.pool.Put(r)
    65  	}
    66  	return n, err
    67  }