71 lines
1.6 KiB
Go
71 lines
1.6 KiB
Go
|
package ws
|
|||
|
|
|
|||
|
|
import (
|
|||
|
|
"bytes"
|
|||
|
|
"compress/gzip"
|
|||
|
|
"io"
|
|||
|
|
"sync"
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
const gzipMagicByte = 0x1F
|
|||
|
|
|
|||
|
|
// Compressor 管理 gzip Writer/Reader 的 sync.Pool 复用
|
|||
|
|
type Compressor struct {
|
|||
|
|
writerPool sync.Pool
|
|||
|
|
readerPool sync.Pool
|
|||
|
|
level int
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// NewCompressor 创建 Compressor（level: gzip.DefaultCompression/BestSpeed/BestCompression 等）
|
|||
|
|
func NewCompressor(level int) *Compressor {
|
|||
|
|
return &Compressor{
|
|||
|
|
level: level,
|
|||
|
|
writerPool: sync.Pool{
|
|||
|
|
New: func() any {
|
|||
|
|
w, _ := gzip.NewWriterLevel(io.Discard, level)
|
|||
|
|
return w
|
|||
|
|
},
|
|||
|
|
},
|
|||
|
|
readerPool: sync.Pool{
|
|||
|
|
New: func() any { return new(gzip.Reader) },
|
|||
|
|
},
|
|||
|
|
}
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Compress 压缩数据，返回 gzip 字节
|
|||
|
|
func (c *Compressor) Compress(data []byte) ([]byte, error) {
|
|||
|
|
buf := &bytes.Buffer{}
|
|||
|
|
w := c.writerPool.Get().(*gzip.Writer)
|
|||
|
|
w.Reset(buf)
|
|||
|
|
if _, err := w.Write(data); err != nil {
|
|||
|
|
c.writerPool.Put(w)
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
if err := w.Close(); err != nil {
|
|||
|
|
c.writerPool.Put(w)
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
c.writerPool.Put(w)
|
|||
|
|
return buf.Bytes(), nil
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// Decompress 解压 gzip 数据
|
|||
|
|
func (c *Compressor) Decompress(data []byte) ([]byte, error) {
|
|||
|
|
r := c.readerPool.Get().(*gzip.Reader)
|
|||
|
|
if err := r.Reset(bytes.NewReader(data)); err != nil {
|
|||
|
|
c.readerPool.Put(r)
|
|||
|
|
return nil, err
|
|||
|
|
}
|
|||
|
|
out, err := io.ReadAll(r)
|
|||
|
|
r.Close()
|
|||
|
|
c.readerPool.Put(r)
|
|||
|
|
return out, err
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// IsCompressed 检查数据是否以 gzip magic byte 开头
|
|||
|
|
func IsCompressed(data []byte) bool {
|
|||
|
|
return len(data) > 0 && data[0] == gzipMagicByte
|
|||
|
|
}
|
|||
|
|
|
|||
|
|
// DefaultCompressor 全局默认压缩器（gzip.DefaultCompression = level 6）
|
|||
|
|
var DefaultCompressor = NewCompressor(gzip.DefaultCompression)
|