go.uber.org/yarpc@v1.72.1/compressor/snappy/snappy.go (about)

     1  // Copyright (c) 2022 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  // Package yarpcsnappy provides a YARPC binding for snappy compression.
    22  package yarpcsnappy
    23  
    24  import (
    25  	"io"
    26  	"sync"
    27  
    28  	"github.com/golang/snappy"
    29  	"go.uber.org/yarpc/api/transport"
    30  )
    31  
    32  // Option is an option argument for the Snappy compressor constructor, New.
    33  type Option interface {
    34  	apply(*Compressor)
    35  }
    36  
    37  // New returns a Snappy compression strategy, suitable for configuring
    38  // an outbound dialer.
    39  //
    40  // The compressor is compatible with the gRPC experimental compressor system.
    41  // However, since gRPC requires global registration of compressors,
    42  // you must arrange for the compressor to be registered in your
    43  // application initialization.
    44  //
    45  //  import (
    46  //      "google.golang.org/grpc/encoding"
    47  //      "go.uber.org/yarpc/compressor/grpc"
    48  //      "go.uber.org/yarpc/compressor/snappy"
    49  //  )
    50  //
    51  //  var SnappyCompressor = yarpcsnappy.New()
    52  //
    53  //  func init()
    54  //      sc := yarpcgrpccompressor.New(SnappyCompressor)
    55  //      encoding.RegisterCompressor(sc)
    56  //  }
    57  //
    58  // If you are constructing your YARPC clients directly through the API,
    59  // create a gRPC dialer with the Compressor option.
    60  //
    61  //  trans := grpc.NewTransport()
    62  //  dialer := trans.NewDialer(grpc.Compressor(SnappyCompressor))
    63  //  peers := roundrobin.New(dialer)
    64  //  outbound := trans.NewOutbound(peers)
    65  //
    66  // If you are using the YARPC configurator to create YARPC objects
    67  // using config files, you will also need to register the compressor
    68  // with your configurator.
    69  //
    70  //  configurator := yarpcconfig.New()
    71  //  configurator.MustRegisterCompressor(SnappyCompressor)
    72  //
    73  // Then, using the compression strategy for outbound requests
    74  // on a particular client, just set the compressor to snappy.
    75  //
    76  //  outbounds:
    77  //    theirsecureservice:
    78  //      grpc:
    79  //        address: ":443"
    80  //        tls:
    81  //          enabled: true
    82  //        compressor: snappy
    83  //
    84  func New(...Option) *Compressor {
    85  	return &Compressor{}
    86  }
    87  
    88  // Compressor represents the snappy streaming compression strategy.
    89  type Compressor struct {
    90  	compressors   sync.Pool
    91  	decompressors sync.Pool
    92  }
    93  
    94  var _ transport.Compressor = (*Compressor)(nil)
    95  
    96  // Name is snappy.
    97  func (*Compressor) Name() string {
    98  	return "snappy"
    99  }
   100  
   101  // Compress creates a snappy compressor.
   102  func (c *Compressor) Compress(w io.Writer) (io.WriteCloser, error) {
   103  	if cw, got := c.compressors.Get().(*writer); got {
   104  		cw.writer.Reset(w)
   105  		return cw, nil
   106  	}
   107  	return &writer{
   108  		writer: snappy.NewBufferedWriter(w),
   109  		pool:   &c.compressors,
   110  	}, nil
   111  }
   112  
   113  type writer struct {
   114  	writer *snappy.Writer
   115  	pool   *sync.Pool
   116  }
   117  
   118  var _ io.WriteCloser = (*writer)(nil)
   119  
   120  func (w *writer) Write(buf []byte) (int, error) {
   121  	return w.writer.Write(buf)
   122  }
   123  
   124  func (w *writer) Close() error {
   125  	defer w.pool.Put(w)
   126  	return w.writer.Close()
   127  }
   128  
   129  // Decompress creates a snappy decompressor.
   130  func (c *Compressor) Decompress(r io.Reader) (io.ReadCloser, error) {
   131  	dr, got := c.decompressors.Get().(*reader)
   132  	if got {
   133  		dr.reader.Reset(r)
   134  		return dr, nil
   135  	}
   136  	return &reader{
   137  		reader: snappy.NewReader(r),
   138  		pool:   &c.decompressors,
   139  	}, nil
   140  }
   141  
   142  type reader struct {
   143  	reader *snappy.Reader
   144  	pool   *sync.Pool
   145  }
   146  
   147  var _ io.ReadCloser = (*reader)(nil)
   148  
   149  func (r *reader) Read(buf []byte) (n int, err error) {
   150  	return r.reader.Read(buf)
   151  }
   152  
   153  func (r *reader) Close() error {
   154  	r.pool.Put(r)
   155  	return nil
   156  }