github.com/thanos-io/thanos@v0.32.5/pkg/extgrpc/snappy/snappy.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package snappy
     5  
     6  import (
     7  	"io"
     8  	"sync"
     9  
    10  	"github.com/klauspost/compress/snappy"
    11  	"google.golang.org/grpc/encoding"
    12  )
    13  
    14  // Name is the name registered for the snappy compressor.
    15  const Name = "snappy"
    16  
    17  var Compressor *compressor = newCompressor()
    18  
    19  func init() {
    20  	encoding.RegisterCompressor(Compressor)
    21  }
    22  
    23  type compressor struct {
    24  	writersPool sync.Pool
    25  	readersPool sync.Pool
    26  }
    27  
    28  func newCompressor() *compressor {
    29  	c := &compressor{}
    30  	c.readersPool = sync.Pool{
    31  		New: func() interface{} {
    32  			return snappy.NewReader(nil)
    33  		},
    34  	}
    35  	c.writersPool = sync.Pool{
    36  		New: func() interface{} {
    37  			return snappy.NewBufferedWriter(nil)
    38  		},
    39  	}
    40  	return c
    41  }
    42  
    43  func (c *compressor) Name() string {
    44  	return Name
    45  }
    46  
    47  func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
    48  	wr := c.writersPool.Get().(*snappy.Writer)
    49  	wr.Reset(w)
    50  	return writeCloser{wr, &c.writersPool}, nil
    51  }
    52  
    53  func (c *compressor) DecompressByteReader(r io.Reader) (io.ByteReader, error) {
    54  	dr := c.readersPool.Get().(*snappy.Reader)
    55  	dr.Reset(r)
    56  	return reader{dr, &c.readersPool}, nil
    57  }
    58  
    59  func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
    60  	dr := c.readersPool.Get().(*snappy.Reader)
    61  	dr.Reset(r)
    62  	return reader{dr, &c.readersPool}, nil
    63  }
    64  
    65  type writeCloser struct {
    66  	writer *snappy.Writer
    67  	pool   *sync.Pool
    68  }
    69  
    70  func (w writeCloser) Write(p []byte) (n int, err error) {
    71  	return w.writer.Write(p)
    72  }
    73  
    74  func (w writeCloser) Close() error {
    75  	defer func() {
    76  		w.writer.Reset(nil)
    77  		w.pool.Put(w.writer)
    78  	}()
    79  
    80  	if w.writer != nil {
    81  		return w.writer.Close()
    82  	}
    83  	return nil
    84  }
    85  
    86  type reader struct {
    87  	reader *snappy.Reader
    88  	pool   *sync.Pool
    89  }
    90  
    91  func (r reader) Read(p []byte) (n int, err error) {
    92  	n, err = r.reader.Read(p)
    93  	if err == io.EOF {
    94  		r.reader.Reset(nil)
    95  		r.pool.Put(r.reader)
    96  	}
    97  	return n, err
    98  }
    99  
   100  func (r reader) ReadByte() (n byte, err error) {
   101  	n, err = r.reader.ReadByte()
   102  	if err == io.EOF {
   103  		r.reader.Reset(nil)
   104  		r.pool.Put(r.reader)
   105  	}
   106  	return n, err
   107  
   108  }