github.com/segmentio/kafka-go@v0.4.48-0.20240318174348-3f6244eb34fd/compress/gzip/gzip.go (about)

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