cbrotli: simplify Go implementation and reduce copying

For Read and Write calls, we can pass the Go-allocated slice buffers
directly to the C API. (They do not contain Go pointers, so this
cross-language aliasing is allowed for the duration of the call.)

io.Reader and io.Writer are the Go primitives for this kind of I/O, so
implement those directly instead of providing separate internal
encoder and decoder APIs. (The internal packages weren't laid out to
be compatible with the go toolchain anyway, making this wrapper
awkward use with `go get` instead of Bazel.)
This commit is contained in:
Bryan C. Mills 2017-03-29 04:51:47 -04:00
parent e5b7c16b98
commit 6c2c387d47
8 changed files with 375 additions and 533 deletions

View File

@ -2,16 +2,19 @@ package(default_visibility = ["//visibility:public"])
licenses(["notice"]) # MIT
load("@io_bazel_rules_go//go:def.bzl", "go_prefix", "go_library", "go_test")
load("@io_bazel_rules_go//go:def.bzl", "go_prefix", "cgo_library", "go_test")
go_prefix("github.com/google/brotli")
go_library(
cgo_library(
name = "cbrotli",
srcs = ["cbrotli.go"],
deps = [
"//go/cbrotli/internal:decoder",
"//go/cbrotli/internal:encoder",
srcs = [
"reader.go",
"writer.go",
],
cdeps = [
"//:brotlidec",
"//:brotlienc",
],
)

View File

@ -1,254 +0,0 @@
// Copyright 2016 Google Inc. All Rights Reserved.
//
// Distributed under MIT license.
// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT
// Package cbrotli compresses and decompresses data with C-Brotli library.
package cbrotli
import (
"bytes"
"io"
"github.com/google/brotli/go/cbrotli/internal/decoder"
"github.com/google/brotli/go/cbrotli/internal/encoder"
)
// An internalError reports an error in the (c-)brotli code itself.
type internalError string
func (e internalError) Error() string {
return "cbrotli: internal error: " + string(e)
}
//------------------------------------------------------------------------------
// Encoder
//------------------------------------------------------------------------------
// WriterOptions configures Writer.
type WriterOptions struct {
// Quality controls the compression-speed vs compression-density trade-offs.
// The higher the quality, the slower the compression. Range is 0 to 11.
Quality int
// LGWin is the base 2 logarithm of the sliding window size.
// Range is 10 to 24. 0 indicates automatic configuration based on Quality.
LGWin int
}
// Writer implements io.WriteCloser, an io.Writer decorator that produces
// Brotli-encoded data.
type Writer struct {
dst io.Writer
encoder encoder.Encoder
closed bool
}
// NewWriter initializes new Writer instance.
// Close MUST be called to free resources.
func NewWriter(dst io.Writer, options WriterOptions) *Writer {
return &Writer{
dst: dst,
encoder: encoder.New(options.Quality, options.LGWin),
}
}
// Close implements io.Closer. Close MUST be invoked to free native resources.
// Also Close implicitly flushes remaining data to the decorated writer.
func (z *Writer) Close() error {
if z.closed {
return nil
}
defer z.encoder.Close()
_, err := z.writeChunk(nil, encoder.Finish)
z.closed = true
return err
}
func (z *Writer) writeChunk(p []byte, op encoder.Operation) (int, error) {
if z.closed {
return 0, internalError("write after close")
}
var totalBytesConsumed int
var err error
for {
bytesConsumed, output, status := z.encoder.CompressStream(p, op)
if status == encoder.Error {
err = internalError("encoder failure")
break
}
p = p[bytesConsumed:]
totalBytesConsumed += bytesConsumed
_, err = z.dst.Write(output)
if err != nil {
break
}
if len(p) == 0 && status == encoder.Done {
break
}
}
return totalBytesConsumed, err
}
// Write implements io.Writer.
func (z *Writer) Write(p []byte) (int, error) {
return z.writeChunk(p, encoder.Process)
}
// Flush outputs encoded data for all input provided to Write. The resulting
// output can be decoded to match all input before Flush, but the stream is
// not yet complete until after Close.
// Flush has a negative impact on compression.
func (z *Writer) Flush() error {
_, err := z.writeChunk(nil, encoder.Finish)
return err
}
// Encode returns content encoded with Brotli.
func Encode(content []byte, options WriterOptions) ([]byte, error) {
var buf bytes.Buffer
writer := NewWriter(&buf, options)
defer writer.Close()
_, err := writer.Write(content)
if err != nil {
return nil, err
}
if err := writer.Close(); err != nil {
return nil, err
}
return buf.Bytes(), nil
}
//------------------------------------------------------------------------------
// Decoder
//------------------------------------------------------------------------------
// Reader implements io.ReadCloser, an io.Reader decorator that decodes
// Brotli-encoded data.
type Reader struct {
src io.Reader
decoder decoder.Decoder
buf []byte // intermediate read buffer pointed to by next
eof bool // true if all compressed stream is decoded
next []byte // buffered data to be passed to decoder
output []byte // data produced by decoder, but not yet consumed
srcErr error // last source reader error
err error // reader state; nil if it is OK to read further
closed bool // true is stream is already closed
}
// NewReader initializes new Reader instance.
// Close MUST be called to free resources.
func NewReader(src io.Reader) *Reader {
return &Reader{
src: src,
decoder: decoder.New(),
buf: make([]byte, 32*1024),
eof: false,
}
}
// Close implements io.Closer. Close MUST be invoked to free native resources.
func (z *Reader) Close() error {
if z.closed {
return nil
}
z.decoder.Close()
z.err = internalError("read after close")
z.closed = true
return nil
}
func isEOF(src io.Reader) bool {
n, err := src.Read(make([]byte, 1))
return n == 0 && err == io.EOF
}
// Read implements io.Reader.
func (z *Reader) Read(p []byte) (int, error) {
// Any error state is unrecoverable.
if z.err != nil {
return 0, z.err
}
// See io.Reader documentation.
if len(p) == 0 {
return 0, nil
}
var totalOutBytes int
// There is no practical limit for amount of bytes being consumed by decoder
// before producing any output. Continue feeding decoder until some data is
// produced
for {
// Push already produced output first.
if outBytes := len(z.output); outBytes != 0 {
outBytes = copy(p, z.output)
p = p[outBytes:]
z.output = z.output[outBytes:]
totalOutBytes += outBytes
// Output buffer is full.
if len(p) == 0 {
break
}
continue
}
// No more produced output left.
// If no more output is expected, then we are finished.
if z.eof {
z.err = io.EOF
break
}
// Replenish buffer (might cause blocking read), only if necessary.
if len(z.next) == 0 && totalOutBytes == 0 && z.srcErr != io.EOF {
var n int
n, z.srcErr = z.src.Read(z.buf)
z.next = z.buf[:n]
if z.srcErr != nil && z.srcErr != io.EOF {
z.err = z.srcErr
break
}
}
// Do decoding.
consumed, output, status := z.decoder.DecompressStream(z.next)
z.output = output
z.next = z.next[consumed:]
if status == decoder.Error {
// When error happens, the remaining output does not matter.
z.err = internalError("decoder failure")
break
} else if status == decoder.Done {
// Decoder stream is closed; no further input is expected.
if len(z.next) != 0 || (z.srcErr != io.EOF && !isEOF(z.src)) {
z.err = internalError("excessive input")
break
}
// No more output is expected; keep pushing output.
z.eof = true
continue
} else {
// If can not move any further...
if consumed == 0 && len(z.output) == 0 {
// Unexpected end of input.
if z.srcErr == io.EOF || totalOutBytes == 0 {
z.err = io.ErrUnexpectedEOF
}
// Postpone blocking reads for the next invocation.
break
}
// Continue pushing output.
}
}
return totalOutBytes, z.err
}
// Decode decodes Brotli encoded data.
func Decode(encodedData []byte) ([]byte, error) {
var buf bytes.Buffer
reader := NewReader(bytes.NewReader(encodedData))
defer reader.Close()
_, err := io.Copy(&buf, reader)
if err != nil {
return nil, err
}
return buf.Bytes(), nil
}

View File

@ -91,7 +91,7 @@ func TestEncoderStreams(t *testing.T) {
// to fill the window.
const lgWin = 16
windowSize := int(math.Pow(2, lgWin))
input := make([]byte, 2*windowSize)
input := make([]byte, 8*windowSize)
rand.Read(input)
out := bytes.Buffer{}
e := NewWriter(&out, WriterOptions{Quality: 11, LGWin: lgWin})
@ -106,7 +106,7 @@ func TestEncoderStreams(t *testing.T) {
// We've fed more data than the sliding window size. Check that some
// compressed data has been output.
if out.Len() == 0 {
t.Errorf("Output length is after %d bytes written", n)
t.Errorf("Output length is 0 after %d bytes written", n)
}
if err := e.Close(); err != nil {
t.Errorf("Close Error after copied %d bytes: %v", n, err)
@ -151,6 +151,9 @@ func TestEncoderFlush(t *testing.T) {
if err := e.Flush(); err != nil {
t.Fatalf("Flush(): %v", err)
}
if out.Len() == 0 {
t.Fatalf("0 bytes written after Flush()")
}
decompressed := make([]byte, 1000)
reader := NewReader(bytes.NewReader(out.Bytes()))
n, err := reader.Read(decompressed)
@ -187,8 +190,8 @@ func TestReader(t *testing.T) {
"Reader output:\n"+
"%q\n"+
"want:\n"+
"%q",
got, content)
"<%d bytes>",
got, len(content))
}
}
@ -204,8 +207,8 @@ func TestDecode(t *testing.T) {
"Decode content:\n"+
"%q\n"+
"want:\n"+
"%q",
decoded, content)
"<%d bytes>",
decoded, len(content))
}
}
@ -213,7 +216,13 @@ func TestDecodeFuzz(t *testing.T) {
// Test that the decoder terminates with corrupted input.
content := bytes.Repeat([]byte("hello world!"), 100)
src := rand.NewSource(0)
encoded, _ := Encode(content, WriterOptions{Quality: 5})
encoded, err := Encode(content, WriterOptions{Quality: 5})
if err != nil {
t.Fatalf("Encode(<%d bytes>, _) = _, %s", len(content), err)
}
if len(encoded) == 0 {
t.Fatalf("Encode(<%d bytes>, _) produced empty output", len(content))
}
for i := 0; i < 100; i++ {
enc := append([]byte{}, encoded...)
for j := 0; j < 5; j++ {
@ -260,12 +269,18 @@ func TestEncodeDecode(t *testing.T) {
t.Errorf("Decode: %v", err)
}
if !bytes.Equal(decoded, input) {
var want string
if len(input) > 320 {
want = fmt.Sprintf("<%d bytes>", len(input))
} else {
want = fmt.Sprintf("%q", input)
}
t.Errorf(""+
"Decode content:\n"+
"%q\n"+
"want:\n"+
"%q",
decoded, input)
"%s",
decoded, want)
}
}
}

