github.com/hoveychen/kafka-go@v0.4.42/compress/snappy/snappy.go (about)

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