Update

vendor/github.com/klauspost/compress/zstd/encoder.go

@@ -38,6 +38,7 @@ type encoder interface {
 	WindowSize(size int64) int32
 	UseBlock(*blockEnc)
 	Reset(d *dict, singleBlock bool)
+	ResetPrefix(prefix []byte)
 }
 
 type encoderState struct {
@@ -60,6 +61,9 @@ type encoderState struct {
 	wg sync.WaitGroup
 	// This waitgroup indicates we have a block encoding/writing.
 	wWg sync.WaitGroup
+
+	// Parallel job state (used when concurrentBlocks is enabled).
+	jobs jobState
 }
 
 // NewWriter will create a new Zstandard encoder.
@@ -74,6 +78,9 @@ func NewWriter(w io.Writer, opts ...EOption) (*Encoder, error) {
 			return nil, err
 		}
 	}
+	if e.o.concurrentBlocks && (e.o.dict != nil || e.o.concurrent <= 1) {
+		e.o.concurrentBlocks = false
+	}
 	if w != nil {
 		e.Reset(w)
 	}
@@ -95,12 +102,31 @@ func (e *Encoder) initialize() {
 // as a new, independent stream.
 func (e *Encoder) Reset(w io.Writer) {
 	s := &e.state
+
+	if e.o.concurrentBlocks {
+		e.shutdownJobWorkers()
+		js := &s.jobs
+		js.jobSize = e.o.jobSize()
+		js.overlapSize = e.o.overlapSize()
+		// js.filling is allocated lazily on first Write/ReadFrom so callers
+		// that only use EncodeAll don't pay the (up to ~32 MB) jobSize cost.
+		js.filling = js.filling[:0]
+		if js.nextPrefix != nil {
+			js.putOverlapBuf(js.nextPrefix)
+			js.nextPrefix = nil
+		}
+		js.jobSeq = 0
+		js.flushedSeq = 0
+		js.flusherErr = nil
+		js.started = false
+	}
+
 	s.wg.Wait()
 	s.wWg.Wait()
 	if cap(s.filling) == 0 {
 		s.filling = make([]byte, 0, e.o.blockSize)
 	}
-	if e.o.concurrent > 1 {
+	if e.o.concurrent > 1 && !e.o.concurrentBlocks {
 		if cap(s.current) == 0 {
 			s.current = make([]byte, 0, e.o.blockSize)
 		}
@@ -131,6 +157,32 @@ func (e *Encoder) Reset(w io.Writer) {
 	s.frameContentSize = 0
 }
 
+// ResetWithOptions will re-initialize the writer and apply the given options
+// as a new, independent stream.
+// Options are applied on top of the existing options.
+// Some options cannot be changed on reset and will return an error.
+func (e *Encoder) ResetWithOptions(w io.Writer, opts ...EOption) error {
+	e.o.resetOpt = true
+	defer func() { e.o.resetOpt = false }()
+	hadDict := e.o.dict != nil
+	for _, o := range opts {
+		if err := o(&e.o); err != nil {
+			return err
+		}
+	}
+	hasDict := e.o.dict != nil
+	if e.o.concurrentBlocks && hasDict {
+		e.o.concurrentBlocks = false
+	}
+	if hadDict != hasDict {
+		// Dict presence changed — encoder type must be recreated.
+		e.state.encoder = nil
+		e.init = sync.Once{}
+	}
+	e.Reset(w)
+	return nil
+}
+
 // ResetContentSize will reset and set a content size for the next stream.
 // If the bytes written does not match the size given an error will be returned
 // when calling Close().
@@ -153,6 +205,49 @@ func (e *Encoder) Write(p []byte) (n int, err error) {
 	if s.eofWritten {
 		return 0, ErrEncoderClosed
 	}
+	if e.o.concurrentBlocks {
+		return e.writeJobs(p)
+	}
+	return e.writeBlocks(p)
+}
+
+func (e *Encoder) writeJobs(p []byte) (n int, err error) {
+	s := &e.state
+	js := &s.jobs
+	jobSize := js.jobSize
+	if cap(js.filling) == 0 && len(p) > 0 {
+		js.filling = make([]byte, 0, jobSize)
+	}
+	for len(p) > 0 {
+		if len(p)+len(js.filling) < jobSize {
+			if e.o.crc {
+				_, _ = s.encoder.CRC().Write(p)
+			}
+			js.filling = append(js.filling, p...)
+			return n + len(p), nil
+		}
+		add := p
+		if len(p)+len(js.filling) > jobSize {
+			add = add[:jobSize-len(js.filling)]
+		}
+		if e.o.crc {
+			_, _ = s.encoder.CRC().Write(add)
+		}
+		js.filling = append(js.filling, add...)
+		p = p[len(add):]
+		n += len(add)
+		if len(js.filling) < jobSize {
+			return n, nil
+		}
+		if err := e.dispatchJob(false); err != nil {
+			return n, err
+		}
+	}
+	return n, nil
+}
+
+func (e *Encoder) writeBlocks(p []byte) (n int, err error) {
+	s := &e.state
 	for len(p) > 0 {
 		if len(p)+len(s.filling) < e.o.blockSize {
 			if e.o.crc {
@@ -351,6 +446,10 @@ func (e *Encoder) ReadFrom(r io.Reader) (n int64, err error) {
 		println("Using ReadFrom")
 	}
 
+	if e.o.concurrentBlocks {
+		return e.readFromJobs(r)
+	}
+
 	// Flush any current writes.
 	if len(e.state.filling) > 0 {
 		if err := e.nextBlock(false); err != nil {
@@ -364,7 +463,6 @@ func (e *Encoder) ReadFrom(r io.Reader) (n int64, err error) {
 		if e.o.crc {
 			_, _ = e.state.encoder.CRC().Write(src[:n2])
 		}
-		// src is now the unfilled part...
 		src = src[n2:]
 		n += int64(n2)
 		switch err {
@@ -397,15 +495,63 @@ func (e *Encoder) ReadFrom(r io.Reader) (n int64, err error) {
 	}
 }
 
+func (e *Encoder) readFromJobs(r io.Reader) (n int64, err error) {
+	js := &e.state.jobs
+	jobSize := js.jobSize
+
+	// Flush any current filling.
+	if len(js.filling) > 0 {
+		if err := e.dispatchJob(false); err != nil {
+			return 0, err
+		}
+	}
+
+	if cap(js.filling) < jobSize {
+		js.filling = make([]byte, 0, jobSize)
+	}
+	js.filling = js.filling[:jobSize]
+	src := js.filling
+	for {
+		n2, err := r.Read(src)
+		if e.o.crc {
+			_, _ = e.state.encoder.CRC().Write(src[:n2])
+		}
+		src = src[n2:]
+		n += int64(n2)
+		switch err {
+		case io.EOF:
+			js.filling = js.filling[:len(js.filling)-len(src)]
+			return n, nil
+		case nil:
+		default:
+			e.state.err = err
+			return n, err
+		}
+		if len(src) > 0 {
+			continue
+		}
+		if err = e.dispatchJob(false); err != nil {
+			return n, err
+		}
+		if cap(js.filling) < jobSize {
+			js.filling = make([]byte, 0, jobSize)
+		}
+		js.filling = js.filling[:jobSize]
+		src = js.filling
+	}
+}
+
 // Flush will send the currently written data to output
 // and block until everything has been written.
 // This should only be used on rare occasions where pushing the currently queued data is critical.
 func (e *Encoder) Flush() error {
 	s := &e.state
+	if e.o.concurrentBlocks {
+		return e.flushJobs()
+	}
 	if len(s.filling) > 0 {
 		err := e.nextBlock(false)
 		if err != nil {
-			// Ignore Flush after Close.
 			if errors.Is(s.err, ErrEncoderClosed) {
 				return nil
 			}
@@ -415,7 +561,6 @@ func (e *Encoder) Flush() error {
 	s.wg.Wait()
 	s.wWg.Wait()
 	if s.err != nil {
-		// Ignore Flush after Close.
 		if errors.Is(s.err, ErrEncoderClosed) {
 			return nil
 		}
@@ -424,6 +569,20 @@ func (e *Encoder) Flush() error {
 	return s.writeErr
 }
 
+func (e *Encoder) flushJobs() error {
+	js := &e.state.jobs
+	if len(js.filling) > 0 {
+		if err := e.dispatchJob(false); err != nil {
+			return err
+		}
+	}
+	e.waitAllJobs()
+	js.mu.Lock()
+	fErr := js.flusherErr
+	js.mu.Unlock()
+	return fErr
+}
+
 // Close will flush the final output and close the stream.
 // The function will block until everything has been written.
 // The Encoder can still be re-used after calling this.
@@ -432,6 +591,16 @@ func (e *Encoder) Close() error {
 	if s.encoder == nil {
 		return nil
 	}
+	if e.o.concurrentBlocks {
+		return e.closeJobs()
+	}
+	if s.w == nil {
+		if len(s.filling) == 0 && !s.headerWritten && !s.eofWritten && s.nInput == 0 {
+			return nil
+		}
+		return errors.New("zstd: encoder has no writer")
+	}
+
 	err := e.nextBlock(true)
 	if err != nil {
 		if errors.Is(s.err, ErrEncoderClosed) {
@@ -482,6 +651,68 @@ func (e *Encoder) Close() error {
 	return s.err
 }
 
+func (e *Encoder) closeJobs() error {
+	s := &e.state
+	js := &s.jobs
+
+	if errors.Is(s.err, ErrEncoderClosed) {
+		return nil
+	}
+
+	if s.w == nil {
+		if len(js.filling) == 0 && !s.headerWritten && !s.eofWritten && s.nInput == 0 {
+			return nil
+		}
+		return errors.New("zstd: encoder has no writer")
+	}
+
+	if err := e.dispatchJob(true); err != nil {
+		e.shutdownJobWorkers()
+		if errors.Is(s.err, ErrEncoderClosed) {
+			return nil
+		}
+		return err
+	}
+
+	if s.frameContentSize > 0 && s.nInput != s.frameContentSize {
+		e.shutdownJobWorkers()
+		return fmt.Errorf("frame content size %d given, but %d bytes was written", s.frameContentSize, s.nInput)
+	}
+
+	if s.fullFrameWritten {
+		e.shutdownJobWorkers()
+		s.err = ErrEncoderClosed
+		return nil
+	}
+
+	e.shutdownJobWorkers()
+	if js.flusherErr != nil {
+		return js.flusherErr
+	}
+
+	// Write CRC
+	if e.o.crc {
+		var tmp [4]byte
+		_, s.err = s.w.Write(s.encoder.AppendCRC(tmp[:0]))
+		s.nWritten += 4
+	}
+
+	// Add padding
+	if s.err == nil && e.o.pad > 0 {
+		add := calcSkippableFrame(s.nWritten, int64(e.o.pad))
+		frame, err := skippableFrame(js.filling[:0], add, rand.Reader)
+		if err != nil {
+			return err
+		}
+		_, s.err = s.w.Write(frame)
+	}
+	if s.err == nil {
+		s.err = ErrEncoderClosed
+		return nil
+	}
+	return s.err
+}
+
 // EncodeAll will encode all input in src and append it to dst.
 // This function can be called concurrently, but each call will only run on a single goroutine.
 // If empty input is given, nothing is returned, unless WithZeroFrames is specified.