blob: 35a3cda91407787066c7d6ba4d3f8cd17cbc8e87 [file] [log] [blame]
Pragya Arya324337e2020-02-20 14:35:08 +05301// Copyright 2019+ Klaus Post. All rights reserved.
2// License information can be found in the LICENSE file.
3// Based on work by Yann Collet, released under BSD License.
4
5package zstd
6
7import (
8 "bytes"
9 "errors"
10 "io"
11 "sync"
12)
13
14// Decoder provides decoding of zstandard streams.
15// The decoder has been designed to operate without allocations after a warmup.
16// This means that you should store the decoder for best performance.
17// To re-use a stream decoder, use the Reset(r io.Reader) error to switch to another stream.
18// A decoder can safely be re-used even if the previous stream failed.
19// To release the resources, you must call the Close() function on a decoder.
20type Decoder struct {
21 o decoderOptions
22
23 // Unreferenced decoders, ready for use.
24 decoders chan *blockDec
25
26 // Unreferenced decoders, ready for use.
27 frames chan *frameDec
28
29 // Streams ready to be decoded.
30 stream chan decodeStream
31
32 // Current read position used for Reader functionality.
33 current decoderState
34
35 // Custom dictionaries
36 dicts map[uint32]struct{}
37
38 // streamWg is the waitgroup for all streams
39 streamWg sync.WaitGroup
40}
41
42// decoderState is used for maintaining state when the decoder
43// is used for streaming.
44type decoderState struct {
45 // current block being written to stream.
46 decodeOutput
47
48 // output in order to be written to stream.
49 output chan decodeOutput
50
51 // cancel remaining output.
52 cancel chan struct{}
53
54 flushed bool
55}
56
57var (
58 // Check the interfaces we want to support.
59 _ = io.WriterTo(&Decoder{})
60 _ = io.Reader(&Decoder{})
61)
62
63// NewReader creates a new decoder.
64// A nil Reader can be provided in which case Reset can be used to start a decode.
65//
66// A Decoder can be used in two modes:
67//
68// 1) As a stream, or
69// 2) For stateless decoding using DecodeAll or DecodeBuffer.
70//
71// Only a single stream can be decoded concurrently, but the same decoder
72// can run multiple concurrent stateless decodes. It is even possible to
73// use stateless decodes while a stream is being decoded.
74//
75// The Reset function can be used to initiate a new stream, which is will considerably
76// reduce the allocations normally caused by NewReader.
77func NewReader(r io.Reader, opts ...DOption) (*Decoder, error) {
78 initPredefined()
79 var d Decoder
80 d.o.setDefault()
81 for _, o := range opts {
82 err := o(&d.o)
83 if err != nil {
84 return nil, err
85 }
86 }
87 d.current.output = make(chan decodeOutput, d.o.concurrent)
88 d.current.flushed = true
89
90 // Create decoders
91 d.decoders = make(chan *blockDec, d.o.concurrent)
92 d.frames = make(chan *frameDec, d.o.concurrent)
93 for i := 0; i < d.o.concurrent; i++ {
94 d.frames <- newFrameDec(d.o)
95 d.decoders <- newBlockDec(d.o.lowMem)
96 }
97
98 if r == nil {
99 return &d, nil
100 }
101 return &d, d.Reset(r)
102}
103
104// Read bytes from the decompressed stream into p.
105// Returns the number of bytes written and any error that occurred.
106// When the stream is done, io.EOF will be returned.
107func (d *Decoder) Read(p []byte) (int, error) {
108 if d.stream == nil {
109 return 0, errors.New("no input has been initialized")
110 }
111 var n int
112 for {
113 if len(d.current.b) > 0 {
114 filled := copy(p, d.current.b)
115 p = p[filled:]
116 d.current.b = d.current.b[filled:]
117 n += filled
118 }
119 if len(p) == 0 {
120 break
121 }
122 if len(d.current.b) == 0 {
123 // We have an error and no more data
124 if d.current.err != nil {
125 break
126 }
127 if !d.nextBlock(n == 0) {
128 return n, nil
129 }
130 }
131 }
132 if len(d.current.b) > 0 {
133 if debug {
134 println("returning", n, "still bytes left:", len(d.current.b))
135 }
136 // Only return error at end of block
137 return n, nil
138 }
139 if d.current.err != nil {
140 d.drainOutput()
141 }
142 if debug {
143 println("returning", n, d.current.err, len(d.decoders))
144 }
145 return n, d.current.err
146}
147
148// Reset will reset the decoder the supplied stream after the current has finished processing.
149// Note that this functionality cannot be used after Close has been called.
150func (d *Decoder) Reset(r io.Reader) error {
151 if d.current.err == ErrDecoderClosed {
152 return d.current.err
153 }
154 if r == nil {
155 return errors.New("nil Reader sent as input")
156 }
157
158 if d.stream == nil {
159 d.stream = make(chan decodeStream, 1)
160 d.streamWg.Add(1)
161 go d.startStreamDecoder(d.stream)
162 }
163
164 d.drainOutput()
165
166 // If bytes buffer and < 1MB, do sync decoding anyway.
167 if bb, ok := r.(*bytes.Buffer); ok && bb.Len() < 1<<20 {
168 if debug {
169 println("*bytes.Buffer detected, doing sync decode, len:", bb.Len())
170 }
171 b := bb.Bytes()
172 dst, err := d.DecodeAll(b, nil)
173 if err == nil {
174 err = io.EOF
175 }
176 d.current.b = dst
177 d.current.err = err
178 d.current.flushed = true
179 if debug {
180 println("sync decode to ", len(dst), "bytes, err:", err)
181 }
182 return nil
183 }
184
185 // Remove current block.
186 d.current.decodeOutput = decodeOutput{}
187 d.current.err = nil
188 d.current.cancel = make(chan struct{})
189 d.current.flushed = false
190 d.current.d = nil
191
192 d.stream <- decodeStream{
193 r: r,
194 output: d.current.output,
195 cancel: d.current.cancel,
196 }
197 return nil
198}
199
200// drainOutput will drain the output until errEndOfStream is sent.
201func (d *Decoder) drainOutput() {
202 if d.current.cancel != nil {
203 println("cancelling current")
204 close(d.current.cancel)
205 d.current.cancel = nil
206 }
207 if d.current.d != nil {
208 if debug {
209 printf("re-adding current decoder %p, decoders: %d", d.current.d, len(d.decoders))
210 }
211 d.decoders <- d.current.d
212 d.current.d = nil
213 d.current.b = nil
214 }
215 if d.current.output == nil || d.current.flushed {
216 println("current already flushed")
217 return
218 }
219 for {
220 select {
221 case v := <-d.current.output:
222 if v.d != nil {
223 if debug {
224 printf("re-adding decoder %p", v.d)
225 }
226 d.decoders <- v.d
227 }
228 if v.err == errEndOfStream {
229 println("current flushed")
230 d.current.flushed = true
231 return
232 }
233 }
234 }
235}
236
237// WriteTo writes data to w until there's no more data to write or when an error occurs.
238// The return value n is the number of bytes written.
239// Any error encountered during the write is also returned.
240func (d *Decoder) WriteTo(w io.Writer) (int64, error) {
241 if d.stream == nil {
242 return 0, errors.New("no input has been initialized")
243 }
244 var n int64
245 for {
246 if len(d.current.b) > 0 {
247 n2, err2 := w.Write(d.current.b)
248 n += int64(n2)
249 if err2 != nil && d.current.err == nil {
250 d.current.err = err2
251 break
252 }
253 }
254 if d.current.err != nil {
255 break
256 }
257 d.nextBlock(true)
258 }
259 err := d.current.err
260 if err != nil {
261 d.drainOutput()
262 }
263 if err == io.EOF {
264 err = nil
265 }
266 return n, err
267}
268
269// DecodeAll allows stateless decoding of a blob of bytes.
270// Output will be appended to dst, so if the destination size is known
271// you can pre-allocate the destination slice to avoid allocations.
272// DecodeAll can be used concurrently.
273// The Decoder concurrency limits will be respected.
274func (d *Decoder) DecodeAll(input, dst []byte) ([]byte, error) {
275 if d.current.err == ErrDecoderClosed {
276 return dst, ErrDecoderClosed
277 }
278
279 // Grab a block decoder and frame decoder.
280 block, frame := <-d.decoders, <-d.frames
281 defer func() {
282 if debug {
283 printf("re-adding decoder: %p", block)
284 }
285 d.decoders <- block
286 frame.rawInput = nil
287 frame.bBuf = nil
288 d.frames <- frame
289 }()
290 frame.bBuf = input
291
292 for {
293 err := frame.reset(&frame.bBuf)
294 if err == io.EOF {
295 return dst, nil
296 }
297 if err != nil {
298 return dst, err
299 }
300 if frame.FrameContentSize > d.o.maxDecodedSize-uint64(len(dst)) {
301 return dst, ErrDecoderSizeExceeded
302 }
303 if frame.FrameContentSize > 0 && frame.FrameContentSize < 1<<30 {
304 // Never preallocate moe than 1 GB up front.
305 if uint64(cap(dst)) < frame.FrameContentSize {
306 dst2 := make([]byte, len(dst), len(dst)+int(frame.FrameContentSize))
307 copy(dst2, dst)
308 dst = dst2
309 }
310 }
311 if cap(dst) == 0 {
312 // Allocate window size * 2 by default if nothing is provided and we didn't get frame content size.
313 size := frame.WindowSize * 2
314 // Cap to 1 MB.
315 if size > 1<<20 {
316 size = 1 << 20
317 }
318 dst = make([]byte, 0, frame.WindowSize)
319 }
320
321 dst, err = frame.runDecoder(dst, block)
322 if err != nil {
323 return dst, err
324 }
325 if len(frame.bBuf) == 0 {
326 break
327 }
328 }
329 return dst, nil
330}
331
332// nextBlock returns the next block.
333// If an error occurs d.err will be set.
334// Optionally the function can block for new output.
335// If non-blocking mode is used the returned boolean will be false
336// if no data was available without blocking.
337func (d *Decoder) nextBlock(blocking bool) (ok bool) {
338 if d.current.d != nil {
339 if debug {
340 printf("re-adding current decoder %p", d.current.d)
341 }
342 d.decoders <- d.current.d
343 d.current.d = nil
344 }
345 if d.current.err != nil {
346 // Keep error state.
347 return blocking
348 }
349
350 if blocking {
351 d.current.decodeOutput = <-d.current.output
352 } else {
353 select {
354 case d.current.decodeOutput = <-d.current.output:
355 default:
356 return false
357 }
358 }
359 if debug {
360 println("got", len(d.current.b), "bytes, error:", d.current.err)
361 }
362 return true
363}
364
365// Close will release all resources.
366// It is NOT possible to reuse the decoder after this.
367func (d *Decoder) Close() {
368 if d.current.err == ErrDecoderClosed {
369 return
370 }
371 d.drainOutput()
372 if d.stream != nil {
373 close(d.stream)
374 d.streamWg.Wait()
375 d.stream = nil
376 }
377 if d.decoders != nil {
378 close(d.decoders)
379 for dec := range d.decoders {
380 dec.Close()
381 }
382 d.decoders = nil
383 }
384 if d.current.d != nil {
385 d.current.d.Close()
386 d.current.d = nil
387 }
388 d.current.err = ErrDecoderClosed
389}
390
391// IOReadCloser returns the decoder as an io.ReadCloser for convenience.
392// Any changes to the decoder will be reflected, so the returned ReadCloser
393// can be reused along with the decoder.
394// io.WriterTo is also supported by the returned ReadCloser.
395func (d *Decoder) IOReadCloser() io.ReadCloser {
396 return closeWrapper{d: d}
397}
398
399// closeWrapper wraps a function call as a closer.
400type closeWrapper struct {
401 d *Decoder
402}
403
404// WriteTo forwards WriteTo calls to the decoder.
405func (c closeWrapper) WriteTo(w io.Writer) (n int64, err error) {
406 return c.d.WriteTo(w)
407}
408
409// Read forwards read calls to the decoder.
410func (c closeWrapper) Read(p []byte) (n int, err error) {
411 return c.d.Read(p)
412}
413
414// Close closes the decoder.
415func (c closeWrapper) Close() error {
416 c.d.Close()
417 return nil
418}
419
420type decodeOutput struct {
421 d *blockDec
422 b []byte
423 err error
424}
425
426type decodeStream struct {
427 r io.Reader
428
429 // Blocks ready to be written to output.
430 output chan decodeOutput
431
432 // cancel reading from the input
433 cancel chan struct{}
434}
435
436// errEndOfStream indicates that everything from the stream was read.
437var errEndOfStream = errors.New("end-of-stream")
438
439// Create Decoder:
440// Spawn n block decoders. These accept tasks to decode a block.
441// Create goroutine that handles stream processing, this will send history to decoders as they are available.
442// Decoders update the history as they decode.
443// When a block is returned:
444// a) history is sent to the next decoder,
445// b) content written to CRC.
446// c) return data to WRITER.
447// d) wait for next block to return data.
448// Once WRITTEN, the decoders reused by the writer frame decoder for re-use.
449func (d *Decoder) startStreamDecoder(inStream chan decodeStream) {
450 defer d.streamWg.Done()
451 frame := newFrameDec(d.o)
452 for stream := range inStream {
453 if debug {
454 println("got new stream")
455 }
456 br := readerWrapper{r: stream.r}
457 decodeStream:
458 for {
459 err := frame.reset(&br)
460 if debug && err != nil {
461 println("Frame decoder returned", err)
462 }
463 if err != nil {
464 stream.output <- decodeOutput{
465 err: err,
466 }
467 break
468 }
469 if debug {
470 println("starting frame decoder")
471 }
472
473 // This goroutine will forward history between frames.
474 frame.frameDone.Add(1)
475 frame.initAsync()
476
477 go frame.startDecoder(stream.output)
478 decodeFrame:
479 // Go through all blocks of the frame.
480 for {
481 dec := <-d.decoders
482 select {
483 case <-stream.cancel:
484 if !frame.sendErr(dec, io.EOF) {
485 // To not let the decoder dangle, send it back.
486 stream.output <- decodeOutput{d: dec}
487 }
488 break decodeStream
489 default:
490 }
491 err := frame.next(dec)
492 switch err {
493 case io.EOF:
494 // End of current frame, no error
495 println("EOF on next block")
496 break decodeFrame
497 case nil:
498 continue
499 default:
500 println("block decoder returned", err)
501 break decodeStream
502 }
503 }
504 // All blocks have started decoding, check if there are more frames.
505 println("waiting for done")
506 frame.frameDone.Wait()
507 println("done waiting...")
508 }
509 frame.frameDone.Wait()
510 println("Sending EOS")
511 stream.output <- decodeOutput{err: errEndOfStream}
512 }
513}