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

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