arche / internal/store/codec.go

commit 154431fd
  1package store
  2
  3import (
  4	"bytes"
  5	"fmt"
  6	"io"
  7
  8	"github.com/klauspost/compress/zstd"
  9	lz4 "github.com/pierrec/lz4/v4"
 10)
 11
 12type codec interface {
 13	Compress(src []byte) []byte
 14	Decompress(src []byte) ([]byte, error)
 15	Close()
 16}
 17
 18func newCodec(name string, dict []byte) (codec, error) {
 19	switch name {
 20	case "", "zstd":
 21		return newZstdCodec(dict)
 22	case "lz4":
 23		return &lz4Codec{}, nil
 24	case "none":
 25		return &noneCodec{}, nil
 26	default:
 27		return nil, fmt.Errorf("unknown compression %q: must be zstd, lz4, or none", name)
 28	}
 29}
 30
 31type zstdCodec struct {
 32	enc *zstd.Encoder
 33	dec *zstd.Decoder
 34}
 35
 36func newZstdCodec(dict []byte) (*zstdCodec, error) {
 37	encOpts := []zstd.EOption{zstd.WithEncoderLevel(zstd.SpeedDefault)}
 38	if len(dict) > 0 {
 39		encOpts = append(encOpts, zstd.WithEncoderDict(dict))
 40	}
 41	enc, err := zstd.NewWriter(nil, encOpts...)
 42	if err != nil {
 43		return nil, fmt.Errorf("zstd encoder: %w", err)
 44	}
 45
 46	decOpts := []zstd.DOption{}
 47	if len(dict) > 0 {
 48		decOpts = append(decOpts, zstd.WithDecoderDicts(dict))
 49	}
 50	dec, err := zstd.NewReader(nil, decOpts...)
 51	if err != nil {
 52		enc.Close()
 53		return nil, fmt.Errorf("zstd decoder: %w", err)
 54	}
 55	return &zstdCodec{enc: enc, dec: dec}, nil
 56}
 57
 58func (c *zstdCodec) Compress(src []byte) []byte {
 59	return c.enc.EncodeAll(src, nil)
 60}
 61
 62func (c *zstdCodec) Decompress(src []byte) ([]byte, error) {
 63	return c.dec.DecodeAll(src, nil)
 64}
 65
 66func (c *zstdCodec) Close() {
 67	c.enc.Close()
 68	c.dec.Close()
 69}
 70
 71type lz4Codec struct{}
 72
 73func (c *lz4Codec) Compress(src []byte) []byte {
 74	var buf bytes.Buffer
 75	w := lz4.NewWriter(&buf)
 76	_, _ = w.Write(src)
 77	_ = w.Close()
 78	return buf.Bytes()
 79}
 80
 81func (c *lz4Codec) Decompress(src []byte) ([]byte, error) {
 82	r := lz4.NewReader(bytes.NewReader(src))
 83	out, err := io.ReadAll(r)
 84	if err != nil {
 85		return nil, fmt.Errorf("lz4 decompress: %w", err)
 86	}
 87	return out, nil
 88}
 89
 90func (c *lz4Codec) Close() {}
 91
 92type noneCodec struct{}
 93
 94func (c *noneCodec) Compress(src []byte) []byte {
 95	out := make([]byte, len(src))
 96	copy(out, src)
 97	return out
 98}
 99
100func (c *noneCodec) Decompress(src []byte) ([]byte, error) {
101	out := make([]byte, len(src))
102	copy(out, src)
103	return out, nil
104}
105
106func (c *noneCodec) Close() {}