github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/rpc/snappy.go (about)

     1  // Copyright 2017 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  package rpc
    12  
    13  import (
    14  	"io"
    15  	"sync"
    16  
    17  	"github.com/golang/snappy"
    18  	"google.golang.org/grpc/encoding"
    19  )
    20  
    21  // NB: The encoding.Compressor implementation needs to be goroutine
    22  // safe as multiple goroutines may be using the same compressor for
    23  // different streams on the same connection.
    24  var snappyWriterPool sync.Pool
    25  var snappyReaderPool sync.Pool
    26  
    27  type snappyWriter struct {
    28  	*snappy.Writer
    29  }
    30  
    31  func (w *snappyWriter) Close() error {
    32  	defer snappyWriterPool.Put(w)
    33  	return w.Writer.Close()
    34  }
    35  
    36  type snappyReader struct {
    37  	*snappy.Reader
    38  }
    39  
    40  func (r *snappyReader) Read(p []byte) (n int, err error) {
    41  	n, err = r.Reader.Read(p)
    42  	if err == io.EOF {
    43  		snappyReaderPool.Put(r)
    44  	}
    45  	return n, err
    46  }
    47  
    48  type snappyCompressor struct {
    49  }
    50  
    51  func (snappyCompressor) Name() string {
    52  	return "snappy"
    53  }
    54  
    55  func (snappyCompressor) Compress(w io.Writer) (io.WriteCloser, error) {
    56  	sw, ok := snappyWriterPool.Get().(*snappyWriter)
    57  	if !ok {
    58  		sw = &snappyWriter{snappy.NewBufferedWriter(w)}
    59  	} else {
    60  		sw.Reset(w)
    61  	}
    62  	return sw, nil
    63  }
    64  
    65  func (snappyCompressor) Decompress(r io.Reader) (io.Reader, error) {
    66  	sr, ok := snappyReaderPool.Get().(*snappyReader)
    67  	if !ok {
    68  		sr = &snappyReader{snappy.NewReader(r)}
    69  	} else {
    70  		sr.Reset(r)
    71  	}
    72  	return sr, nil
    73  }
    74  
    75  func init() {
    76  	encoding.RegisterCompressor(snappyCompressor{})
    77  }