github.com/QuangHoangHao/kafka-go@v0.4.36/compress/snappy/snappy.go (about)

     1  package snappy
     2  
     3  import (
     4  	"io"
     5  	"sync"
     6  
     7  	"github.com/klauspost/compress/snappy"
     8  )
     9  
    10  // Framing is an enumeration type used to enable or disable xerial framing of
    11  // snappy messages.
    12  type Framing int
    13  
    14  const (
    15  	Framed Framing = iota
    16  	Unframed
    17  )
    18  
    19  var (
    20  	readerPool sync.Pool
    21  	writerPool sync.Pool
    22  )
    23  
    24  // Codec is the implementation of a compress.Codec which supports creating
    25  // readers and writers for kafka messages compressed with snappy.
    26  type Codec struct {
    27  	// An optional framing to apply to snappy compression.
    28  	//
    29  	// Default to Framed.
    30  	Framing Framing
    31  }
    32  
    33  // Code implements the compress.Codec interface.
    34  func (c *Codec) Code() int8 { return 2 }
    35  
    36  // Name implements the compress.Codec interface.
    37  func (c *Codec) Name() string { return "snappy" }
    38  
    39  // NewReader implements the compress.Codec interface.
    40  func (c *Codec) NewReader(r io.Reader) io.ReadCloser {
    41  	x, _ := readerPool.Get().(*xerialReader)
    42  	if x != nil {
    43  		x.Reset(r)
    44  	} else {
    45  		x = &xerialReader{
    46  			reader: r,
    47  			decode: snappy.Decode,
    48  		}
    49  	}
    50  	return &reader{xerialReader: x}
    51  }
    52  
    53  // NewWriter implements the compress.Codec interface.
    54  func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
    55  	x, _ := writerPool.Get().(*xerialWriter)
    56  	if x != nil {
    57  		x.Reset(w)
    58  	} else {
    59  		x = &xerialWriter{
    60  			writer: w,
    61  			encode: snappy.Encode,
    62  		}
    63  	}
    64  	x.framed = c.Framing == Framed
    65  	return &writer{xerialWriter: x}
    66  }
    67  
    68  type reader struct{ *xerialReader }
    69  
    70  func (r *reader) Close() (err error) {
    71  	if x := r.xerialReader; x != nil {
    72  		r.xerialReader = nil
    73  		x.Reset(nil)
    74  		readerPool.Put(x)
    75  	}
    76  	return
    77  }
    78  
    79  type writer struct{ *xerialWriter }
    80  
    81  func (w *writer) Close() (err error) {
    82  	if x := w.xerialWriter; x != nil {
    83  		w.xerialWriter = nil
    84  		err = x.Flush()
    85  		x.Reset(nil)
    86  		writerPool.Put(x)
    87  	}
    88  	return
    89  }