A better way to wrap go's zlib into a streaming decompressor
This commit is contained in:
Kovid Goyal
parent
75a5d88bc2
commit
73ee5b32c9
2 changed files with 109 additions and 50 deletions
|
|
@ -3,7 +3,6 @@
|
|||
package transfer
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"compress/zlib"
|
||||
"fmt"
|
||||
"io"
|
||||
|
|
@ -29,48 +28,6 @@ const (
|
|||
state_canceled
|
||||
)
|
||||
|
||||
type decompressor = func([]byte, bool) ([]byte, error)
|
||||
|
||||
func identity_decompressor(data []byte, is_last bool) ([]byte, error) { return data, nil }
|
||||
|
||||
type zlib_decompressor struct {
|
||||
z io.ReadCloser
|
||||
b bytes.Buffer
|
||||
buf []byte
|
||||
}
|
||||
|
||||
func (self *zlib_decompressor) add_bytes(b []byte, is_last bool) (ans []byte, err error) {
|
||||
self.b.Write(b)
|
||||
pos, n := 0, 0
|
||||
for {
|
||||
if cap(self.buf) < pos+1024 {
|
||||
newcap := utils.Max(2*cap(self.buf), pos+8192)
|
||||
self.buf = append(self.buf[:pos], make([]byte, newcap-pos)...)
|
||||
}
|
||||
n, err = self.z.Read(self.buf[pos:cap(self.buf)])
|
||||
pos += n
|
||||
switch err {
|
||||
case io.EOF:
|
||||
n = 0
|
||||
case nil:
|
||||
default:
|
||||
return nil, err
|
||||
}
|
||||
if n == 0 {
|
||||
break
|
||||
}
|
||||
}
|
||||
if is_last {
|
||||
if err = self.z.Close(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
}
|
||||
if self.b.Len() == 0 {
|
||||
self.b.Reset()
|
||||
}
|
||||
return self.buf[:pos], nil
|
||||
}
|
||||
|
||||
type output_file interface {
|
||||
write([]byte) error
|
||||
close() error
|
||||
|
|
@ -115,7 +72,7 @@ type remote_file struct {
|
|||
parent string
|
||||
expanded_local_path string
|
||||
file_id string
|
||||
decompressor decompressor
|
||||
decompressor utils.StreamDecompressor
|
||||
remote_symlink_value string
|
||||
actual_file output_file
|
||||
}
|
||||
|
|
@ -137,15 +94,16 @@ func new_remote_file(opts *Options, ftc *FileTransmissionCommand) (*remote_file,
|
|||
ans := &remote_file{
|
||||
expected_size: ftc.Size, ftype: ftc.Ftype, mtime: ftc.Mtime, spec_id: spec_id,
|
||||
permissions: ftc.Permissions, remote_path: ftc.Name, display_name: wcswidth.StripEscapeCodes(ftc.Name),
|
||||
remote_id: ftc.Status, remote_target: string(ftc.Data), parent: ftc.Parent, decompressor: identity_decompressor,
|
||||
remote_id: ftc.Status, remote_target: string(ftc.Data), parent: ftc.Parent,
|
||||
}
|
||||
compression_capable := ftc.Ftype == FileType_regular && ftc.Size > 4096 && should_be_compressed(ftc.Name, opts.Compress)
|
||||
if compression_capable {
|
||||
z := &zlib_decompressor{}
|
||||
if z.z, err = zlib.NewReader(&z.b); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
ans.decompressor = z.add_bytes
|
||||
ans.decompressor, err = utils.NewStreamDecompressor(zlib.NewReader)
|
||||
} else {
|
||||
ans.decompressor, err = utils.NewStreamDecompressor(nil)
|
||||
}
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return ans, nil
|
||||
}
|
||||
|
|
|
|||
101
tools/utils/stream_decompressor.go
Normal file
101
tools/utils/stream_decompressor.go
Normal file
|
|
@ -0,0 +1,101 @@
|
|||
// License: GPLv3 Copyright: 2023, Kovid Goyal, <kovid at kovidgoyal.net>
|
||||
|
||||
package utils
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
)
|
||||
|
||||
var _ = fmt.Print
|
||||
|
||||
// Decompress the data in chunk. output_callback() will be called zero or more times with decompressed data. Note that
|
||||
// it may be called in a different goroutine. The data provided to output_callback() is only valid for the lifetime of
|
||||
// output_callback(), so it should copy out the data.
|
||||
type StreamDecompressor = func(chunk []byte, is_last bool, output_callback func([]byte) error) error
|
||||
|
||||
type output struct {
|
||||
chunk []byte
|
||||
err error
|
||||
}
|
||||
|
||||
type stream_decompressor struct {
|
||||
impl io.ReadCloser
|
||||
pipe_r *io.PipeReader
|
||||
pipe_w *io.PipeWriter
|
||||
obuf [8192]byte
|
||||
err error
|
||||
callback func([]byte) error
|
||||
}
|
||||
|
||||
func (self *stream_decompressor) process() {
|
||||
for {
|
||||
n, err := self.impl.Read(self.obuf[:])
|
||||
if n > 0 {
|
||||
if ocerr := self.callback(self.obuf[:n]); ocerr != nil {
|
||||
self.pipe_r.CloseWithError(ocerr)
|
||||
break
|
||||
}
|
||||
}
|
||||
if err != nil {
|
||||
self.pipe_r.CloseWithError(err)
|
||||
break
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (self *stream_decompressor) next(chunk []byte, is_last bool, output_callback func([]byte) error) (err error) {
|
||||
if self.err != nil {
|
||||
return self.err
|
||||
}
|
||||
self.callback = output_callback
|
||||
if _, err = self.pipe_w.Write(chunk); err != nil {
|
||||
self.err = err
|
||||
return err
|
||||
}
|
||||
if is_last {
|
||||
defer func() {
|
||||
self.pipe_r.Close()
|
||||
self.pipe_w.Close()
|
||||
if self.err == nil {
|
||||
self.err = io.EOF
|
||||
}
|
||||
}()
|
||||
self.err = self.impl.Close()
|
||||
return self.err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Wrap Go's awful decompressor routines to allow feeding them
|
||||
// data in chunks. For example:
|
||||
// sd, err := NewStreamDecompressor(zlib.NewReader)
|
||||
// sd(chunk, false, output_callback)
|
||||
// ...
|
||||
// sd(last_chunk, true, output_callback)
|
||||
// after this call calling sd() further will just return io.EOF
|
||||
func NewStreamDecompressor(constructor func(io.Reader) (io.ReadCloser, error)) (StreamDecompressor, error) {
|
||||
if constructor == nil { // identity decompressor
|
||||
var err error
|
||||
return func(chunk []byte, is_last bool, cb func([]byte) error) error {
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
err = cb(chunk)
|
||||
retval := err
|
||||
if is_last && err != nil {
|
||||
err = io.EOF
|
||||
}
|
||||
return retval
|
||||
}, nil
|
||||
}
|
||||
s := stream_decompressor{}
|
||||
s.pipe_r, s.pipe_w = io.Pipe()
|
||||
rc, err := constructor(s.pipe_r)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
s.impl = rc
|
||||
go s.process()
|
||||
return s.next, nil
|
||||
}
|
||||
Loading…
Reference in a new issue