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

     1  package gzip
     2  
     3  import (
     4  	"compress/gzip"
     5  	"io"
     6  	"sync"
     7  )
     8  
     9  var (
    10  	readerPool sync.Pool
    11  )
    12  
    13  // Codec is the implementation of a compress.Codec which supports creating
    14  // readers and writers for kafka messages compressed with gzip.
    15  type Codec struct {
    16  	// The compression level to configure on writers created by this codec.
    17  	// Acceptable values are defined in the standard gzip package.
    18  	//
    19  	// Default to gzip.DefaultCompressionLevel.
    20  	Level int
    21  
    22  	writerPool sync.Pool
    23  }
    24  
    25  // Code implements the compress.Codec interface.
    26  func (c *Codec) Code() int8 { return 1 }
    27  
    28  // Name implements the compress.Codec interface.
    29  func (c *Codec) Name() string { return "gzip" }
    30  
    31  // NewReader implements the compress.Codec interface.
    32  func (c *Codec) NewReader(r io.Reader) io.ReadCloser {
    33  	var err error
    34  	z, _ := readerPool.Get().(*gzip.Reader)
    35  	if z != nil {
    36  		err = z.Reset(r)
    37  	} else {
    38  		z, err = gzip.NewReader(r)
    39  	}
    40  	if err != nil {
    41  		if z != nil {
    42  			readerPool.Put(z)
    43  		}
    44  		return &errorReader{err: err}
    45  	}
    46  	return &reader{Reader: z}
    47  }
    48  
    49  // NewWriter implements the compress.Codec interface.
    50  func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
    51  	x := c.writerPool.Get()
    52  	z, _ := x.(*gzip.Writer)
    53  	if z == nil {
    54  		x, err := gzip.NewWriterLevel(w, c.level())
    55  		if err != nil {
    56  			return &errorWriter{err: err}
    57  		}
    58  		z = x
    59  	} else {
    60  		z.Reset(w)
    61  	}
    62  	return &writer{codec: c, Writer: z}
    63  }
    64  
    65  func (c *Codec) level() int {
    66  	if c.Level != 0 {
    67  		return c.Level
    68  	}
    69  	return gzip.DefaultCompression
    70  }
    71  
    72  type reader struct{ *gzip.Reader }
    73  
    74  func (r *reader) Close() (err error) {
    75  	if z := r.Reader; z != nil {
    76  		r.Reader = nil
    77  		err = z.Close()
    78  		// Pass it an empty reader, which is a zero-size value implementing the
    79  		// flate.Reader interface to avoid the construction of a bufio.Reader in
    80  		// the call to Reset.
    81  		//
    82  		// Note: we could also not reset the reader at all, but that would cause
    83  		// the underlying reader to be retained until the gzip.Reader is freed,
    84  		// which may not be desirable.
    85  		z.Reset(emptyReader{})
    86  		readerPool.Put(z)
    87  	}
    88  	return
    89  }
    90  
    91  type writer struct {
    92  	codec *Codec
    93  	*gzip.Writer
    94  }
    95  
    96  func (w *writer) Close() (err error) {
    97  	if z := w.Writer; z != nil {
    98  		w.Writer = nil
    99  		err = z.Close()
   100  		z.Reset(nil)
   101  		w.codec.writerPool.Put(z)
   102  	}
   103  	return
   104  }
   105  
   106  type emptyReader struct{}
   107  
   108  func (emptyReader) ReadByte() (byte, error) { return 0, io.EOF }
   109  
   110  func (emptyReader) Read([]byte) (int, error) { return 0, io.EOF }
   111  
   112  type errorReader struct{ err error }
   113  
   114  func (r *errorReader) Close() error { return r.err }
   115  
   116  func (r *errorReader) Read([]byte) (int, error) { return 0, r.err }
   117  
   118  type errorWriter struct{ err error }
   119  
   120  func (w *errorWriter) Close() error { return w.err }
   121  
   122  func (w *errorWriter) Write([]byte) (int, error) { return 0, w.err }