zstd_stream.go

package zstd

/*
#include "zstd.h"

typedef struct compressStream2_result_s {
	size_t return_code;
	size_t bytes_consumed;
	size_t bytes_written;
} compressStream2_result;

static void ZSTD_compressStream2_wrapper(compressStream2_result* result, ZSTD_CCtx* ctx,
		void* dst, size_t maxDstSize, const void* src, size_t srcSize) {
	ZSTD_outBuffer outBuffer = { dst, maxDstSize, 0 };
	ZSTD_inBuffer inBuffer = { src, srcSize, 0 };
	size_t retCode = ZSTD_compressStream2(ctx, &outBuffer, &inBuffer, ZSTD_e_continue);

	result->return_code = retCode;
	result->bytes_consumed = inBuffer.pos;
	result->bytes_written = outBuffer.pos;
}

static void ZSTD_compressStream2_flush(compressStream2_result* result, ZSTD_CCtx* ctx,
		void* dst, size_t maxDstSize, const void* src, size_t srcSize) {
	ZSTD_outBuffer outBuffer = { dst, maxDstSize, 0 };
	ZSTD_inBuffer inBuffer = { src, srcSize, 0 };
	size_t retCode = ZSTD_compressStream2(ctx, &outBuffer, &inBuffer, ZSTD_e_flush);

	result->return_code = retCode;
	result->bytes_consumed = inBuffer.pos;
	result->bytes_written = outBuffer.pos;
}

static void ZSTD_compressStream2_finish(compressStream2_result* result, ZSTD_CCtx* ctx,
		void* dst, size_t maxDstSize, const void* src, size_t srcSize) {
	ZSTD_outBuffer outBuffer = { dst, maxDstSize, 0 };
	ZSTD_inBuffer inBuffer = { src, srcSize, 0 };
	size_t retCode = ZSTD_compressStream2(ctx, &outBuffer, &inBuffer, ZSTD_e_end);

	result->return_code = retCode;
	result->bytes_consumed = inBuffer.pos;
	result->bytes_written = outBuffer.pos;
}

// decompressStream2_result is the same as compressStream2_result, but keep 2 separate struct for easier changes
typedef struct decompressStream2_result_s {
	size_t return_code;
	size_t bytes_consumed;
	size_t bytes_written;
} decompressStream2_result;

static void ZSTD_decompressStream_wrapper(decompressStream2_result* result, ZSTD_DCtx* ctx,
		void* dst, size_t maxDstSize, const void* src, size_t srcSize) {
	ZSTD_outBuffer outBuffer = { dst, maxDstSize, 0 };
	ZSTD_inBuffer inBuffer = { src, srcSize, 0 };
	size_t retCode = ZSTD_decompressStream(ctx, &outBuffer, &inBuffer);

	result->return_code = retCode;
	result->bytes_consumed = inBuffer.pos;
	result->bytes_written = outBuffer.pos;
}
*/
import "C"
import (
	"errors"
	"fmt"
	"io"
	"runtime"
	"sync"
	"unsafe"
)

var errShortRead = errors.New("short read")
var errReaderClosed = errors.New("Reader is closed")
var ErrNoParallelSupport = errors.New("No parallel support")

// Writer is an io.WriteCloser that zstd-compresses its input.
type Writer struct {
	CompressionLevel int

	ctx              *C.ZSTD_CCtx
	dict             []byte
	dstBuffer        []byte
	firstError       error
	underlyingWriter io.Writer
	resultBuffer     *C.compressStream2_result
}

func resize(in []byte, newSize int) []byte {
	if in == nil {
		return make([]byte, newSize)
	}
	if newSize <= cap(in) {
		return in[:newSize]
	}
	toAdd := newSize - len(in)
	return append(in, make([]byte, toAdd)...)
}

// NewWriter creates a new Writer with default compression options.  Writes to
// the writer will be written in compressed form to w.
func NewWriter(w io.Writer) *Writer {
	return NewWriterLevelDict(w, DefaultCompression, nil)
}

// NewWriterLevel is like NewWriter but specifies the compression level instead
// of assuming default compression.
//
// The level can be DefaultCompression or any integer value between BestSpeed
// and BestCompression inclusive.
func NewWriterLevel(w io.Writer, level int) *Writer {
	return NewWriterLevelDict(w, level, nil)

}

