cubefs

1
// Copyright 2019+ Klaus Post. All rights reserved.
2
// License information can be found in the LICENSE file.
3

4
package zstd
5

6
import (
7
	"errors"
8
	"io"
9
	"sync"
10
)
11

12
// ZipMethodWinZip is the method for Zstandard compressed data inside Zip files for WinZip.
13
// See https://www.winzip.com/win/en/comp_info.html
14
const ZipMethodWinZip = 93
15

16
// ZipMethodPKWare is the original method number used by PKWARE to indicate Zstandard compression.
17
// Deprecated: This has been deprecated by PKWARE, use ZipMethodWinZip instead for compression.
18
// See https://pkware.cachefly.net/webdocs/APPNOTE/APPNOTE-6.3.9.TXT
19
const ZipMethodPKWare = 20
20

21
var zipReaderPool sync.Pool
22

23
// newZipReader cannot be used since we would leak goroutines...
24
func newZipReader(r io.Reader) io.ReadCloser {
25
	dec, ok := zipReaderPool.Get().(*Decoder)
26
	if ok {
27
		dec.Reset(r)
28
	} else {
29
		d, err := NewReader(r, WithDecoderConcurrency(1), WithDecoderLowmem(true))
30
		if err != nil {
31
			panic(err)
32
		}
33
		dec = d
34
	}
35
	return &pooledZipReader{dec: dec}
36
}
37

38
type pooledZipReader struct {
39
	mu  sync.Mutex // guards Close and Read
40
	dec *Decoder
41
}
42

43
func (r *pooledZipReader) Read(p []byte) (n int, err error) {
44
	r.mu.Lock()
45
	defer r.mu.Unlock()
46
	if r.dec == nil {
47
		return 0, errors.New("Read after Close")
48
	}
49
	dec, err := r.dec.Read(p)
50

51
	return dec, err
52
}
53

54
func (r *pooledZipReader) Close() error {
55
	r.mu.Lock()
56
	defer r.mu.Unlock()
57
	var err error
58
	if r.dec != nil {
59
		err = r.dec.Reset(nil)
60
		zipReaderPool.Put(r.dec)
61
		r.dec = nil
62
	}
63
	return err
64
}
65

66
type pooledZipWriter struct {
67
	mu   sync.Mutex // guards Close and Read
68
	enc  *Encoder
69
	pool *sync.Pool
70
}
71

72
func (w *pooledZipWriter) Write(p []byte) (n int, err error) {
73
	w.mu.Lock()
74
	defer w.mu.Unlock()
75
	if w.enc == nil {
76
		return 0, errors.New("Write after Close")
77
	}
78
	return w.enc.Write(p)
79
}
80

81
func (w *pooledZipWriter) Close() error {
82
	w.mu.Lock()
83
	defer w.mu.Unlock()
84
	var err error
85
	if w.enc != nil {
86
		err = w.enc.Close()
87
		w.pool.Put(w.enc)
88
		w.enc = nil
89
	}
90
	return err
91
}
92

93
// ZipCompressor returns a compressor that can be registered with zip libraries.
94
// The provided encoder options will be used on all encodes.
95
func ZipCompressor(opts ...EOption) func(w io.Writer) (io.WriteCloser, error) {
96
	var pool sync.Pool
97
	return func(w io.Writer) (io.WriteCloser, error) {
98
		enc, ok := pool.Get().(*Encoder)
99
		if ok {
100
			enc.Reset(w)
101
		} else {
102
			var err error
103
			enc, err = NewWriter(w, opts...)
104
			if err != nil {
105
				return nil, err
106
			}
107
		}
108
		return &pooledZipWriter{enc: enc, pool: &pool}, nil
109
	}
110
}
111

112
// ZipDecompressor returns a decompressor that can be registered with zip libraries.
113
// See ZipCompressor for example.
114
func ZipDecompressor() func(r io.Reader) io.ReadCloser {
115
	return func(r io.Reader) io.ReadCloser {
116
		d, err := NewReader(r, WithDecoderConcurrency(1), WithDecoderLowmem(true))
117
		if err != nil {
118
			panic(err)
119
		}
120
		return d.IOReadCloser()
121
	}
122
}
123