View File

@ -1,19 +0,0 @@
package(default_visibility = ["//visibility:public"])
licenses(["notice"]) # MIT
load("@io_bazel_rules_go//go:def.bzl", "cgo_library")
cgo_library(
name = "decoder",
srcs = ["decoder.go"],
visibility = ["//go/cbrotli:__subpackages__"],
cdeps = ["//:brotlidec"],
)
cgo_library(
name = "encoder",
srcs = ["encoder.go"],
visibility = ["//go/cbrotli:__subpackages__"],
cdeps = ["//:brotlienc"],
)

View File

@ -1,110 +0,0 @@
// Copyright 2016 Google Inc. All Rights Reserved.
//
// Distributed under MIT license.
// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT
// Package decoder wraps the brotli decoder C API used by package brotli.
package decoder
/*
#include <brotli/decode.h>
// Wrap BrotliDecoderDecompressStream so that it doesn't take variable (in-out)
// pointers. Instead of updated pointer, deltas are saved in auxiliary struct.
struct DecompressStreamResult {
size_t bytes_consumed;
const uint8_t* output_data;
size_t output_data_size;
BrotliDecoderResult status;
};
struct DecompressStreamResult DecompressStream(BrotliDecoderState* s,
const uint8_t* encoded_data, size_t encoded_data_size) {
struct DecompressStreamResult result;
size_t available_in = encoded_data_size;
const uint8_t* next_in = encoded_data;
size_t available_out = 0;
result.status = BrotliDecoderDecompressStream(s,
&available_in, &next_in, &available_out, 0, 0);
result.bytes_consumed = encoded_data_size - available_in;
result.output_data = 0;
result.output_data_size = 0;
if (result.status != BROTLI_DECODER_RESULT_ERROR) {
result.output_data = BrotliDecoderTakeOutput(s, &result.output_data_size);
if (BrotliDecoderIsFinished(s)) {
result.status = BROTLI_DECODER_RESULT_SUCCESS;
}
}
return result;
}
*/
import "C"
import (
"unsafe"
)
// Status represents internal state after DecompressStream invokation
type Status int
const (
// Error happened
Error Status = iota
// Done means that no more output will be produced
Done
// Ok means that more output might be produced with no additional input
Ok
)
// Decoder is the Brotli c-decoder handle.
type Decoder struct {
state *C.BrotliDecoderState
}
// New returns a new Brotli c-decoder handle.
// Close MUST be called to free resources.
func New() Decoder {
return Decoder{state: C.BrotliDecoderCreateInstance(nil, nil, nil)}
}
// Close frees resources used by decoder.
func (z *Decoder) Close() {
C.BrotliDecoderDestroyInstance(z.state)
z.state = nil
}
func goStatus(cStatus C.BrotliDecoderResult) (status Status) {
switch cStatus {
case C.BROTLI_DECODER_RESULT_SUCCESS:
return Done
case C.BROTLI_DECODER_RESULT_NEEDS_MORE_INPUT:
return Ok
case C.BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT:
return Ok
}
return Error
}
// cBytes casts a Go []byte into a C uint8_t*. We pass &buf[0] directly to C,
// which is legal because C doesn't save the pointer longer than the call and
// the byte array itself doesn't contain any pointers.
func cBytes(buf []byte) (*C.uint8_t, C.size_t) {
if len(buf) == 0 {
return (*C.uint8_t)(nil), 0
}
return (*C.uint8_t)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))
}
// DecompressStream reads Brotli-encoded bytes from in, and returns produced
// bytes. Output contents should not be modified. Liveness of output is
// hard-limited by Decoder liveness; slice becomes invalid when any Decoder
// method is invoked.
func (z *Decoder) DecompressStream(in []byte) (
bytesConsumed int, output []byte, status Status) {
cin, cinSize := cBytes(in)
result := C.DecompressStream(z.state, cin, cinSize)
output = C.GoBytes(
unsafe.Pointer(result.output_data), C.int(result.output_data_size))
return int(result.bytes_consumed), output, goStatus(result.status)
}