// NewWriterLevelDict is like NewWriterLevel but specifies a dictionary to
// compress with.  If the dictionary is empty or nil it is ignored. The dictionary
// should not be modified until the writer is closed.
func NewWriterLevelDict(w io.Writer, level int, dict []byte) *Writer {
	var err error
	ctx := C.ZSTD_createCStream()

	// Load dictionnary if any
	if dict != nil {
		err = getError(int(C.ZSTD_CCtx_loadDictionary(ctx,
			unsafe.Pointer(&dict[0]),
			C.size_t(len(dict)),
		)))
	}

	if err == nil {
		// Only set level if the ctx is not in error already
		err = getError(int(C.ZSTD_CCtx_setParameter(ctx, C.ZSTD_c_compressionLevel, C.int(level))))
	}

	return &Writer{
		CompressionLevel: level,
		ctx:              ctx,
		dict:             dict,
		dstBuffer:        make([]byte, CompressBound(1024)),
		firstError:       err,
		underlyingWriter: w,
		resultBuffer:     new(C.compressStream2_result),
	}
}

// Write writes a compressed form of p to the underlying io.Writer.
func (w *Writer) Write(p []byte) (int, error) {
	if w.firstError != nil {
		return 0, w.firstError
	}
	if len(p) == 0 {
		return 0, nil
	}
	total := len(p)
	// Check if dstBuffer is enough
	w.dstBuffer = w.dstBuffer[0:cap(w.dstBuffer)]
	if len(w.dstBuffer) < CompressBound(len(p)) {
		w.dstBuffer = make([]byte, CompressBound(len(p)))
	}

	dstoff := 0
	consumed := 0
	for len(p) > 0 {
		C.ZSTD_compressStream2_wrapper(
			w.resultBuffer,
			w.ctx,
			unsafe.Pointer(&w.dstBuffer[dstoff]),
			C.size_t(len(w.dstBuffer[dstoff:])),
			unsafe.Pointer(&p[0]),
			C.size_t(len(p)),
		)
		ret := int(w.resultBuffer.return_code)
		if err := getError(ret); err != nil {
			// The stream is dead after this.
			w.firstError = err
			return 0, err
		}
		p = p[w.resultBuffer.bytes_consumed:]
		dstoff += int(w.resultBuffer.bytes_written)
		consumed += int(w.resultBuffer.bytes_consumed)
		if len(p) > 0 && dstoff == len(w.dstBuffer) {
			// We have bytes remaining to compress and our output buffer
			// filled up. This shouldn't happen since we calculated it
			// in advance using CompressBound, but we need to handle it
			// in case there was some miscalculation, or the internal
			// stream buffer contained enough data from previous writes
			// to overflow dstBuffer. (it's not clear from the docs
			// whether this is possible)
			//
			// Allocate space for whatever we haven't compressed yet.
			newbuf := make([]byte, len(w.dstBuffer)+CompressBound(total-consumed))
			copy(newbuf, w.dstBuffer)
			w.dstBuffer = newbuf

		}
	}

	// Write to underlying buffer
	_, err := w.underlyingWriter.Write(w.dstBuffer[:dstoff])

	// Same behaviour as zlib, we can't know how much data we wrote, only
	// if there was an error
	if err != nil {
		return 0, err
	}
	return total, err
}

// Flush writes any unwritten data to the underlying io.Writer.
func (w *Writer) Flush() error {
	if w.firstError != nil {
		return w.firstError
	}

	ret := 1 // So we loop at least once
	for ret > 0 {
		C.ZSTD_compressStream2_flush(
			w.resultBuffer,
			w.ctx,
			unsafe.Pointer(&w.dstBuffer[0]),
			C.size_t(len(w.dstBuffer)),
			unsafe.Pointer(uintptr(0)),
			C.size_t(0),
		)
		ret = int(w.resultBuffer.return_code)
		if err := getError(ret); err != nil {
			return err
		}
		written := int(w.resultBuffer.bytes_written)
		_, err := w.underlyingWriter.Write(w.dstBuffer[:written])
		if err != nil {
			return err
		}

		if ret > 0 { // We have a hint if we need to resize the dstBuffer
			w.dstBuffer = w.dstBuffer[:cap(w.dstBuffer)]
			if len(w.dstBuffer) < ret {
				w.dstBuffer = make([]byte, ret)
			}
		}
	}

	return nil
}

// Close closes the Writer, flushing any unwritten data to the underlying
// io.Writer and freeing objects, but does not close the underlying io.Writer.
func (w *Writer) Close() error {
	if w.firstError != nil {
		return w.firstError
	}

	ret := 1 // So we loop at least once
	for ret > 0 {
		C.ZSTD_compressStream2_finish(
			w.resultBuffer,
			w.ctx,
			unsafe.Pointer(&w.dstBuffer[0]),
			C.size_t(len(w.dstBuffer)),
			unsafe.Pointer(uintptr(0)),
			C.size_t(0),
		)
		ret = int(w.resultBuffer.return_code)
		if err := getError(ret); err != nil {
			return err
		}
		written := int(w.resultBuffer.bytes_written)
		_, err := w.underlyingWriter.Write(w.dstBuffer[:written])
		if err != nil {
			C.ZSTD_freeCStream(w.ctx)
			return err
		}

		if ret > 0 { // We have a hint if we need to resize the dstBuffer
			w.dstBuffer = w.dstBuffer[:cap(w.dstBuffer)]
			if len(w.dstBuffer) < ret {
				w.dstBuffer = make([]byte, ret)
			}
		}
	}

	return getError(int(C.ZSTD_freeCStream(w.ctx)))
}

