Initial QSfera import
This commit is contained in:
+124
@@ -0,0 +1,124 @@
|
||||
package compress
|
||||
|
||||
import (
|
||||
"encoding"
|
||||
"fmt"
|
||||
"io"
|
||||
"strconv"
|
||||
"strings"
|
||||
|
||||
"github.com/segmentio/kafka-go/compress/gzip"
|
||||
"github.com/segmentio/kafka-go/compress/lz4"
|
||||
"github.com/segmentio/kafka-go/compress/snappy"
|
||||
"github.com/segmentio/kafka-go/compress/zstd"
|
||||
)
|
||||
|
||||
// Compression represents the compression applied to a record set.
|
||||
type Compression int8
|
||||
|
||||
const (
|
||||
None Compression = 0
|
||||
Gzip Compression = 1
|
||||
Snappy Compression = 2
|
||||
Lz4 Compression = 3
|
||||
Zstd Compression = 4
|
||||
)
|
||||
|
||||
func (c Compression) Codec() Codec {
|
||||
if i := int(c); i >= 0 && i < len(Codecs) {
|
||||
return Codecs[i]
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (c Compression) String() string {
|
||||
if codec := c.Codec(); codec != nil {
|
||||
return codec.Name()
|
||||
}
|
||||
return "uncompressed"
|
||||
}
|
||||
|
||||
func (c Compression) MarshalText() ([]byte, error) {
|
||||
return []byte(c.String()), nil
|
||||
}
|
||||
|
||||
func (c *Compression) UnmarshalText(b []byte) error {
|
||||
switch string(b) {
|
||||
case "none", "uncompressed":
|
||||
*c = None
|
||||
return nil
|
||||
}
|
||||
|
||||
for _, codec := range Codecs[None+1:] {
|
||||
if codec.Name() == string(b) {
|
||||
*c = Compression(codec.Code())
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
i, err := strconv.ParseInt(string(b), 10, 64)
|
||||
if err == nil && i >= 0 && i < int64(len(Codecs)) {
|
||||
*c = Compression(i)
|
||||
return nil
|
||||
}
|
||||
|
||||
s := &strings.Builder{}
|
||||
s.WriteString("none, uncompressed")
|
||||
|
||||
for i, codec := range Codecs[None+1:] {
|
||||
if i < (len(Codecs) - 1) {
|
||||
s.WriteString(", ")
|
||||
} else {
|
||||
s.WriteString(", or ")
|
||||
}
|
||||
s.WriteString(codec.Name())
|
||||
}
|
||||
|
||||
return fmt.Errorf("compression format must be one of %s, not %q", s, b)
|
||||
}
|
||||
|
||||
var (
|
||||
_ encoding.TextMarshaler = Compression(0)
|
||||
_ encoding.TextUnmarshaler = (*Compression)(nil)
|
||||
)
|
||||
|
||||
// Codec represents a compression codec to encode and decode the messages.
|
||||
// See : https://cwiki.apache.org/confluence/display/KAFKA/Compression
|
||||
//
|
||||
// A Codec must be safe for concurrent access by multiple go routines.
|
||||
type Codec interface {
|
||||
// Code returns the compression codec code
|
||||
Code() int8
|
||||
|
||||
// Human-readable name for the codec.
|
||||
Name() string
|
||||
|
||||
// Constructs a new reader which decompresses data from r.
|
||||
NewReader(r io.Reader) io.ReadCloser
|
||||
|
||||
// Constructs a new writer which writes compressed data to w.
|
||||
NewWriter(w io.Writer) io.WriteCloser
|
||||
}
|
||||
|
||||
var (
|
||||
// The global gzip codec installed on the Codecs table.
|
||||
GzipCodec gzip.Codec
|
||||
|
||||
// The global snappy codec installed on the Codecs table.
|
||||
SnappyCodec snappy.Codec
|
||||
|
||||
// The global lz4 codec installed on the Codecs table.
|
||||
Lz4Codec lz4.Codec
|
||||
|
||||
// The global zstd codec installed on the Codecs table.
|
||||
ZstdCodec zstd.Codec
|
||||
|
||||
// The global table of compression codecs supported by the kafka protocol.
|
||||
Codecs = [...]Codec{
|
||||
None: nil,
|
||||
Gzip: &GzipCodec,
|
||||
Snappy: &SnappyCodec,
|
||||
Lz4: &Lz4Codec,
|
||||
Zstd: &ZstdCodec,
|
||||
}
|
||||
)
|
||||
+123
@@ -0,0 +1,123 @@
|
||||
package gzip
|
||||
|
||||
import (
|
||||
"io"
|
||||
"sync"
|
||||
|
||||
"github.com/klauspost/compress/gzip"
|
||||
)
|
||||
|
||||
var (
|
||||
readerPool sync.Pool
|
||||
)
|
||||
|
||||
// Codec is the implementation of a compress.Codec which supports creating
|
||||
// readers and writers for kafka messages compressed with gzip.
|
||||
type Codec struct {
|
||||
// The compression level to configure on writers created by this codec.
|
||||
// Acceptable values are defined in the standard gzip package.
|
||||
//
|
||||
// Default to gzip.DefaultCompressionLevel.
|
||||
Level int
|
||||
|
||||
writerPool sync.Pool
|
||||
}
|
||||
|
||||
// Code implements the compress.Codec interface.
|
||||
func (c *Codec) Code() int8 { return 1 }
|
||||
|
||||
// Name implements the compress.Codec interface.
|
||||
func (c *Codec) Name() string { return "gzip" }
|
||||
|
||||
// NewReader implements the compress.Codec interface.
|
||||
func (c *Codec) NewReader(r io.Reader) io.ReadCloser {
|
||||
var err error
|
||||
z, _ := readerPool.Get().(*gzip.Reader)
|
||||
if z != nil {
|
||||
err = z.Reset(r)
|
||||
} else {
|
||||
z, err = gzip.NewReader(r)
|
||||
}
|
||||
if err != nil {
|
||||
if z != nil {
|
||||
readerPool.Put(z)
|
||||
}
|
||||
return &errorReader{err: err}
|
||||
}
|
||||
return &reader{Reader: z}
|
||||
}
|
||||
|
||||
// NewWriter implements the compress.Codec interface.
|
||||
func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
|
||||
x := c.writerPool.Get()
|
||||
z, _ := x.(*gzip.Writer)
|
||||
if z == nil {
|
||||
x, err := gzip.NewWriterLevel(w, c.level())
|
||||
if err != nil {
|
||||
return &errorWriter{err: err}
|
||||
}
|
||||
z = x
|
||||
} else {
|
||||
z.Reset(w)
|
||||
}
|
||||
return &writer{codec: c, Writer: z}
|
||||
}
|
||||
|
||||
func (c *Codec) level() int {
|
||||
if c.Level != 0 {
|
||||
return c.Level
|
||||
}
|
||||
return gzip.DefaultCompression
|
||||
}
|
||||
|
||||
type reader struct{ *gzip.Reader }
|
||||
|
||||
func (r *reader) Close() (err error) {
|
||||
if z := r.Reader; z != nil {
|
||||
r.Reader = nil
|
||||
err = z.Close()
|
||||
// Pass it an empty reader, which is a zero-size value implementing the
|
||||
// flate.Reader interface to avoid the construction of a bufio.Reader in
|
||||
// the call to Reset.
|
||||
//
|
||||
// Note: we could also not reset the reader at all, but that would cause
|
||||
// the underlying reader to be retained until the gzip.Reader is freed,
|
||||
// which may not be desirable.
|
||||
z.Reset(emptyReader{})
|
||||
readerPool.Put(z)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
type writer struct {
|
||||
codec *Codec
|
||||
*gzip.Writer
|
||||
}
|
||||
|
||||
func (w *writer) Close() (err error) {
|
||||
if z := w.Writer; z != nil {
|
||||
w.Writer = nil
|
||||
err = z.Close()
|
||||
z.Reset(nil)
|
||||
w.codec.writerPool.Put(z)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
type emptyReader struct{}
|
||||
|
||||
func (emptyReader) ReadByte() (byte, error) { return 0, io.EOF }
|
||||
|
||||
func (emptyReader) Read([]byte) (int, error) { return 0, io.EOF }
|
||||
|
||||
type errorReader struct{ err error }
|
||||
|
||||
func (r *errorReader) Close() error { return r.err }
|
||||
|
||||
func (r *errorReader) Read([]byte) (int, error) { return 0, r.err }
|
||||
|
||||
type errorWriter struct{ err error }
|
||||
|
||||
func (w *errorWriter) Close() error { return w.err }
|
||||
|
||||
func (w *errorWriter) Write([]byte) (int, error) { return 0, w.err }
|
||||
+68
@@ -0,0 +1,68 @@
|
||||
package lz4
|
||||
|
||||
import (
|
||||
"io"
|
||||
"sync"
|
||||
|
||||
"github.com/pierrec/lz4/v4"
|
||||
)
|
||||
|
||||
var (
|
||||
readerPool sync.Pool
|
||||
writerPool sync.Pool
|
||||
)
|
||||
|
||||
// Codec is the implementation of a compress.Codec which supports creating
|
||||
// readers and writers for kafka messages compressed with lz4.
|
||||
type Codec struct{}
|
||||
|
||||
// Code implements the compress.Codec interface.
|
||||
func (c *Codec) Code() int8 { return 3 }
|
||||
|
||||
// Name implements the compress.Codec interface.
|
||||
func (c *Codec) Name() string { return "lz4" }
|
||||
|
||||
// NewReader implements the compress.Codec interface.
|
||||
func (c *Codec) NewReader(r io.Reader) io.ReadCloser {
|
||||
z, _ := readerPool.Get().(*lz4.Reader)
|
||||
if z != nil {
|
||||
z.Reset(r)
|
||||
} else {
|
||||
z = lz4.NewReader(r)
|
||||
}
|
||||
return &reader{Reader: z}
|
||||
}
|
||||
|
||||
// NewWriter implements the compress.Codec interface.
|
||||
func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
|
||||
z, _ := writerPool.Get().(*lz4.Writer)
|
||||
if z != nil {
|
||||
z.Reset(w)
|
||||
} else {
|
||||
z = lz4.NewWriter(w)
|
||||
}
|
||||
return &writer{Writer: z}
|
||||
}
|
||||
|
||||
type reader struct{ *lz4.Reader }
|
||||
|
||||
func (r *reader) Close() (err error) {
|
||||
if z := r.Reader; z != nil {
|
||||
r.Reader = nil
|
||||
z.Reset(nil)
|
||||
readerPool.Put(z)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
type writer struct{ *lz4.Writer }
|
||||
|
||||
func (w *writer) Close() (err error) {
|
||||
if z := w.Writer; z != nil {
|
||||
w.Writer = nil
|
||||
err = z.Close()
|
||||
z.Reset(nil)
|
||||
writerPool.Put(z)
|
||||
}
|
||||
return
|
||||
}
|
||||
+110
@@ -0,0 +1,110 @@
|
||||
package snappy
|
||||
|
||||
import (
|
||||
"io"
|
||||
"sync"
|
||||
|
||||
"github.com/klauspost/compress/s2"
|
||||
"github.com/klauspost/compress/snappy"
|
||||
)
|
||||
|
||||
// Framing is an enumeration type used to enable or disable xerial framing of
|
||||
// snappy messages.
|
||||
type Framing int
|
||||
|
||||
const (
|
||||
Framed Framing = iota
|
||||
Unframed
|
||||
)
|
||||
|
||||
// Compression level.
|
||||
type Compression int
|
||||
|
||||
const (
|
||||
DefaultCompression Compression = iota
|
||||
FasterCompression
|
||||
BetterCompression
|
||||
BestCompression
|
||||
)
|
||||
|
||||
var (
|
||||
readerPool sync.Pool
|
||||
writerPool sync.Pool
|
||||
)
|
||||
|
||||
// Codec is the implementation of a compress.Codec which supports creating
|
||||
// readers and writers for kafka messages compressed with snappy.
|
||||
type Codec struct {
|
||||
// An optional framing to apply to snappy compression.
|
||||
//
|
||||
// Default to Framed.
|
||||
Framing Framing
|
||||
|
||||
// Compression level.
|
||||
Compression Compression
|
||||
}
|
||||
|
||||
// Code implements the compress.Codec interface.
|
||||
func (c *Codec) Code() int8 { return 2 }
|
||||
|
||||
// Name implements the compress.Codec interface.
|
||||
func (c *Codec) Name() string { return "snappy" }
|
||||
|
||||
// NewReader implements the compress.Codec interface.
|
||||
func (c *Codec) NewReader(r io.Reader) io.ReadCloser {
|
||||
x, _ := readerPool.Get().(*xerialReader)
|
||||
if x != nil {
|
||||
x.Reset(r)
|
||||
} else {
|
||||
x = &xerialReader{
|
||||
reader: r,
|
||||
decode: snappy.Decode,
|
||||
}
|
||||
}
|
||||
return &reader{xerialReader: x}
|
||||
}
|
||||
|
||||
// NewWriter implements the compress.Codec interface.
|
||||
func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
|
||||
x, _ := writerPool.Get().(*xerialWriter)
|
||||
if x != nil {
|
||||
x.Reset(w)
|
||||
} else {
|
||||
x = &xerialWriter{writer: w}
|
||||
}
|
||||
x.framed = c.Framing == Framed
|
||||
switch c.Compression {
|
||||
case FasterCompression:
|
||||
x.encode = s2.EncodeSnappy
|
||||
case BetterCompression:
|
||||
x.encode = s2.EncodeSnappyBetter
|
||||
case BestCompression:
|
||||
x.encode = s2.EncodeSnappyBest
|
||||
default:
|
||||
x.encode = snappy.Encode // aka. s2.EncodeSnappyBetter
|
||||
}
|
||||
return &writer{xerialWriter: x}
|
||||
}
|
||||
|
||||
type reader struct{ *xerialReader }
|
||||
|
||||
func (r *reader) Close() (err error) {
|
||||
if x := r.xerialReader; x != nil {
|
||||
r.xerialReader = nil
|
||||
x.Reset(nil)
|
||||
readerPool.Put(x)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
type writer struct{ *xerialWriter }
|
||||
|
||||
func (w *writer) Close() (err error) {
|
||||
if x := w.xerialWriter; x != nil {
|
||||
w.xerialWriter = nil
|
||||
err = x.Flush()
|
||||
x.Reset(nil)
|
||||
writerPool.Put(x)
|
||||
}
|
||||
return
|
||||
}
|
||||
+330
@@ -0,0 +1,330 @@
|
||||
package snappy
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/binary"
|
||||
"errors"
|
||||
"io"
|
||||
|
||||
"github.com/klauspost/compress/snappy"
|
||||
)
|
||||
|
||||
const defaultBufferSize = 32 * 1024
|
||||
|
||||
// An implementation of io.Reader which consumes a stream of xerial-framed
|
||||
// snappy-encoeded data. The framing is optional, if no framing is detected
|
||||
// the reader will simply forward the bytes from its underlying stream.
|
||||
type xerialReader struct {
|
||||
reader io.Reader
|
||||
header [16]byte
|
||||
input []byte
|
||||
output []byte
|
||||
offset int64
|
||||
nbytes int64
|
||||
decode func([]byte, []byte) ([]byte, error)
|
||||
}
|
||||
|
||||
func (x *xerialReader) Reset(r io.Reader) {
|
||||
x.reader = r
|
||||
x.input = x.input[:0]
|
||||
x.output = x.output[:0]
|
||||
x.header = [16]byte{}
|
||||
x.offset = 0
|
||||
x.nbytes = 0
|
||||
}
|
||||
|
||||
func (x *xerialReader) Read(b []byte) (int, error) {
|
||||
for {
|
||||
if x.offset < int64(len(x.output)) {
|
||||
n := copy(b, x.output[x.offset:])
|
||||
x.offset += int64(n)
|
||||
return n, nil
|
||||
}
|
||||
|
||||
n, err := x.readChunk(b)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
if n > 0 {
|
||||
return n, nil
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (x *xerialReader) WriteTo(w io.Writer) (int64, error) {
|
||||
wn := int64(0)
|
||||
|
||||
for {
|
||||
for x.offset < int64(len(x.output)) {
|
||||
n, err := w.Write(x.output[x.offset:])
|
||||
wn += int64(n)
|
||||
x.offset += int64(n)
|
||||
if err != nil {
|
||||
return wn, err
|
||||
}
|
||||
}
|
||||
|
||||
if _, err := x.readChunk(nil); err != nil {
|
||||
if errors.Is(err, io.EOF) {
|
||||
err = nil
|
||||
}
|
||||
return wn, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (x *xerialReader) readChunk(dst []byte) (int, error) {
|
||||
x.output = x.output[:0]
|
||||
x.offset = 0
|
||||
prefix := 0
|
||||
|
||||
if x.nbytes == 0 {
|
||||
n, err := x.readFull(x.header[:])
|
||||
if err != nil && n == 0 {
|
||||
return 0, err
|
||||
}
|
||||
prefix = n
|
||||
}
|
||||
|
||||
if isXerialHeader(x.header[:]) {
|
||||
if cap(x.input) < 4 {
|
||||
x.input = make([]byte, 4, defaultBufferSize)
|
||||
} else {
|
||||
x.input = x.input[:4]
|
||||
}
|
||||
|
||||
_, err := x.readFull(x.input)
|
||||
if err != nil {
|
||||
return 0, err
|
||||
}
|
||||
|
||||
frame := int(binary.BigEndian.Uint32(x.input))
|
||||
if cap(x.input) < frame {
|
||||
x.input = make([]byte, frame, align(frame, defaultBufferSize))
|
||||
} else {
|
||||
x.input = x.input[:frame]
|
||||
}
|
||||
|
||||
if _, err := x.readFull(x.input); err != nil {
|
||||
return 0, err
|
||||
}
|
||||
} else {
|
||||
if cap(x.input) == 0 {
|
||||
x.input = make([]byte, 0, defaultBufferSize)
|
||||
} else {
|
||||
x.input = x.input[:0]
|
||||
}
|
||||
|
||||
if prefix > 0 {
|
||||
x.input = append(x.input, x.header[:prefix]...)
|
||||
}
|
||||
|
||||
for {
|
||||
if len(x.input) == cap(x.input) {
|
||||
b := make([]byte, len(x.input), 2*cap(x.input))
|
||||
copy(b, x.input)
|
||||
x.input = b
|
||||
}
|
||||
|
||||
n, err := x.read(x.input[len(x.input):cap(x.input)])
|
||||
x.input = x.input[:len(x.input)+n]
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) && len(x.input) > 0 {
|
||||
break
|
||||
}
|
||||
return 0, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
var n int
|
||||
var err error
|
||||
|
||||
if x.decode == nil {
|
||||
x.output, x.input, err = x.input, x.output, nil
|
||||
} else if n, err = snappy.DecodedLen(x.input); n <= len(dst) && err == nil {
|
||||
// If the output buffer is large enough to hold the decode value,
|
||||
// write it there directly instead of using the intermediary output
|
||||
// buffer.
|
||||
_, err = x.decode(dst, x.input)
|
||||
} else {
|
||||
var b []byte
|
||||
n = 0
|
||||
b, err = x.decode(x.output[:cap(x.output)], x.input)
|
||||
if err == nil {
|
||||
x.output = b
|
||||
}
|
||||
}
|
||||
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (x *xerialReader) read(b []byte) (int, error) {
|
||||
n, err := x.reader.Read(b)
|
||||
x.nbytes += int64(n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (x *xerialReader) readFull(b []byte) (int, error) {
|
||||
n, err := io.ReadFull(x.reader, b)
|
||||
x.nbytes += int64(n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
// An implementation of a xerial-framed snappy-encoded output stream.
|
||||
// Each Write made to the writer is framed with a xerial header.
|
||||
type xerialWriter struct {
|
||||
writer io.Writer
|
||||
header [16]byte
|
||||
input []byte
|
||||
output []byte
|
||||
nbytes int64
|
||||
framed bool
|
||||
encode func([]byte, []byte) []byte
|
||||
}
|
||||
|
||||
func (x *xerialWriter) Reset(w io.Writer) {
|
||||
x.writer = w
|
||||
x.input = x.input[:0]
|
||||
x.output = x.output[:0]
|
||||
x.nbytes = 0
|
||||
}
|
||||
|
||||
func (x *xerialWriter) ReadFrom(r io.Reader) (int64, error) {
|
||||
wn := int64(0)
|
||||
|
||||
if cap(x.input) == 0 {
|
||||
x.input = make([]byte, 0, defaultBufferSize)
|
||||
}
|
||||
|
||||
for {
|
||||
if x.full() {
|
||||
x.grow()
|
||||
}
|
||||
|
||||
n, err := r.Read(x.input[len(x.input):cap(x.input)])
|
||||
wn += int64(n)
|
||||
x.input = x.input[:len(x.input)+n]
|
||||
|
||||
if x.fullEnough() {
|
||||
if err := x.Flush(); err != nil {
|
||||
return wn, err
|
||||
}
|
||||
}
|
||||
|
||||
if err != nil {
|
||||
if errors.Is(err, io.EOF) {
|
||||
err = nil
|
||||
}
|
||||
return wn, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (x *xerialWriter) Write(b []byte) (int, error) {
|
||||
wn := 0
|
||||
|
||||
if cap(x.input) == 0 {
|
||||
x.input = make([]byte, 0, defaultBufferSize)
|
||||
}
|
||||
|
||||
for len(b) > 0 {
|
||||
if x.full() {
|
||||
x.grow()
|
||||
}
|
||||
|
||||
n := copy(x.input[len(x.input):cap(x.input)], b)
|
||||
b = b[n:]
|
||||
wn += n
|
||||
x.input = x.input[:len(x.input)+n]
|
||||
|
||||
if x.fullEnough() {
|
||||
if err := x.Flush(); err != nil {
|
||||
return wn, err
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
return wn, nil
|
||||
}
|
||||
|
||||
func (x *xerialWriter) Flush() error {
|
||||
if len(x.input) == 0 {
|
||||
return nil
|
||||
}
|
||||
|
||||
var b []byte
|
||||
if x.encode == nil {
|
||||
b = x.input
|
||||
} else {
|
||||
x.output = x.encode(x.output[:cap(x.output)], x.input)
|
||||
b = x.output
|
||||
}
|
||||
|
||||
x.input = x.input[:0]
|
||||
x.output = x.output[:0]
|
||||
|
||||
if x.framed && x.nbytes == 0 {
|
||||
writeXerialHeader(x.header[:])
|
||||
_, err := x.write(x.header[:])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
if x.framed {
|
||||
writeXerialFrame(x.header[:4], len(b))
|
||||
_, err := x.write(x.header[:4])
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
}
|
||||
|
||||
_, err := x.write(b)
|
||||
return err
|
||||
}
|
||||
|
||||
func (x *xerialWriter) write(b []byte) (int, error) {
|
||||
n, err := x.writer.Write(b)
|
||||
x.nbytes += int64(n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (x *xerialWriter) full() bool {
|
||||
return len(x.input) == cap(x.input)
|
||||
}
|
||||
|
||||
func (x *xerialWriter) fullEnough() bool {
|
||||
return x.framed && (cap(x.input)-len(x.input)) < 1024
|
||||
}
|
||||
|
||||
func (x *xerialWriter) grow() {
|
||||
tmp := make([]byte, len(x.input), 2*cap(x.input))
|
||||
copy(tmp, x.input)
|
||||
x.input = tmp
|
||||
}
|
||||
|
||||
func align(n, a int) int {
|
||||
if (n % a) == 0 {
|
||||
return n
|
||||
}
|
||||
return ((n / a) + 1) * a
|
||||
}
|
||||
|
||||
var (
|
||||
xerialHeader = [...]byte{130, 83, 78, 65, 80, 80, 89, 0}
|
||||
xerialVersionInfo = [...]byte{0, 0, 0, 1, 0, 0, 0, 1}
|
||||
)
|
||||
|
||||
func isXerialHeader(src []byte) bool {
|
||||
return len(src) >= 16 && bytes.Equal(src[:8], xerialHeader[:])
|
||||
}
|
||||
|
||||
func writeXerialHeader(b []byte) {
|
||||
copy(b[:8], xerialHeader[:])
|
||||
copy(b[8:], xerialVersionInfo[:])
|
||||
}
|
||||
|
||||
func writeXerialFrame(b []byte, n int) {
|
||||
binary.BigEndian.PutUint32(b, uint32(n))
|
||||
}
|
||||
+168
@@ -0,0 +1,168 @@
|
||||
// Package zstd implements Zstandard compression.
|
||||
package zstd
|
||||
|
||||
import (
|
||||
"io"
|
||||
"sync"
|
||||
|
||||
"github.com/klauspost/compress/zstd"
|
||||
)
|
||||
|
||||
// Codec is the implementation of a compress.Codec which supports creating
|
||||
// readers and writers for kafka messages compressed with zstd.
|
||||
type Codec struct {
|
||||
// The compression level configured on writers created by the codec.
|
||||
//
|
||||
// Default to 3.
|
||||
Level int
|
||||
|
||||
encoderPool sync.Pool // *encoder
|
||||
}
|
||||
|
||||
// Code implements the compress.Codec interface.
|
||||
func (c *Codec) Code() int8 { return 4 }
|
||||
|
||||
// Name implements the compress.Codec interface.
|
||||
func (c *Codec) Name() string { return "zstd" }
|
||||
|
||||
// NewReader implements the compress.Codec interface.
|
||||
func (c *Codec) NewReader(r io.Reader) io.ReadCloser {
|
||||
p := new(reader)
|
||||
if p.dec, _ = decoderPool.Get().(*zstd.Decoder); p.dec != nil {
|
||||
p.dec.Reset(r)
|
||||
} else {
|
||||
z, err := zstd.NewReader(r,
|
||||
zstd.WithDecoderConcurrency(1),
|
||||
)
|
||||
if err != nil {
|
||||
p.err = err
|
||||
} else {
|
||||
p.dec = z
|
||||
}
|
||||
}
|
||||
return p
|
||||
}
|
||||
|
||||
func (c *Codec) level() int {
|
||||
if c.Level != 0 {
|
||||
return c.Level
|
||||
}
|
||||
return 3
|
||||
}
|
||||
|
||||
func (c *Codec) zstdLevel() zstd.EncoderLevel {
|
||||
return zstd.EncoderLevelFromZstd(c.level())
|
||||
}
|
||||
|
||||
var decoderPool sync.Pool // *zstd.Decoder
|
||||
|
||||
type reader struct {
|
||||
dec *zstd.Decoder
|
||||
err error
|
||||
}
|
||||
|
||||
// Close implements the io.Closer interface.
|
||||
func (r *reader) Close() error {
|
||||
if r.dec != nil {
|
||||
r.dec.Reset(devNull{}) // don't retain the underlying reader
|
||||
decoderPool.Put(r.dec)
|
||||
r.dec = nil
|
||||
r.err = io.ErrClosedPipe
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// Read implements the io.Reader interface.
|
||||
func (r *reader) Read(p []byte) (int, error) {
|
||||
if r.err != nil {
|
||||
return 0, r.err
|
||||
}
|
||||
if r.dec == nil {
|
||||
return 0, io.EOF
|
||||
}
|
||||
return r.dec.Read(p)
|
||||
}
|
||||
|
||||
// WriteTo implements the io.WriterTo interface.
|
||||
func (r *reader) WriteTo(w io.Writer) (int64, error) {
|
||||
if r.err != nil {
|
||||
return 0, r.err
|
||||
}
|
||||
if r.dec == nil {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
return r.dec.WriteTo(w)
|
||||
}
|
||||
|
||||
// NewWriter implements the compress.Codec interface.
|
||||
func (c *Codec) NewWriter(w io.Writer) io.WriteCloser {
|
||||
p := new(writer)
|
||||
if enc, _ := c.encoderPool.Get().(*zstd.Encoder); enc == nil {
|
||||
z, err := zstd.NewWriter(w,
|
||||
zstd.WithEncoderLevel(c.zstdLevel()),
|
||||
zstd.WithEncoderConcurrency(1),
|
||||
zstd.WithZeroFrames(true),
|
||||
)
|
||||
if err != nil {
|
||||
p.err = err
|
||||
} else {
|
||||
p.enc = z
|
||||
}
|
||||
} else {
|
||||
p.enc = enc
|
||||
p.enc.Reset(w)
|
||||
}
|
||||
p.c = c
|
||||
return p
|
||||
}
|
||||
|
||||
type writer struct {
|
||||
c *Codec
|
||||
enc *zstd.Encoder
|
||||
err error
|
||||
}
|
||||
|
||||
// Close implements the io.Closer interface.
|
||||
func (w *writer) Close() error {
|
||||
if w.enc != nil {
|
||||
// Close needs to be called to write the end of stream marker and flush
|
||||
// the buffers. The zstd package documents that the encoder is re-usable
|
||||
// after being closed.
|
||||
err := w.enc.Close()
|
||||
if err != nil {
|
||||
w.err = err
|
||||
}
|
||||
w.enc.Reset(devNull{}) // don't retain the underlying writer
|
||||
w.c.encoderPool.Put(w.enc)
|
||||
w.enc = nil
|
||||
return err
|
||||
}
|
||||
return w.err
|
||||
}
|
||||
|
||||
// WriteTo implements the io.WriterTo interface.
|
||||
func (w *writer) Write(p []byte) (int, error) {
|
||||
if w.err != nil {
|
||||
return 0, w.err
|
||||
}
|
||||
if w.enc == nil {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
return w.enc.Write(p)
|
||||
}
|
||||
|
||||
// ReadFrom implements the io.ReaderFrom interface.
|
||||
func (w *writer) ReadFrom(r io.Reader) (int64, error) {
|
||||
if w.err != nil {
|
||||
return 0, w.err
|
||||
}
|
||||
if w.enc == nil {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
return w.enc.ReadFrom(r)
|
||||
}
|
||||
|
||||
type devNull struct{}
|
||||
|
||||
func (devNull) Read([]byte) (int, error) { return 0, io.EOF }
|
||||
func (devNull) Write([]byte) (int, error) { return 0, nil }
|
||||
Reference in New Issue
Block a user