mostynb / go-grpc-compression Goto Github PK

package main

import (
	"bytes"
	"io"
	"os"
	"runtime"
	"sync"

	"github.com/klauspost/compress/zstd"
	_ "github.com/mostynb/go-grpc-compression/nonclobbering/zstd"
	"google.golang.org/grpc/encoding"
)

// Name is the name registered for the gzip compressor.
const Name = "zstd-stream"

func init() {
	c := &compressor{}
	c.poolDecompressor.New = func() interface{} {
		zr, _ := zstd.NewReader(nil)
		runtime.SetFinalizer(zr, (*zstd.Decoder).Close)
		return &reader{
			Reader: zr,
			pool:   &c.poolDecompressor,
		}
	}
	c.poolCompressor.New = func() interface{} {
		zw, _ := zstd.NewWriter(nil)
		runtime.SetFinalizer(zw, (*zstd.Encoder).Close)
		return &writer{
			Writer: zw,
			pool:   &c.poolCompressor,
		}
	}
	encoding.RegisterCompressor(c)
}

type writer struct {
	Writer *zstd.Encoder
	pool   *sync.Pool
}

func (c *compressor) Compress(w io.Writer) (io.WriteCloser, error) {
	z := c.poolCompressor.Get().(*writer)
	z.Writer.Reset(w)
	return z, nil
}

func (z *writer) Close() error {
	defer z.pool.Put(z)
	return z.Writer.Close()
}

func (z *writer) Write(data []byte) (int, error) {
	return z.Writer.Write(data)
}

type reader struct {
	Reader *zstd.Decoder
	pool   *sync.Pool
}

func (c *compressor) Decompress(r io.Reader) (io.Reader, error) {
	z := c.poolDecompressor.Get().(*reader)

	if err := z.Reader.Reset(r); err != nil {
		c.poolDecompressor.Put(z)
		return nil, err
	}
	return z, nil
}

func (z *reader) Read(p []byte) (n int, err error) {
	n, err = z.Reader.Read(p)
	if err == io.EOF {
		z.pool.Put(z)
	}
	return n, err
}

func (c *compressor) Name() string {
	return Name
}

type compressor struct {
	poolCompressor   sync.Pool
	poolDecompressor sync.Pool
}

func main() {
	data, err := os.ReadFile("testdata/test.json")
	if err != nil {
		panic(err)
	}

	compress := func() []byte {
		var buf = bytes.Buffer{}
		c := encoding.GetCompressor("zstd-pool")
		w, _ := c.Compress(&buf)
		_, err = w.Write(data)
		err = w.Close()
		return buf.Bytes()
	}

	for i := 0; i < 10000; i++ {
		compress()
	}
}

real    0m12.891s
user    0m14.182s
sys     0m0.111s

real    0m17.284s
user    0m17.298s
sys     0m0.082s

func main() {
	data, err := os.ReadFile("testdata/test.json")
	if err != nil {
		panic(err)
	}

	compress := func() []byte {
		var buf = bytes.Buffer{}
		c := encoding.GetCompressor("zstd")
		w, _ := c.Compress(&buf)
		_, err = w.Write(data)
		err = w.Close()
		return buf.Bytes()
	}

	compressed := compress()

	decompress := func() {
		c := encoding.GetCompressor("zstd-stream")
		r, err := c.Decompress(bytes.NewReader(compressed))
		if err != nil {
			panic(err)
		}
		_, err = io.ReadAll(r)
		if err != nil {
			panic(err)
		}
	}
	for i := 0; i < 100000; i++ {
		decompress()
	}
}

real    0m34.640s
user    0m40.313s
sys     0m1.434s

real    0m35.556s
user    0m37.684s
sys     0m0.803s