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 }