View File

@ -1,135 +0,0 @@
// Copyright 2016 Google Inc. All Rights Reserved.
//
// Distributed under MIT license.
// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT
// Package encoder wraps the brotli encoder C API used by package brotli.
package encoder
/*
#include <brotli/encode.h>
// Wrap BrotliEncoderCompressStream so that it doesn't take variable (in-out)
// pointers. Instead of updated pointer, deltas are saved in auxiliary struct.
struct CompressStreamResult {
size_t bytes_consumed;
const uint8_t* output_data;
size_t output_data_size;
int success;
int has_more;
};
struct CompressStreamResult CompressStream(
BrotliEncoderState* s, BrotliEncoderOperation op,
const uint8_t* data, size_t data_size) {
struct CompressStreamResult result;
size_t available_in = data_size;
const uint8_t* next_in = data;
size_t available_out = 0;
result.success = BrotliEncoderCompressStream(s, op,
&available_in, &next_in, &available_out, 0, 0) ? 1 : 0;
result.bytes_consumed = data_size - available_in;
result.output_data = 0;
result.output_data_size = 0;
if (result.success) {
result.output_data = BrotliEncoderTakeOutput(s, &result.output_data_size);
}
result.has_more = BrotliEncoderHasMoreOutput(s) ? 1 : 0;
return result;
}
*/
import "C"
import (
"unsafe"
)
// Operation represents type of request to CompressStream
type Operation int
const (
// Process input
Process Operation = iota
// Flush input processed so far
Flush
// Finish stream
Finish
)
// Status represents internal state after CompressStream invocation
type Status int
const (
// Error happened
Error Status = iota
// Done means that no more output will be produced
Done
// Ok means that more output might be produced with no additional input
Ok
)
// Encoder is the Brotli c-encoder handle.
type Encoder struct {
state *C.BrotliEncoderState
}
// New returns a new Brotli c-encoder handle.
// quality and lgWin are described in third_party/Brotli/enc/encode.h.
// Close MUST be called to free resources.
func New(quality, lgWin int) Encoder {
state := C.BrotliEncoderCreateInstance(nil, nil, nil)
// TODO(b/18187008): Check if LGBLOCK or MODE are useful to Flywheel.
C.BrotliEncoderSetParameter(
state, C.BROTLI_PARAM_QUALITY, (C.uint32_t)(quality))
C.BrotliEncoderSetParameter(
state, C.BROTLI_PARAM_LGWIN, (C.uint32_t)(lgWin))
return Encoder{state}
}
// Close frees resources used by encoder.
func (z *Encoder) Close() {
C.BrotliEncoderDestroyInstance(z.state)
z.state = nil
}
// cBytes casts a Go []byte into a C uint8_t*. We pass &buf[0] directly to C,
// which is legal because C doesn't save the pointer longer than the call and
// the byte array itself doesn't contain any pointers.
func cBytes(buf []byte) (*C.uint8_t, C.size_t) {
if len(buf) == 0 {
return (*C.uint8_t)(nil), 0
}
return (*C.uint8_t)(unsafe.Pointer(&buf[0])), C.size_t(len(buf))
}
func cOperation(op Operation) (cOp C.BrotliEncoderOperation) {
switch op {
case Flush:
return C.BROTLI_OPERATION_FLUSH
case Finish:
return C.BROTLI_OPERATION_FINISH
}
return C.BROTLI_OPERATION_PROCESS
}
// CompressStream processes data and produces Brotli-encoded bytes. Encoder may
// consume considerable amount of input before the first output bytes come out.
// Flush and Finish operations force Encoder to produce output that corresponds
// to input consumed so far. Output contents should not be modified. Liveness of
// output is hard-limited by Encoder liveness; slice becomes invalid when any
// Encoder method is invoked.
func (z *Encoder) CompressStream(in []byte, op Operation) (
bytesConsumed int, output []byte, status Status) {
cin, cinSize := cBytes(in)
result := C.CompressStream(z.state, cOperation(op), cin, cinSize)
output = C.GoBytes(
unsafe.Pointer(result.output_data), C.int(result.output_data_size))
var outcome Status
if result.success == 0 {
outcome = Error
} else if result.has_more != 0 {
outcome = Ok
} else {
outcome = Done
}
return int(result.bytes_consumed), output, outcome
}

