Merge pull request #86 from delthas/1.x

[zstd_stream] Add Writer.Flush()

Author: Viq111

README.md

@@ -60,6 +60,9 @@ NewWriterLevelDict(w io.Writer, level int, dict []byte) *Writer
 // Write compresses the input data and write it to the underlying writer
 (w *Writer) Write(p []byte) (int, error)
 
+// Flush writes any unwritten data to the underlying writer
+(w *Writer) Flush() error
+
 // Close flushes the buffer and frees C zstd objects
 (w *Writer) Close() error
 ```

zstd_stream.go

@@ -21,6 +21,16 @@ static void ZSTD_compressStream2_wrapper(compressStream2_result* result, ZSTD_CC
 	result->bytes_written = outBuffer.pos;
 }
 
+static void ZSTD_compressStream2_flush(compressStream2_result* result, ZSTD_CCtx* ctx, uintptr_t dst, size_t maxDstSize, const uintptr_t src, size_t srcSize) {
+	ZSTD_outBuffer outBuffer = { (void*)dst, maxDstSize, 0 };
+	ZSTD_inBuffer inBuffer = { (void*)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, uintptr_t dst, size_t maxDstSize, const uintptr_t src, size_t srcSize) {
 	ZSTD_outBuffer outBuffer = { (void*)dst, maxDstSize, 0 };
 	ZSTD_inBuffer inBuffer = { (void*)src, srcSize, 0 };
@@ -203,6 +213,49 @@ func (w *Writer) Write(p []byte) (int, error) {
 	return len(p), 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 {
+		srcPtr := C.uintptr_t(uintptr(0)) // Do not point anywhere, if src is empty
+		if len(w.srcBuffer) > 0 {
+			srcPtr = C.uintptr_t(uintptr(unsafe.Pointer(&w.srcBuffer[0])))
+		}
+
+		C.ZSTD_compressStream2_flush(
+			w.resultBuffer,
+			w.ctx,
+			C.uintptr_t(uintptr(unsafe.Pointer(&w.dstBuffer[0]))),
+			C.size_t(len(w.dstBuffer)),
+			srcPtr,
+			C.size_t(len(w.srcBuffer)),
+		)
+		ret = int(w.resultBuffer.return_code)
+		if err := getError(ret); err != nil {
+			return err
+		}
+		w.srcBuffer = w.srcBuffer[w.resultBuffer.bytes_consumed:]
+		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 {
@@ -231,7 +284,11 @@ func (w *Writer) Close() error {
 		}
 		w.srcBuffer = w.srcBuffer[w.resultBuffer.bytes_consumed:]
 		written := int(w.resultBuffer.bytes_written)
-		w.underlyingWriter.Write(w.dstBuffer[: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)]

zstd_stream_test.go

@@ -2,8 +2,10 @@ package zstd
 
 import (
 	"bytes"
+	"errors"
 	"io"
 	"io/ioutil"
+	"log"
 	"runtime/debug"
 	"testing"
 )
@@ -149,6 +151,81 @@ func TestStreamEmptyPayload(t *testing.T) {
 	}
 }
 
+func TestStreamFlush(t *testing.T) {
+	var w bytes.Buffer
+	writer := NewWriter(&w)
+	reader := NewReader(&w)
+
+	payload := "cc" // keep the payload short to make sure it will not be automatically flushed by zstd
+	buf := make([]byte, len(payload))
+
+	for i := 0; i < 5; i++ {
+		_, err := writer.Write([]byte(payload))
+		failOnError(t, "Failed writing to compress object", err)
+
+		err = writer.Flush()
+		failOnError(t, "Failed flushing compress object", err)
+
+		_, err = io.ReadFull(reader, buf)
+		failOnError(t, "Failed reading uncompress object", err)
+
+		if string(buf) != payload {
+			debug.PrintStack()
+			log.Fatal("Uncompressed object mismatch")
+		}
+	}
+
+	failOnError(t, "Failed to close compress object", writer.Close())
+	failOnError(t, "Failed to close uncompress object", reader.Close())
+}
+
+type closeableWriter struct{
+	w io.Writer
+	closed bool
+}
+
+func (c *closeableWriter) Write(p []byte) (n int, err error) {
+	if c.closed {
+		return 0, errors.New("io: Write on a closed closeableWriter")
+	}
+	return c.w.Write(p)
+}
+
+func (c *closeableWriter) Close() error {
+	c.closed = true
+	return nil
+}
+
+func TestStreamFlushError(t *testing.T) {
+	var bw bytes.Buffer
+	w := closeableWriter{w: &bw}
+	writer := NewWriter(&w)
+
+	_, err := writer.Write([]byte("cc"))
+	failOnError(t, "Failed writing to compress object", err)
+
+	w.Close()
+	if err = writer.Flush(); err == nil {
+		debug.PrintStack()
+		t.Fatal("Writer.Flush returned no error when writing to underlying io.Writer failed")
+	}
+}
+
+func TestStreamCloseError(t *testing.T) {
+	var bw bytes.Buffer
+	w := closeableWriter{w: &bw}
+	writer := NewWriter(&w)
+
+	_, err := writer.Write([]byte("cc"))
+	failOnError(t, "Failed writing to compress object", err)
+
+	w.Close()
+	if err = writer.Close(); err == nil {
+		debug.PrintStack()
+		t.Fatal("Writer.Close returned no error when writing to underlying io.Writer failed")
+	}
+}
+
 type breakingReader struct{}
 
 func (r *breakingReader) Read(p []byte) (int, error) {