mirror of
https://github.com/prometheus/prometheus.git
synced 2025-03-05 20:59:13 -08:00
603 lines
14 KiB
Go
603 lines
14 KiB
Go
package remote
|
|
|
|
import (
|
|
"bytes"
|
|
"compress/flate"
|
|
"compress/gzip"
|
|
"compress/lzw"
|
|
"io"
|
|
"sync"
|
|
|
|
reS2 "github.com/klauspost/compress/s2"
|
|
reSnappy "github.com/klauspost/compress/snappy"
|
|
"github.com/klauspost/compress/zstd"
|
|
reZstd "github.com/klauspost/compress/zstd"
|
|
|
|
"github.com/andybalholm/brotli"
|
|
"github.com/golang/snappy"
|
|
)
|
|
|
|
type Compression interface {
|
|
Compress(data []byte) ([]byte, error)
|
|
Decompress(data []byte) ([]byte, error)
|
|
}
|
|
|
|
// hacky globals to easily tweak the compression algorithm and run some benchmarks
|
|
type CompAlgorithm int
|
|
|
|
var UseAlgorithm = Snappy
|
|
|
|
const (
|
|
Snappy CompAlgorithm = iota
|
|
SnappyAlt
|
|
S2
|
|
ZstdFast
|
|
ZstdDefault
|
|
ZstdBestComp
|
|
GzipFast
|
|
GzipComp
|
|
Lzw
|
|
FlateFast
|
|
FlateComp
|
|
BrotliFast
|
|
BrotliComp
|
|
BrotliDefault
|
|
)
|
|
|
|
// sync.Pool-ed createComp
|
|
var compPool = sync.Pool{
|
|
// New optionally specifies a function to generate
|
|
// a value when Get would otherwise return nil.
|
|
New: func() interface{} { return createComp() },
|
|
}
|
|
|
|
func GetPooledComp() Compression {
|
|
return compPool.Get().(Compression)
|
|
}
|
|
|
|
func PutPooledComp(c Compression) {
|
|
compPool.Put(c)
|
|
}
|
|
|
|
var createComp func() Compression = func() Compression {
|
|
switch UseAlgorithm {
|
|
case Snappy:
|
|
return &snappyCompression{}
|
|
case SnappyAlt:
|
|
return &snappyAltCompression{}
|
|
case S2:
|
|
return &s2Compression{}
|
|
case ZstdDefault:
|
|
return &zstdCompression{level: zstd.SpeedDefault}
|
|
case ZstdFast:
|
|
return &zstdCompression{level: zstd.SpeedFastest}
|
|
case ZstdBestComp:
|
|
return &zstdCompression{level: zstd.SpeedBestCompression}
|
|
case GzipFast:
|
|
return &gzipCompression{level: gzip.BestSpeed}
|
|
case GzipComp:
|
|
return &gzipCompression{level: gzip.BestCompression}
|
|
case Lzw:
|
|
return &lzwCompression{}
|
|
case FlateFast:
|
|
return &flateCompression{level: flate.BestSpeed}
|
|
case FlateComp:
|
|
return &flateCompression{level: flate.BestCompression}
|
|
case BrotliFast:
|
|
return &brotliCompression{quality: brotli.BestSpeed}
|
|
case BrotliDefault:
|
|
return &brotliCompression{quality: brotli.DefaultCompression}
|
|
case BrotliComp:
|
|
return &brotliCompression{quality: brotli.BestCompression}
|
|
default:
|
|
panic("unknown compression algorithm")
|
|
}
|
|
}
|
|
|
|
type noopCompression struct{}
|
|
|
|
func (n *noopCompression) Compress(data []byte) ([]byte, error) {
|
|
return data, nil
|
|
}
|
|
|
|
func (n *noopCompression) Decompress(data []byte) ([]byte, error) {
|
|
return data, nil
|
|
}
|
|
|
|
type snappyCompression struct {
|
|
buf []byte
|
|
}
|
|
|
|
func (s *snappyCompression) Compress(data []byte) ([]byte, error) {
|
|
s.buf = s.buf[0:cap(s.buf)]
|
|
compressed := snappy.Encode(s.buf, data)
|
|
if n := snappy.MaxEncodedLen(len(data)); n > len(s.buf) {
|
|
s.buf = make([]byte, n)
|
|
}
|
|
return compressed, nil
|
|
}
|
|
func (s *snappyCompression) Decompress(data []byte) ([]byte, error) {
|
|
s.buf = s.buf[0:cap(s.buf)]
|
|
uncompressed, err := snappy.Decode(s.buf, data)
|
|
if len(uncompressed) > cap(s.buf) {
|
|
s.buf = uncompressed
|
|
}
|
|
return uncompressed, err
|
|
}
|
|
|
|
type snappyAltCompression struct {
|
|
buf []byte
|
|
}
|
|
|
|
func (s *snappyAltCompression) Compress(data []byte) ([]byte, error) {
|
|
s.buf = s.buf[0:cap(s.buf)]
|
|
res := reSnappy.Encode(s.buf, data)
|
|
if n := reSnappy.MaxEncodedLen(len(data)); n > len(s.buf) {
|
|
s.buf = make([]byte, n)
|
|
}
|
|
return res, nil
|
|
}
|
|
func (s *snappyAltCompression) Decompress(data []byte) ([]byte, error) {
|
|
s.buf = s.buf[0:cap(s.buf)]
|
|
uncompressed, err := reSnappy.Decode(s.buf, data)
|
|
if len(uncompressed) > cap(s.buf) {
|
|
s.buf = uncompressed
|
|
}
|
|
return uncompressed, err
|
|
}
|
|
|
|
type s2Compression struct {
|
|
buf []byte
|
|
}
|
|
|
|
func (s *s2Compression) Compress(data []byte) ([]byte, error) {
|
|
res := reS2.Encode(s.buf, data)
|
|
if n := reS2.MaxEncodedLen(len(data)); n > len(s.buf) {
|
|
s.buf = make([]byte, n)
|
|
}
|
|
return res, nil
|
|
}
|
|
|
|
func (s *s2Compression) Decompress(data []byte) ([]byte, error) {
|
|
s.buf = s.buf[0:cap(s.buf)]
|
|
uncompressed, err := reS2.Decode(s.buf, data)
|
|
if len(uncompressed) > cap(s.buf) {
|
|
s.buf = uncompressed
|
|
}
|
|
return uncompressed, err
|
|
}
|
|
|
|
type zstdCompression struct {
|
|
level zstd.EncoderLevel
|
|
buf []byte
|
|
w *reZstd.Encoder
|
|
}
|
|
|
|
func (z *zstdCompression) Compress(data []byte) ([]byte, error) {
|
|
var err error
|
|
if z.w == nil {
|
|
// TODO: should be initialized on creation
|
|
z.w, err = reZstd.NewWriter(nil, reZstd.WithEncoderLevel(reZstd.EncoderLevel(z.level)))
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
z.buf = z.buf[:0]
|
|
writer := bytes.NewBuffer(z.buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
z.w.Reset(writer)
|
|
|
|
z.w.Write(data)
|
|
err = z.w.Close()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
res := writer.Bytes()
|
|
if len(res) > cap(z.buf) {
|
|
z.buf = res
|
|
}
|
|
return res, nil
|
|
}
|
|
|
|
func (z *zstdCompression) Decompress(data []byte) ([]byte, error) {
|
|
decoder, err := reZstd.NewReader(nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
z.buf = z.buf[:0]
|
|
buf, err := decoder.DecodeAll(data, z.buf)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(buf) > cap(z.buf) {
|
|
z.buf = buf
|
|
}
|
|
return buf, nil
|
|
}
|
|
|
|
type gzipCompression struct {
|
|
level int
|
|
buf []byte
|
|
w *gzip.Writer
|
|
}
|
|
|
|
func (g *gzipCompression) Compress(data []byte) ([]byte, error) {
|
|
var err error
|
|
if g.w == nil {
|
|
g.w, err = gzip.NewWriterLevel(nil, g.level)
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
g.buf = g.buf[:0]
|
|
buf := bytes.NewBuffer(g.buf)
|
|
g.w.Reset(buf)
|
|
_, err = g.w.Write(data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
err = g.w.Close()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(buf.Bytes()) > cap(g.buf) {
|
|
g.buf = buf.Bytes()
|
|
}
|
|
return buf.Bytes(), nil
|
|
}
|
|
|
|
func (g *gzipCompression) Decompress(data []byte) ([]byte, error) {
|
|
r := bytes.NewReader(data)
|
|
gzReader, err := gzip.NewReader(r)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
defer gzReader.Close()
|
|
decompressedData, err := io.ReadAll(gzReader)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
return decompressedData, nil
|
|
|
|
// TODO: debug this
|
|
// r := bytes.NewReader(data)
|
|
// var err error
|
|
// if g.r == nil {
|
|
// g.r, err = gzip.NewReader(r)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// }
|
|
// g.r.Reset(r)
|
|
// defer g.r.Close()
|
|
// return io.ReadAll(g.r)
|
|
}
|
|
|
|
type lzwCompression struct {
|
|
buf []byte
|
|
w *lzw.Writer
|
|
}
|
|
|
|
func (l *lzwCompression) Compress(data []byte) ([]byte, error) {
|
|
if l.w == nil {
|
|
l.w = lzw.NewWriter(nil, lzw.LSB, 8).(*lzw.Writer)
|
|
}
|
|
compressed := bytes.NewBuffer(l.buf)
|
|
l.w.Reset(compressed, lzw.LSB, 8)
|
|
_, err := l.w.Write(data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
err = l.w.Close()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(compressed.Bytes()) > cap(l.buf) {
|
|
l.buf = compressed.Bytes()
|
|
}
|
|
return compressed.Bytes(), nil
|
|
}
|
|
|
|
func (l *lzwCompression) Decompress(data []byte) ([]byte, error) {
|
|
reader := bytes.NewReader(data)
|
|
r := lzw.NewReader(reader, lzw.LSB, 8)
|
|
defer r.Close()
|
|
return io.ReadAll(r)
|
|
}
|
|
|
|
type flateCompression struct {
|
|
level int
|
|
buf []byte
|
|
w *flate.Writer
|
|
}
|
|
|
|
func (f *flateCompression) Compress(data []byte) ([]byte, error) {
|
|
var err error
|
|
if f.w == nil {
|
|
f.w, err = flate.NewWriter(nil, f.level)
|
|
}
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
f.buf = f.buf[:0]
|
|
compressed := bytes.NewBuffer(f.buf)
|
|
f.w.Reset(compressed)
|
|
_, err = f.w.Write(data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
err = f.w.Close()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(compressed.Bytes()) > cap(f.buf) {
|
|
f.buf = compressed.Bytes()
|
|
}
|
|
return compressed.Bytes(), nil
|
|
}
|
|
|
|
func (f *flateCompression) Decompress(data []byte) ([]byte, error) {
|
|
reader := bytes.NewReader(data)
|
|
r := flate.NewReader(reader)
|
|
defer r.Close()
|
|
return io.ReadAll(r)
|
|
}
|
|
|
|
type brotliCompression struct {
|
|
quality int
|
|
buf []byte
|
|
w *brotli.Writer
|
|
r *brotli.Reader
|
|
}
|
|
|
|
func (b *brotliCompression) Compress(data []byte) ([]byte, error) {
|
|
if b.w == nil {
|
|
b.w = brotli.NewWriterLevel(nil, b.quality)
|
|
}
|
|
|
|
b.buf = (b.buf)[:0]
|
|
compressed := bytes.NewBuffer(b.buf)
|
|
b.w.Reset(compressed)
|
|
_, err := b.w.Write(data)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
err = b.w.Flush()
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if len(compressed.Bytes()) > cap(b.buf) {
|
|
b.buf = compressed.Bytes()
|
|
}
|
|
return compressed.Bytes(), nil
|
|
}
|
|
|
|
func (b *brotliCompression) Decompress(data []byte) ([]byte, error) {
|
|
if b.r == nil {
|
|
b.r = brotli.NewReader(nil)
|
|
}
|
|
b.r.Reset(bytes.NewReader(data))
|
|
return io.ReadAll(b.r)
|
|
}
|
|
|
|
// func compressSnappy(bytes []byte, buf *[]byte) ([]byte, error) {
|
|
// // snappy uses len() to see if it needs to allocate a new slice. Make the
|
|
// // buffer as long as possible.
|
|
// *buf = (*buf)[0:cap(*buf)]
|
|
// compressed := snappy.Encode(*buf, bytes)
|
|
// if n := snappy.MaxEncodedLen(len(bytes)); buf != nil && n > len(*buf) {
|
|
// // grow the buffer for the next time
|
|
// *buf = make([]byte, n)
|
|
// }
|
|
|
|
// return compressed, nil
|
|
// }
|
|
|
|
// func compressSnappyAlt(bytes []byte, buf *[]byte) ([]byte, error) {
|
|
// res := reSnappy.Encode(*buf, bytes)
|
|
// if n := reSnappy.MaxEncodedLen(len(bytes)); buf != nil && n > len(*buf) {
|
|
// // grow the buffer for the next time
|
|
// *buf = make([]byte, n)
|
|
// }
|
|
// return res, nil
|
|
// }
|
|
|
|
// func compressS2(bytes []byte, buf *[]byte) ([]byte, error) {
|
|
// res := reS2.Encode(*buf, bytes)
|
|
// if n := reS2.MaxEncodedLen(len(bytes)); buf != nil && n > len(*buf) {
|
|
// // grow the buffer for the next time
|
|
// *buf = make([]byte, n)
|
|
// }
|
|
// return res, nil
|
|
// }
|
|
|
|
// func compressZstdWithLevel(level reZstd.EncoderLevel) func(data []byte, buf *[]byte) ([]byte, error) {
|
|
// // TODO: use a pool or something. just testing for now
|
|
// encoder, err := reZstd.NewWriter(nil, reZstd.WithEncoderLevel(level))
|
|
// return func(data []byte, buf *[]byte) ([]byte, error) {
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// *buf = (*buf)[:0]
|
|
// writer := bytes.NewBuffer(*buf)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// encoder.Reset(writer)
|
|
// encoder.Write(data)
|
|
// err = encoder.Close()
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// res := writer.Bytes()
|
|
// if len(res) > cap(*buf) {
|
|
// *buf = res
|
|
// }
|
|
// return res, nil
|
|
// }
|
|
// }
|
|
|
|
// func compressGzipWithLevel(level int) func([]byte, *[]byte) ([]byte, error) {
|
|
// // TODO: use a pool or something. just testing for now
|
|
// gzWriter, err := gzip.NewWriterLevel(nil, level)
|
|
|
|
// return func(data []byte, buf2 *[]byte) ([]byte, error) {
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// *buf2 = (*buf2)[:0]
|
|
// buf := bytes.NewBuffer(*buf2)
|
|
// gzWriter.Reset(buf)
|
|
// _, err = gzWriter.Write(data)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// err = gzWriter.Close()
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
|
|
// if len(buf.Bytes()) > cap(*buf2) {
|
|
// *buf2 = buf.Bytes()
|
|
// }
|
|
|
|
// return buf.Bytes(), nil
|
|
// }
|
|
// }
|
|
|
|
// func compressLzw() func(data []byte, buf *[]byte) ([]byte, error) {
|
|
// writer := lzw.NewWriter(nil, lzw.LSB, 8).(*lzw.Writer)
|
|
|
|
// return func(data []byte, buf *[]byte) ([]byte, error) {
|
|
// compressed := bytes.NewBuffer(*buf)
|
|
// writer.Reset(compressed, lzw.LSB, 8)
|
|
// _, err := writer.Write(data)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// err = writer.Close()
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// if len(compressed.Bytes()) > cap(*buf) {
|
|
// *buf = compressed.Bytes()
|
|
// }
|
|
// return compressed.Bytes(), nil
|
|
// }
|
|
// }
|
|
|
|
// func compressFlateWithLevel(level int) func(data []byte, buf *[]byte) ([]byte, error) {
|
|
// writer, err := flate.NewWriter(nil, level)
|
|
|
|
// return func(data []byte, buf *[]byte) ([]byte, error) {
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// *buf = (*buf)[:0]
|
|
// compressed := bytes.NewBuffer(*buf)
|
|
// writer.Reset(compressed)
|
|
// _, err = writer.Write(data)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// err = writer.Close()
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// if len(compressed.Bytes()) > cap(*buf) {
|
|
// *buf = compressed.Bytes()
|
|
// }
|
|
// return compressed.Bytes(), nil
|
|
// }
|
|
// }
|
|
|
|
// func compressBrotliWithQuality(q int) func(data []byte, _ *[]byte) ([]byte, error) {
|
|
// writer := brotli.NewWriterLevel(nil, q)
|
|
// return func(data []byte, buf *[]byte) ([]byte, error) {
|
|
// // return brotli.Encode(data, brotli.WriterOptions{Quality: q})
|
|
// *buf = (*buf)[:0]
|
|
// compressed := bytes.NewBuffer(*buf)
|
|
// writer.Reset(compressed)
|
|
// _, err := writer.Write(data)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// err = writer.Flush()
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// if len(compressed.Bytes()) > cap(*buf) {
|
|
// *buf = compressed.Bytes()
|
|
// }
|
|
// return compressed.Bytes(), nil
|
|
// }
|
|
|
|
// }
|
|
|
|
// func decompressBrotli(compressed []byte) ([]byte, error) {
|
|
// reader := bytes.NewReader(compressed)
|
|
// r := brotli.NewReader(reader)
|
|
// return io.ReadAll(r)
|
|
// }
|
|
|
|
// func decompressFlate(compressed []byte) ([]byte, error) {
|
|
// reader := bytes.NewReader(compressed)
|
|
// r := flate.NewReader(reader)
|
|
// defer r.Close()
|
|
// return io.ReadAll(r)
|
|
// }
|
|
|
|
// func decompressLzw(compressed []byte) ([]byte, error) {
|
|
// reader := bytes.NewReader(compressed)
|
|
// r := lzw.NewReader(reader, lzw.LSB, 8)
|
|
// defer r.Close()
|
|
// return io.ReadAll(r)
|
|
// }
|
|
|
|
// func decompressSnappy(compressed []byte) ([]byte, error) {
|
|
// uncompressed, err := snappy.Decode(nil, compressed)
|
|
// return uncompressed, err
|
|
// }
|
|
|
|
// func decompressSnappyAlt(compressed []byte) ([]byte, error) {
|
|
// uncompressed, err := reSnappy.Decode(nil, compressed)
|
|
// return uncompressed, err
|
|
// }
|
|
|
|
// func decompressS2(compressed []byte) ([]byte, error) {
|
|
// uncompressed, err := reS2.Decode(nil, compressed)
|
|
// return uncompressed, err
|
|
// }
|
|
|
|
// func decompressZstd(compressed []byte) ([]byte, error) {
|
|
// reader := bytes.NewReader(compressed)
|
|
// decoder, err := reZstd.NewReader(reader)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// defer decoder.Close()
|
|
// return io.ReadAll(decoder)
|
|
// }
|
|
|
|
// func decompressGzip(compressed []byte) ([]byte, error) {
|
|
// r := bytes.NewReader(compressed)
|
|
// gzReader, err := gzip.NewReader(r)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
// defer gzReader.Close()
|
|
// decompressedData, err := io.ReadAll(gzReader)
|
|
// if err != nil {
|
|
// return nil, err
|
|
// }
|
|
|
|
// return decompressedData, nil
|
|
// }
|