154
go/cbrotli/reader.go Normal file
View File

@ -0,0 +1,154 @@
// Copyright 2016 Google Inc. All Rights Reserved.
//
// Distributed under MIT license.
// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT
// Package cbrotli compresses and decompresses data with C-Brotli library.
package cbrotli
/*
#include <stddef.h>
#include <stdint.h>
#include <brotli/decode.h>
static BrotliDecoderResult DecompressStream(BrotliDecoderState* s,
uint8_t* out, size_t out_len,
const uint8_t* in, size_t in_len,
size_t* out_written,
size_t* in_consumed) {
size_t in_remaining = in_len;
size_t out_remaining = out_len;
BrotliDecoderResult result = BrotliDecoderDecompressStream(
s, &in_remaining, &in, &out_remaining, &out, NULL);
*out_written = out_len - out_remaining;
*in_consumed = in_len - in_remaining;
return result;
}
*/
import "C"
import (
"bytes"
"errors"
"io"
)
type decodeError C.BrotliDecoderErrorCode
func (err decodeError) Error() string {
return "cbrotli: " + C.GoString(C.BrotliDecoderErrorString(C.BrotliDecoderErrorCode(err)))
}
var errExcessiveInput = errors.New("cbrotli: excessive input")
// Reader implements io.ReadCloser by reading Brotli-encoded data from an
// underlying Reader.
type Reader struct {
src io.Reader
state *C.BrotliDecoderState
buf []byte // scratch space for reading from src
in []byte // current chunk to decode; usually aliases buf
}
// initialBufSize is a "good" buffer size that avoids excessive round-trips
// between C and Go but doesn't waste too much memory on buffering.
// It is arbitrarily chosen to be equal to the constant used in io.Copy.
//
// TODO(bcmills): This constant should be based on empirical measurements.
const initialReadBufSize = 32 * 1024
// NewReader initializes new Reader instance.
// Close MUST be called to free resources.
func NewReader(src io.Reader) *Reader {
return &Reader{
src: src,
state: C.BrotliDecoderCreateInstance(nil, nil, nil),
buf: make([]byte, initialReadBufSize),
}
}
func (r *Reader) Close() error {
if r.state != nil {
C.BrotliDecoderDestroyInstance(r.state)
r.state = nil
}
return nil
}
func (r *Reader) Read(p []byte) (n int, err error) {
if len(r.in) == 0 {
n, err := r.src.Read(r.buf)
if n == 0 && err != nil {
return 0, err
}
r.in = r.buf[:n]
}
if len(p) == 0 {
return 0, nil
}
for n == 0 {
var written, consumed C.size_t
result := C.DecompressStream(r.state,
(*C.uint8_t)(&p[0]), C.size_t(len(p)),
(*C.uint8_t)(&r.in[0]), C.size_t(len(r.in)),
&written, &consumed)
r.in = r.in[int(consumed):]
n = int(written)
switch result {
case C.BROTLI_DECODER_RESULT_SUCCESS:
if len(r.in) > 0 {
return n, errExcessiveInput
}
return n, nil
case C.BROTLI_DECODER_RESULT_ERROR:
return n, decodeError(C.BrotliDecoderGetErrorCode(r.state))
case C.BROTLI_DECODER_RESULT_NEEDS_MORE_OUTPUT:
if n == 0 {
return 0, io.ErrShortBuffer
}
return n, nil
case C.BROTLI_DECODER_NEEDS_MORE_INPUT:
}
if len(r.in) >= len(r.buf)/2 {
// Too much buffer was left over: expand it to make faster progress.
//
// TODO(bcmills): Is there some way to get better feedback about buffer
// sizes from the C API?
r.buf = make([]byte, len(r.in)*2)
}
// Move unread bytes to the start of the buffer.
r.in = r.buf[:copy(r.buf, r.in)]
// Top off the buffer.
encN, err := r.src.Read(r.in[len(r.in):cap(r.in)])
if encN == 0 && n == 0 {
// We need more input to make progress, but there isn't any available.
if err == io.EOF {
return 0, io.ErrUnexpectedEOF
}
return 0, err
}
r.in = r.in[:len(r.in)+encN]
}
return n, nil
}
// Decode decodes Brotli encoded data.
func Decode(encodedData []byte) ([]byte, error) {
r := NewReader(bytes.NewReader(encodedData))
defer r.Close()
var buf bytes.Buffer
_, err := io.Copy(&buf, r)
if err != nil {
return nil, err
}
return buf.Bytes(), nil
}

