353 lines
7.4 KiB
Go
353 lines
7.4 KiB
Go
// Copyright 2019+ Klaus Post. All rights reserved.
|
|
// License information can be found in the LICENSE file.
|
|
// Based on work by Yann Collet, released under BSD License.
|
|
|
|
package zstd
|
|
|
|
import (
|
|
"fmt"
|
|
rdebug "runtime/debug"
|
|
"sync"
|
|
)
|
|
|
|
type encJob struct {
|
|
prefix []byte // overlap from previous job (nil for first)
|
|
input []byte // job's own input data (swapped from filling)
|
|
last bool // last block of last job gets last=true
|
|
output []byte // compressed blocks (filled by worker)
|
|
err error // encoding error
|
|
done chan struct{} // closed when complete
|
|
}
|
|
|
|
type jobState struct {
|
|
jobSize int
|
|
overlapSize int
|
|
filling []byte // accumulates input up to jobSize
|
|
nextPrefix []byte // overlap prefix prepared for the next dispatched job
|
|
|
|
jobSeq int // next job sequence number
|
|
|
|
jobCh chan *encJob // dispatch to workers
|
|
resultCh chan *encJob // ordered results to flusher
|
|
|
|
workerWg sync.WaitGroup
|
|
flusherWg sync.WaitGroup
|
|
|
|
mu sync.Mutex
|
|
flushedSeq int // last flushed sequence number
|
|
cond *sync.Cond
|
|
|
|
flusherErr error
|
|
started bool
|
|
|
|
inputPool sync.Pool // *[]byte buffers of jobSize cap
|
|
outputPool sync.Pool // *[]byte buffers for compressed output
|
|
overlapPool sync.Pool // *[]byte buffers for overlap prefixes
|
|
}
|
|
|
|
func (e *Encoder) startJobWorkers() {
|
|
js := &e.state.jobs
|
|
n := e.o.concurrent
|
|
js.jobCh = make(chan *encJob, n)
|
|
js.resultCh = make(chan *encJob, n)
|
|
js.flushedSeq = 0
|
|
js.cond = sync.NewCond(&js.mu)
|
|
|
|
// Workers borrow encoders from the shared e.encoders pool per-job.
|
|
// Ensure the pool is initialized before any worker tries to borrow.
|
|
e.init.Do(e.initialize)
|
|
|
|
for range n {
|
|
js.workerWg.Add(1)
|
|
go e.jobWorker()
|
|
}
|
|
js.flusherWg.Add(1)
|
|
go e.jobFlusher()
|
|
js.started = true
|
|
}
|
|
|
|
func (e *Encoder) jobWorker() {
|
|
js := &e.state.jobs
|
|
defer js.workerWg.Done()
|
|
for job := range js.jobCh {
|
|
enc := <-e.encoders
|
|
e.compressJob(enc, job)
|
|
e.encoders <- enc
|
|
close(job.done)
|
|
}
|
|
}
|
|
|
|
func (e *Encoder) compressJob(enc encoder, job *encJob) {
|
|
defer func() {
|
|
if r := recover(); r != nil {
|
|
job.err = fmt.Errorf("panic in parallel job: %v", r)
|
|
rdebug.PrintStack()
|
|
}
|
|
}()
|
|
|
|
if len(job.prefix) > 0 {
|
|
enc.ResetPrefix(job.prefix)
|
|
} else {
|
|
enc.Reset(nil, false)
|
|
}
|
|
|
|
data := job.input
|
|
if len(data) == 0 && job.last {
|
|
blk := enc.Block()
|
|
blk.reset(nil)
|
|
blk.last = true
|
|
blk.encodeRaw(nil)
|
|
job.output = append(job.output, blk.output...)
|
|
return
|
|
}
|
|
|
|
blk := enc.Block()
|
|
for len(data) > 0 {
|
|
todo := data
|
|
if len(todo) > e.o.blockSize {
|
|
todo = todo[:e.o.blockSize]
|
|
}
|
|
data = data[len(todo):]
|
|
|
|
blk.pushOffsets()
|
|
enc.Encode(blk, todo)
|
|
blk.last = len(data) == 0 && job.last
|
|
|
|
err := blk.encode(todo, e.o.noEntropy, !e.o.allLitEntropy)
|
|
if err != nil {
|
|
job.err = err
|
|
return
|
|
}
|
|
job.output = append(job.output, blk.output...)
|
|
blk.reset(nil)
|
|
}
|
|
}
|
|
|
|
func (js *jobState) getInputBuf(size int) []byte {
|
|
if v := js.inputPool.Get(); v != nil {
|
|
bp := v.(*[]byte)
|
|
b := *bp
|
|
if cap(b) >= size {
|
|
return b[:0]
|
|
}
|
|
}
|
|
return make([]byte, 0, size)
|
|
}
|
|
|
|
func (js *jobState) putInputBuf(b []byte) {
|
|
if cap(b) > 0 {
|
|
b = b[:0]
|
|
js.inputPool.Put(&b)
|
|
}
|
|
}
|
|
|
|
func (js *jobState) getOutputBuf(size int) []byte {
|
|
if v := js.outputPool.Get(); v != nil {
|
|
bp := v.(*[]byte)
|
|
b := *bp
|
|
if cap(b) >= size {
|
|
return b[:0]
|
|
}
|
|
}
|
|
return make([]byte, 0, size)
|
|
}
|
|
|
|
func (js *jobState) putOutputBuf(b []byte) {
|
|
if cap(b) > 0 {
|
|
b = b[:0]
|
|
js.outputPool.Put(&b)
|
|
}
|
|
}
|
|
|
|
func (js *jobState) getOverlapBuf(size int) []byte {
|
|
if v := js.overlapPool.Get(); v != nil {
|
|
bp := v.(*[]byte)
|
|
b := *bp
|
|
if cap(b) >= size {
|
|
return b[:size]
|
|
}
|
|
}
|
|
return make([]byte, size)
|
|
}
|
|
|
|
func (js *jobState) putOverlapBuf(b []byte) {
|
|
if cap(b) > 0 {
|
|
b = b[:0]
|
|
js.overlapPool.Put(&b)
|
|
}
|
|
}
|
|
|
|
func (e *Encoder) jobFlusher() {
|
|
js := &e.state.jobs
|
|
defer js.flusherWg.Done()
|
|
for job := range js.resultCh {
|
|
<-job.done
|
|
// Worker has fully exited compressJob, so the prefix is no longer
|
|
// in use. Return it to the pool regardless of outcome.
|
|
if job.prefix != nil {
|
|
js.putOverlapBuf(job.prefix)
|
|
job.prefix = nil
|
|
}
|
|
if job.err != nil {
|
|
js.mu.Lock()
|
|
js.flusherErr = job.err
|
|
js.cond.Broadcast()
|
|
js.mu.Unlock()
|
|
for range js.resultCh {
|
|
}
|
|
return
|
|
}
|
|
if len(job.output) > 0 {
|
|
_, err := e.state.w.Write(job.output)
|
|
if err != nil {
|
|
js.mu.Lock()
|
|
js.flusherErr = err
|
|
js.cond.Broadcast()
|
|
js.mu.Unlock()
|
|
for range js.resultCh {
|
|
}
|
|
return
|
|
}
|
|
e.state.nWritten += int64(len(job.output))
|
|
}
|
|
// Return buffers to pools.
|
|
js.putInputBuf(job.input)
|
|
js.putOutputBuf(job.output)
|
|
job.input = nil
|
|
job.output = nil
|
|
|
|
js.mu.Lock()
|
|
js.flushedSeq++
|
|
js.cond.Broadcast()
|
|
js.mu.Unlock()
|
|
}
|
|
}
|
|
|
|
func (e *Encoder) shutdownJobWorkers() {
|
|
js := &e.state.jobs
|
|
if !js.started {
|
|
return
|
|
}
|
|
close(js.jobCh)
|
|
js.workerWg.Wait()
|
|
close(js.resultCh)
|
|
js.flusherWg.Wait()
|
|
js.started = false
|
|
}
|
|
|
|
// waitAllJobs blocks until all dispatched jobs have been flushed.
|
|
func (e *Encoder) waitAllJobs() {
|
|
js := &e.state.jobs
|
|
if !js.started {
|
|
return
|
|
}
|
|
js.mu.Lock()
|
|
for js.flushedSeq < js.jobSeq && js.flusherErr == nil {
|
|
js.cond.Wait()
|
|
}
|
|
js.mu.Unlock()
|
|
}
|
|
|
|
func (e *Encoder) dispatchJob(final bool) error {
|
|
s := &e.state
|
|
js := &s.jobs
|
|
|
|
js.mu.Lock()
|
|
fErr := js.flusherErr
|
|
js.mu.Unlock()
|
|
if fErr != nil {
|
|
return fErr
|
|
}
|
|
|
|
if !s.headerWritten {
|
|
// Single-block optimization: fall through to encodeAll path.
|
|
if final && len(js.filling) > 0 && len(js.filling) <= e.o.blockSize {
|
|
s.current = e.encodeAll(s.encoder, js.filling, s.current[:0])
|
|
var n2 int
|
|
n2, s.err = s.w.Write(s.current)
|
|
if s.err != nil {
|
|
return s.err
|
|
}
|
|
s.nWritten += int64(n2)
|
|
s.nInput += int64(len(js.filling))
|
|
s.current = s.current[:0]
|
|
js.filling = js.filling[:0]
|
|
s.headerWritten = true
|
|
s.fullFrameWritten = true
|
|
s.eofWritten = true
|
|
return nil
|
|
}
|
|
if final && len(js.filling) == 0 && !e.o.fullZero {
|
|
s.headerWritten = true
|
|
s.fullFrameWritten = true
|
|
s.eofWritten = true
|
|
return nil
|
|
}
|
|
|
|
var tmp [maxHeaderSize]byte
|
|
fh := frameHeader{
|
|
ContentSize: uint64(s.frameContentSize),
|
|
WindowSize: uint32(s.encoder.WindowSize(s.frameContentSize)),
|
|
SingleSegment: false,
|
|
Checksum: e.o.crc,
|
|
DictID: 0,
|
|
}
|
|
dst := fh.appendTo(tmp[:0])
|
|
var n2 int
|
|
n2, s.err = s.w.Write(dst)
|
|
if s.err != nil {
|
|
return s.err
|
|
}
|
|
s.nWritten += int64(n2)
|
|
s.headerWritten = true
|
|
}
|
|
|
|
if len(js.filling) == 0 && !final {
|
|
return nil
|
|
}
|
|
|
|
if !js.started {
|
|
e.startJobWorkers()
|
|
}
|
|
|
|
// Estimate output size for pooled buffer.
|
|
outputEst := max(len(js.filling)/2, 512)
|
|
|
|
job := &encJob{
|
|
last: final,
|
|
done: make(chan struct{}),
|
|
output: js.getOutputBuf(outputEst),
|
|
}
|
|
|
|
// Each job owns its prefix slice; the flusher returns it to the pool
|
|
// after <-job.done, so workers and dispatch never share a buffer.
|
|
if js.nextPrefix != nil {
|
|
job.prefix = js.nextPrefix
|
|
js.nextPrefix = nil
|
|
}
|
|
|
|
// Build the next job's prefix from the tail of this job's input.
|
|
if !final && len(js.filling) > 0 {
|
|
overlapLen := min(js.overlapSize, len(js.filling))
|
|
np := js.getOverlapBuf(overlapLen)
|
|
copy(np, js.filling[len(js.filling)-overlapLen:])
|
|
js.nextPrefix = np
|
|
}
|
|
|
|
// Swap filling buffer into job — zero-copy for the input data.
|
|
job.input = js.filling
|
|
js.filling = js.getInputBuf(js.jobSize)
|
|
|
|
s.nInput += int64(len(job.input))
|
|
js.jobSeq++
|
|
|
|
if final {
|
|
s.eofWritten = true
|
|
}
|
|
|
|
js.resultCh <- job
|
|
js.jobCh <- job
|
|
|
|
return nil
|
|
}
|