github.com/hack0072008/kafka-go@v1.0.1/compress/zstd/zstd.go (about)

     1  // Package zstd implements Zstandard compression.
     2  package zstd
     3  
     4  import (
     5  	"io"
     6  	"runtime"
     7  	"sync"
     8  
     9  	"github.com/klauspost/compress/zstd"
    10  )
    11  
    12  // Codec is the implementation of a compress.Codec which supports creating
    13  // readers and writers for kafka messages compressed with zstd.
    14  type Codec struct {
    15  	// The compression level configured on writers created by the codec.
    16  	//
    17  	// Default to 3.
    18  	Level int
    19  
    20  	encoderPool sync.Pool // *encoder
    21  }
    22  
    23  // Code implements the compress.Codec interface.
    24  func (c *Codec) Code() int8 { return 4 }
    25  
    26  // Name implements the compress.Codec interface.
    27  func (c *Codec) Name() string { return "zstd" }
    28  
    29  // NewReader implements the compress.Codec interface.
    30  func (c *Codec) NewReader(r io.Reader) io.ReadCloser {
    31  	p := new(reader)
    32  	if dec, _ := decoderPool.Get().(*decoder); dec == nil {
    33  		z, err := zstd.NewReader(r)
    34  		if err != nil {
    35  			p.err = err
    36  		} else {
    37  			p.dec = &decoder{z}
    38  			// We need a finalizer because the reader spawns goroutines
    39  			// that will only be stopped if the Close method is called.
    40  			runtime.SetFinalizer(p.dec, (*decoder).finalize)
    41  		}
    42  	} else {
    43  		p.dec = dec
    44  		p.err = dec.Reset(r)
    45  	}
    46  	return p
    47  }
    48  
    49  func (c *Codec) level() int {
    50  	if c.Level != 0 {
    51  		return c.Level
    52  	}
    53  	return 3
    54  }
    55  
    56  func (c *Codec) zstdLevel() zstd.EncoderLevel {
    57  	return zstd.EncoderLevelFromZstd(c.level())
    58  }
    59  
    60  var decoderPool sync.Pool // *decoder
    61  
    62  type decoder struct {
    63  	*zstd.Decoder
    64  }
    65  
    66  func (d *decoder) finalize() {
    67  	d.Close()
    68  }
    69  
    70  type reader struct {
    71  	dec *decoder
    72  	err error
    73  }
    74  
    75  // Close implements the io.Closer interface.
    76  func (r *reader) Close() error {
    77  	if r.dec != nil {
    78  		r.dec.Reset(devNull{}) // don't retain the underlying reader
    79  		decoderPool.Put(r.dec)
    80  		r.dec = nil
    81  		r.err = io.ErrClosedPipe
    82  	}
    83  	return nil
    84  }
    85  
    86  // Read implements the io.Reader interface.
    87  func (r *reader) Read(p []byte) (int, error) {
    88  	if r.err != nil {
    89  		return 0, r.err
    90  	}
    91  	return r.dec.Read(p)
    92  }
    93  
    94  // WriteTo implements the io.WriterTo interface.
    95  func (r *reader) WriteTo(w io.Writer) (int64, error) {
    96  	if r.err != nil {
    97  		return 0, r.err
    98  	}
    99  	return r.dec.WriteTo(w)
   100  }
   101  
   102  // NewWriter implements the compress.Codec interface.
   103  func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
   104  	p := new(writer)
   105  	if enc, _ := c.encoderPool.Get().(*encoder); enc == nil {
   106  		z, err := zstd.NewWriter(w, zstd.WithEncoderLevel(c.zstdLevel()))
   107  		if err != nil {
   108  			p.err = err
   109  		} else {
   110  			p.enc = &encoder{z}
   111  			// We need a finalizer because the writer spawns goroutines
   112  			// that will only be stopped if the Close method is called.
   113  			runtime.SetFinalizer(p.enc, (*encoder).finalize)
   114  		}
   115  	} else {
   116  		p.enc = enc
   117  		p.enc.Reset(w)
   118  	}
   119  	p.c = c
   120  	return p
   121  }
   122  
   123  type encoder struct {
   124  	*zstd.Encoder
   125  }
   126  
   127  func (e *encoder) finalize() {
   128  	e.Close()
   129  }
   130  
   131  type writer struct {
   132  	c   *Codec
   133  	enc *encoder
   134  	err error
   135  }
   136  
   137  // Close implements the io.Closer interface.
   138  func (w *writer) Close() error {
   139  	if w.enc != nil {
   140  		// Close needs to be called to write the end of stream marker and flush
   141  		// the buffers. The zstd package documents that the encoder is re-usable
   142  		// after being closed.
   143  		err := w.enc.Close()
   144  		if err != nil {
   145  			w.err = err
   146  		}
   147  		w.enc.Reset(devNull{}) // don't retain the underlying writer
   148  		w.c.encoderPool.Put(w.enc)
   149  		w.enc = nil
   150  		return err
   151  	}
   152  	return nil
   153  }
   154  
   155  // WriteTo implements the io.WriterTo interface.
   156  func (w *writer) Write(p []byte) (int, error) {
   157  	if w.err != nil {
   158  		return 0, w.err
   159  	}
   160  	if w.enc == nil {
   161  		return 0, io.ErrClosedPipe
   162  	}
   163  	return w.enc.Write(p)
   164  }
   165  
   166  // ReadFrom implements the io.ReaderFrom interface.
   167  func (w *writer) ReadFrom(r io.Reader) (int64, error) {
   168  	if w.err != nil {
   169  		return 0, w.err
   170  	}
   171  	if w.enc == nil {
   172  		return 0, io.ErrClosedPipe
   173  	}
   174  	return w.enc.ReadFrom(r)
   175  }
   176  
   177  type devNull struct{}
   178  
   179  func (devNull) Read([]byte) (int, error)  { return 0, io.EOF }
   180  func (devNull) Write([]byte) (int, error) { return 0, nil }