188
go/cbrotli/writer.go Normal file
View File

@ -0,0 +1,188 @@
// Copyright 2016 Google Inc. All Rights Reserved.
//
// Distributed under MIT license.
// See file LICENSE for detail or copy at https://opensource.org/licenses/MIT
package cbrotli
/*
#include <stdbool.h>
#include <stddef.h>
#include <stdint.h>
#include <brotli/encode.h>
static bool CompressStream(BrotliEncoderState* s, BrotliEncoderOperation op,
uint8_t* out, size_t out_len,
const uint8_t* in, size_t in_len,
size_t* out_written, size_t* in_consumed) {
size_t in_remaining = in_len;
size_t out_remaining = out_len;
bool ok = !!BrotliEncoderCompressStream(s, op, &in_remaining, &in,
&out_remaining, &out, NULL);
*out_written = out_len - out_remaining;
if (in_consumed != NULL) {
*in_consumed = in_len - in_remaining;
}
return ok;
}
*/
import "C"
import (
"bytes"
"errors"
"io"
)
// WriterOptions configures Writer.
type WriterOptions struct {
// Quality controls the compression-speed vs compression-density trade-offs.
// The higher the quality, the slower the compression. Range is 0 to 11.
Quality int
// LGWin is the base 2 logarithm of the sliding window size.
// Range is 10 to 24. 0 indicates automatic configuration based on Quality.
LGWin int
// BufferSize is the number of bytes to use to buffer encoded output.
// 0 indicates an implementation-defined default.
BufferSize int
}
// Writer implements io.WriteCloser by writing Brotli-encoded data to an
// underlying Writer.
type Writer struct {
dst io.Writer
state *C.BrotliEncoderState
buf, encoded []byte
}
var (
errEncode = errors.New("cbrotli: encode error")
errWriterClosed = errors.New("cbrotli: Writer is closed")
)
// NewWriter initializes new Writer instance.
// Close MUST be called to free resources.
func NewWriter(dst io.Writer, options WriterOptions) *Writer {
state := C.BrotliEncoderCreateInstance(nil, nil, nil)
// TODO(b/18187008): Check if LGBLOCK or MODE are useful to Flywheel.
C.BrotliEncoderSetParameter(
state, C.BROTLI_PARAM_QUALITY, (C.uint32_t)(options.Quality))
C.BrotliEncoderSetParameter(
state, C.BROTLI_PARAM_LGWIN, (C.uint32_t)(options.LGWin))
// TODO(bcmills): If the underlying io.Writer implements io.ReaderFrom, use
// that instead of copying through a secondary buffer.
bufSize := options.BufferSize
if bufSize <= 0 {
// TODO(bcmills): Are there better default buffer sizes to use here?
//
// Ideally it would be nice to dynamically resize the buffer based on
// feedback from the encoder.
if options.LGWin > 0 {
bufSize = 1 << uint(options.LGWin)
} else {
bufSize = initialReadBufSize
}
}
return &Writer{
dst: dst,
state: state,
buf: make([]byte, bufSize),
}
}
func (w *Writer) flushBuf() error {
if len(w.encoded) == 0 {
return nil
}
n, err := w.dst.Write(w.encoded)
w.encoded = w.encoded[n:]
return err
}
func (w *Writer) untilEmpty(op C.BrotliEncoderOperation) error {
if w.state == nil {
return errWriterClosed
}
if err := w.flushBuf(); err != nil {
return err
}
for {
var written C.size_t
ok := C.CompressStream(w.state, op,
(*C.uint8_t)(&w.buf[0]), C.size_t(len(w.buf)), nil, 0, &written, nil)
w.encoded = w.buf[:int(written)]
if err := w.flushBuf(); err != nil {
return err
}
if written == 0 {
if !ok {
return errEncode
}
return nil
}
}
}
// Flush outputs encoded data for all input provided to Write. The resulting
// output can be decoded to match all input before Flush, but the stream is
// not yet complete until after Close.
// Flush has a negative impact on compression.
func (w *Writer) Flush() error {
return w.untilEmpty(C.BROTLI_OPERATION_FLUSH)
}
// Close flushes remaining data to the decorated writer
// and frees C resources.
func (w *Writer) Close() error {
err := w.Flush()
if err == nil {
err = w.untilEmpty(C.BROTLI_OPERATION_FINISH)
}
C.BrotliEncoderDestroyInstance(w.state)
w.state = nil
return err
}
// Write implements io.Writer. Flush or Close must be called to ensure that the
// encoded bytes are actually flushed to the underlying Writer.
func (w *Writer) Write(p []byte) (n int, err error) {
if w.state == nil {
return 0, errWriterClosed
}
for len(p) > 0 {
if err := w.flushBuf(); err != nil {
return n, err
}
var written, consumed C.size_t
ok := C.CompressStream(w.state, C.BROTLI_OPERATION_PROCESS,
(*C.uint8_t)(&w.buf[0]), C.size_t(len(w.buf)),
(*C.uint8_t)(&p[0]), C.size_t(len(p)),
&written, &consumed)
w.encoded = w.buf[:int(written)]
n += int(consumed)
p = p[int(consumed):]
if !ok {
return n, errEncode
}
}
return n, nil
}
// Encode returns content encoded with Brotli.
func Encode(content []byte, options WriterOptions) ([]byte, error) {
var buf bytes.Buffer
writer := NewWriter(&buf, options)
_, err := writer.Write(content)
if closeErr := writer.Close(); err == nil && closeErr != nil {
err = closeErr
}
if err != nil {
return nil, err
}
return buf.Bytes(), nil
}