// Set the number of workers to run the compression in parallel using multiple threads
// If > 1, the Write() call will become asynchronous. This means data will be buffered until processed.
// If you call Write() too fast, you might incur a memory buffer up to as large as your input.
// Consider calling Flush() periodically if you need to compress a very large file that would not fit all in memory.
// By default only one worker is used.
func (w *Writer) SetNbWorkers(n int) error {
	if w.firstError != nil {
		return w.firstError
	}
	if err := getError(int(C.ZSTD_CCtx_setParameter(w.ctx, C.ZSTD_c_nbWorkers, C.int(n)))); err != nil {
		w.firstError = err
		// First error case, a shared libary is used, and the library was compiled without parallel support
		if err.Error() == "Unsupported parameter" {
			return ErrNoParallelSupport
		} else {
			// This could happen if a very large number is passed in, and possibly zstd refuse to create as many threads, or the OS fails to do so
			return err
		}
	}
	return nil
}

// cSize is the recommended size of reader.compressionBuffer. This func and
// invocation allow for a one-time check for validity.
var cSize = func() int {
	v := int(C.ZSTD_DStreamInSize())
	if v <= 0 {
		panic(fmt.Errorf("ZSTD_DStreamInSize() returned invalid size: %v", v))
	}
	return v
}()

// dSize is the recommended size of reader.decompressionBuffer. This func and
// invocation allow for a one-time check for validity.
var dSize = func() int {
	v := int(C.ZSTD_DStreamOutSize())
	if v <= 0 {
		panic(fmt.Errorf("ZSTD_DStreamOutSize() returned invalid size: %v", v))
	}
	return v
}()

// cPool is a pool of buffers for use in reader.compressionBuffer. Buffers are
// taken from the pool in NewReaderDict, returned in reader.Close(). Returns a
// pointer to a slice to avoid the extra allocation of returning the slice as a
// value.
var cPool = sync.Pool{
	New: func() interface{} {
		buff := make([]byte, cSize)
		return &buff
	},
}

// dPool is a pool of buffers for use in reader.decompressionBuffer. Buffers are
// taken from the pool in NewReaderDict, returned in reader.Close(). Returns a
// pointer to a slice to avoid the extra allocation of returning the slice as a
// value.
var dPool = sync.Pool{
	New: func() interface{} {
		buff := make([]byte, dSize)
		return &buff
	},
}

// reader is an io.ReadCloser that decompresses when read from.
type reader struct {
	ctx                 *C.ZSTD_DCtx
	compressionBuffer   []byte
	compressionLeft     int
	decompressionBuffer []byte
	decompOff           int
	decompSize          int
	dict                []byte
	firstError          error
	recommendedSrcSize  int
	resultBuffer        *C.decompressStream2_result
	underlyingReader    io.Reader
}

// NewReader creates a new io.ReadCloser.  Reads from the returned ReadCloser
// read and decompress data from r.  It is the caller's responsibility to call
// Close on the ReadCloser when done.  If this is not done, underlying objects
// in the zstd library will not be freed.
func NewReader(r io.Reader) io.ReadCloser {
	return NewReaderDict(r, nil)
}

// NewReaderDict is like NewReader but uses a preset dictionary.  NewReaderDict
// ignores the dictionary if it is nil.
func NewReaderDict(r io.Reader, dict []byte) io.ReadCloser {
	var err error
	ctx := C.ZSTD_createDStream()
	if len(dict) == 0 {
		err = getError(int(C.ZSTD_initDStream(ctx)))
	} else {
		err = getError(int(C.ZSTD_DCtx_reset(ctx, C.ZSTD_reset_session_only)))
		if err == nil {
			// Only load dictionary if we succesfully inited the context
			err = getError(int(C.ZSTD_DCtx_loadDictionary(
				ctx,
				unsafe.Pointer(&dict[0]),
				C.size_t(len(dict)))))
		}
	}
	compressionBufferP := cPool.Get().(*[]byte)
	decompressionBufferP := dPool.Get().(*[]byte)
	return &reader{
		ctx:                 ctx,
		dict:                dict,
		compressionBuffer:   *compressionBufferP,
		decompressionBuffer: *decompressionBufferP,
		firstError:          err,
		recommendedSrcSize:  cSize,
		resultBuffer:        new(C.decompressStream2_result),
		underlyingReader:    r,
	}
}

