1 files changed, 28 insertions, 3 deletions

diff --git a/vendor/github.com/klauspost/compress/s2/writer.go b/vendor/github.com/klauspost/compress/s2/writer.go
index 0a46f2b98..fd15078f7 100644
--- a/vendor/github.com/klauspost/compress/s2/writer.go
+++ b/vendor/github.com/klauspost/compress/s2/writer.go
@@ -83,11 +83,14 @@ type Writer struct {
 	snappy            bool
 	flushOnWrite      bool
 	appendIndex       bool
+	bufferCB          func([]byte)
 	level             uint8
 }
 
 type result struct {
 	b []byte
+	// return when writing
+	ret []byte
 	// Uncompressed start offset
 	startOffset int64
 }
@@ -146,6 +149,10 @@ func (w *Writer) Reset(writer io.Writer) {
 		for write := range toWrite {
 			// Wait for the data to be available.
 			input := <-write
+			if input.ret != nil && w.bufferCB != nil {
+				w.bufferCB(input.ret)
+				input.ret = nil
+			}
 			in := input.b
 			if len(in) > 0 {
 				if w.err(nil) == nil {
@@ -341,7 +348,8 @@ func (w *Writer) AddSkippableBlock(id uint8, data []byte) (err error) {
 // but the input buffer cannot be written to by the caller
 // until Flush or Close has been called when concurrency != 1.
 //
-// If you cannot control that, use the regular Write function.
+// Use the WriterBufferDone to receive a callback when the buffer is done
+// Processing.
 //
 // Note that input is not buffered.
 // This means that each write will result in discrete blocks being created.
@@ -364,6 +372,9 @@ func (w *Writer) EncodeBuffer(buf []byte) (err error) {
 	}
 	if w.concurrency == 1 {
 		_, err := w.writeSync(buf)
+		if w.bufferCB != nil {
+			w.bufferCB(buf)
+		}
 		return err
 	}
 
@@ -378,7 +389,7 @@ func (w *Writer) EncodeBuffer(buf []byte) (err error) {
 			hWriter <- result{startOffset: w.uncompWritten, b: magicChunkBytes}
 		}
 	}
-
+	orgBuf := buf
 	for len(buf) > 0 {
 		// Cut input.
 		uncompressed := buf
@@ -397,6 +408,9 @@ func (w *Writer) EncodeBuffer(buf []byte) (err error) {
 			startOffset: w.uncompWritten,
 		}
 		w.uncompWritten += int64(len(uncompressed))
+		if len(buf) == 0 && w.bufferCB != nil {
+			res.ret = orgBuf
+		}
 		go func() {
 			race.ReadSlice(uncompressed)
 
@@ -922,7 +936,7 @@ func WriterBetterCompression() WriterOption {
 }
 
 // WriterBestCompression will enable better compression.
-// EncodeBetter compresses better than Encode but typically with a
+// EncodeBest compresses better than Encode but typically with a
 // big speed decrease on compression.
 func WriterBestCompression() WriterOption {
 	return func(w *Writer) error {
@@ -941,6 +955,17 @@ func WriterUncompressed() WriterOption {
 	}
 }
 
+// WriterBufferDone will perform a callback when EncodeBuffer has finished
+// writing a buffer to the output and the buffer can safely be reused.
+// If the buffer was split into several blocks, it will be sent after the last block.
+// Callbacks will not be done concurrently.
+func WriterBufferDone(fn func(b []byte)) WriterOption {
+	return func(w *Writer) error {
+		w.bufferCB = fn
+		return nil
+	}
+}
+
 // WriterBlockSize allows to override the default block size.
 // Blocks will be this size or smaller.
 // Minimum size is 4KB and maximum size is 4MB.