go-git / utils / sync / zlib.go

package sync

import (
	"bytes"
	"compress/zlib"
	"io"
	"sync"
)

var (
	zlibInitBytes = []byte{0x78, 0x9c, 0x01, 0x00, 0x00, 0xff, 0xff, 0x00, 0x00, 0x00, 0x01}
	zlibReader    = sync.Pool{
		New: func() interface{} {
			return NewZlibReader(nil)
		},
	}
	zlibWriter = sync.Pool{
		New: func() interface{} {
			return zlib.NewWriter(nil)
		},
	}
)

type zlibReadCloser interface {
	io.ReadCloser
	zlib.Resetter
}

func NewZlibReader(dict *[]byte) ZLibReader {
	r, _ := zlib.NewReader(bytes.NewReader(zlibInitBytes))
	return ZLibReader{
		Reader: r.(zlibReadCloser),
		dict:   dict,
	}
}

type ZLibReader struct {
	dict   *[]byte
	Reader zlibReadCloser
}

func (z ZLibReader) Reset(r io.Reader) error {
	var dict []byte
	if z.dict != nil {
		dict = *z.dict
	}
	return z.Reader.Reset(r, dict)
}

// GetZlibReader returns a ZLibReader that is managed by a sync.Pool.
// Returns a ZLibReader that is reset using a dictionary that is
// also managed by a sync.Pool.
//
// After use, the ZLibReader should be put back into the sync.Pool
// by calling PutZlibReader.
func GetZlibReader(r io.Reader) (ZLibReader, error) {
	z := zlibReader.Get().(ZLibReader)
	z.dict = GetByteSlice()

	err := z.Reader.Reset(r, *z.dict)

	return z, err
}

// PutZlibReader puts z back into its sync.Pool, first closing the reader.
// The Byte slice dictionary is also put back into its sync.Pool.
func PutZlibReader(z ZLibReader) {
	z.Reader.Close()
	PutByteSlice(z.dict)
	zlibReader.Put(z)
}

// GetZlibWriter returns a *zlib.Writer that is managed by a sync.Pool.
// Returns a writer that is reset with w and ready for use.
//
// After use, the *zlib.Writer should be put back into the sync.Pool
// by calling PutZlibWriter.
func GetZlibWriter(w io.Writer) *zlib.Writer {
	z := zlibWriter.Get().(*zlib.Writer)
	z.Reset(w)
	return z
}

// PutZlibWriter puts w back into its sync.Pool.
func PutZlibWriter(w *zlib.Writer) {
	zlibWriter.Put(w)
}