// Close frees the allocated C objects
func (r *reader) Close() error {
	if r.firstError != nil {
		return r.firstError
	}

	cb := r.compressionBuffer
	db := r.decompressionBuffer
	// Ensure that we won't resuse buffer
	r.firstError = errReaderClosed
	r.compressionBuffer = nil
	r.decompressionBuffer = nil

	cPool.Put(&cb)
	dPool.Put(&db)
	return getError(int(C.ZSTD_freeDStream(r.ctx)))
}

func (r *reader) Read(p []byte) (int, error) {
	if r.firstError != nil {
		return 0, r.firstError
	}

	if len(p) == 0 {
		return 0, nil
	}

	// If we already have some uncompressed bytes, return without blocking
	if r.decompSize > r.decompOff {
		if r.decompSize-r.decompOff > len(p) {
			copy(p, r.decompressionBuffer[r.decompOff:])
			r.decompOff += len(p)
			return len(p), nil
		}
		// From https://golang.org/pkg/io/#Reader
		// > Read conventionally returns what is available instead of waiting for more.
		copy(p, r.decompressionBuffer[r.decompOff:r.decompSize])
		got := r.decompSize - r.decompOff
		r.decompOff = r.decompSize
		return got, nil
	}

	// Repeatedly read from the underlying reader until we get
	// at least one zstd block, so that we don't block if the
	// other end has flushed a block.
	for {
		// - If the last decompression didn't entirely fill the decompression buffer,
		//   zstd flushed all it could, and needs new data. In that case, do 1 Read.
		// - If the last decompression did entirely fill the decompression buffer,
		//   it might have needed more room to decompress the input. In that case,
		//   don't do any unnecessary Read that might block.
		needsData := r.decompSize < len(r.decompressionBuffer)

		var src []byte
		if !needsData {
			src = r.compressionBuffer[:r.compressionLeft]
		} else {
			src = r.compressionBuffer
			var n int
			var err error
			// Read until data arrives or an error occurs.
			for n == 0 && err == nil {
				n, err = r.underlyingReader.Read(src[r.compressionLeft:])
			}
			if err != nil && err != io.EOF { // Handle underlying reader errors first
				return 0, fmt.Errorf("failed to read from underlying reader: %w", err)
			}
			if n == 0 {
				// Ideally, we'd return with ErrUnexpectedEOF in all cases where the stream was unexpectedly EOF'd
				// during a block or frame, i.e. when there are incomplete, pending compression data.
				// However, it's hard to detect those cases with zstd. Namely, there is no way to know the size of
				// the current buffered compression data in the zstd stream internal buffers.
				// Best effort: throw ErrUnexpectedEOF if we still have some pending buffered compression data that
				// zstd doesn't want to accept.
				// If we don't have any buffered compression data but zstd still has some in its internal buffers,
				// we will return with EOF instead.
				if r.compressionLeft > 0 {
					return 0, io.ErrUnexpectedEOF
				}
				return 0, io.EOF
			}
			src = src[:r.compressionLeft+n]
		}

		// C code
		var srcPtr *byte // Do not point anywhere, if src is empty
		if len(src) > 0 {
			srcPtr = &src[0]
		}

		C.ZSTD_decompressStream_wrapper(
			r.resultBuffer,
			r.ctx,
			unsafe.Pointer(&r.decompressionBuffer[0]),
			C.size_t(len(r.decompressionBuffer)),
			unsafe.Pointer(srcPtr),
			C.size_t(len(src)),
		)
		retCode := int(r.resultBuffer.return_code)

		// Keep src here even though we reuse later, the code might be deleted at some point
		runtime.KeepAlive(src)
		if err := getError(retCode); err != nil {
			return 0, fmt.Errorf("failed to decompress: %w", err)
		}

		// Put everything in buffer
		bytesConsumed := int(r.resultBuffer.bytes_consumed)
		if bytesConsumed < len(src) {
			left := src[bytesConsumed:]
			copy(r.compressionBuffer, left)
		}
		r.compressionLeft = len(src) - bytesConsumed
		r.decompSize = int(r.resultBuffer.bytes_written)
		r.decompOff = copy(p, r.decompressionBuffer[:r.decompSize])

		// Resize buffers
		nsize := retCode // Hint for next src buffer size
		if nsize <= 0 {
			// Reset to recommended size
			nsize = r.recommendedSrcSize
		}
		if nsize < r.compressionLeft {
			nsize = r.compressionLeft
		}
		r.compressionBuffer = resize(r.compressionBuffer, nsize)

		if r.decompOff > 0 {
			return r.decompOff, nil
		}
	}
}