Add LZ4 compression

protocol/lz4stream.go

@@ -0,0 +1,119 @@
+package protocol
+
+import (
+	"bytes"
+	"encoding/binary"
+	"errors"
+	"io"
+	"sync"
+
+	lz4 "github.com/bkaradzic/go-lz4"
+)
+
+const lz4Magic = 0x5e63b278
+
+type lz4Writer struct {
+	wr  io.Writer
+	mut sync.Mutex
+	buf []byte
+}
+
+func newLZ4Writer(w io.Writer) *lz4Writer {
+	return &lz4Writer{wr: w}
+}
+
+func (w *lz4Writer) Write(bs []byte) (int, error) {
+	w.mut.Lock()
+	defer w.mut.Unlock()
+
+	var err error
+	w.buf, err = lz4.Encode(w.buf[:cap(w.buf)], bs)
+	if err != nil {
+		return 0, err
+	}
+
+	var hdr [8]byte
+	binary.BigEndian.PutUint32(hdr[0:], lz4Magic)
+	binary.BigEndian.PutUint32(hdr[4:], uint32(len(w.buf)))
+	_, err = w.wr.Write(hdr[:])
+	if err != nil {
+		return 0, err
+	}
+
+	_, err = w.wr.Write(w.buf)
+	if err != nil {
+		return 0, err
+	}
+
+	if debug {
+		l.Debugf("lz4 write; %d / %d bytes", len(bs), 8+len(w.buf))
+	}
+	return len(bs), nil
+}
+
+type lz4Reader struct {
+	rd     io.Reader
+	mut    sync.Mutex
+	buf    []byte
+	ebuf   []byte
+	obuf   *bytes.Buffer
+	ibytes uint64
+	obytes uint64
+}
+
+func newLZ4Reader(r io.Reader) *lz4Reader {
+	return &lz4Reader{rd: r}
+}
+
+func (r *lz4Reader) Read(bs []byte) (int, error) {
+	r.mut.Lock()
+	defer r.mut.Unlock()
+
+	if r.obuf == nil {
+		r.obuf = bytes.NewBuffer(nil)
+	}
+
+	if r.obuf.Len() == 0 {
+		if err := r.moreBits(); err != nil {
+			return 0, err
+		}
+	}
+
+	n, err := r.obuf.Read(bs)
+	if debug {
+		l.Debugf("lz4 read; %d bytes", n)
+	}
+	return n, err
+}
+
+func (r *lz4Reader) moreBits() error {
+	var hdr [8]byte
+	_, err := io.ReadFull(r.rd, hdr[:])
+	if binary.BigEndian.Uint32(hdr[0:]) != lz4Magic {
+		return errors.New("bad magic")
+	}
+
+	ln := int(binary.BigEndian.Uint32(hdr[4:]))
+	if len(r.buf) < ln {
+		r.buf = make([]byte, int(ln))
+	} else {
+		r.buf = r.buf[:ln]
+	}
+
+	_, err = io.ReadFull(r.rd, r.buf)
+	if err != nil {
+		return err
+	}
+
+	r.ebuf, err = lz4.Decode(r.ebuf[:cap(r.ebuf)], r.buf)
+	if err != nil {
+		return err
+	}
+
+	if debug {
+		l.Debugf("lz4 moreBits: %d / %d bytes", ln+8, len(r.ebuf))
+	}
+
+	_, err = r.obuf.Write(r.ebuf)
+	return err
+}