forked from nytimes/gziphandler
/
zstd.go
68 lines (57 loc) · 1.17 KB
/
zstd.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
package zstd
import (
"fmt"
"io"
"sync"
"github.com/CAFxX/httpcompression/contrib/internal/utils"
"github.com/klauspost/compress/zstd"
)
const (
Encoding = "zstd"
DefaultCompression = zstd.SpeedDefault
)
type compressor struct {
pool sync.Pool
opts []zstd.EOption
}
func New(opts ...zstd.EOption) (c *compressor, err error) {
defer func() {
if r := recover(); r != nil {
c, err = nil, fmt.Errorf("panic: %v", r)
}
}()
opts = append([]zstd.EOption(nil), opts...)
tw, err := zstd.NewWriter(io.Discard, opts...)
if err != nil {
return nil, err
}
if err := utils.CheckWriter(tw); err != nil {
return nil, fmt.Errorf("zstd: writer initialization: %w", err)
}
c = &compressor{opts: opts}
return c, nil
}
func (c *compressor) Get(w io.Writer) io.WriteCloser {
if gw, ok := c.pool.Get().(*zstdWriter); ok {
gw.Reset(w)
return gw
}
gw, err := zstd.NewWriter(w, c.opts...)
if err != nil {
return utils.ErrorWriteCloser{Err: err}
}
return &zstdWriter{
Encoder: gw,
c: c,
}
}
type zstdWriter struct {
*zstd.Encoder
c *compressor
}
func (w *zstdWriter) Close() error {
err := w.Encoder.Close()
w.Reset(nil)
w.c.pool.Put(w)
return err
}