rpc-go/codec/compression_codec.go

package codec

import (
	"bytes"
	"compress/gzip"
	"io/ioutil"

	css "git.loafle.net/commons/server-go/socket"
)

type GZipCodec struct {
}

func (_ *GZipCodec) Encode(message []byte) ([]byte, error) {
	var b bytes.Buffer
	gWriter, err := gzip.NewWriterLevel(&b, gzip.BestSpeed)
	if nil != err {
		return nil, err
	}

	if _, err := gWriter.Write(message); nil != err {
		return nil, err
	}
	if err := gWriter.Flush(); nil != err {
		return nil, err
	}
	if err := gWriter.Close(); nil != err {
		return nil, err
	}

	return b.Bytes(), nil
}

func (_ *GZipCodec) Decode(message []byte) ([]byte, error) {
	gReader, err := gzip.NewReader(bytes.NewBuffer(message))
	if nil != err {
		return nil, err
	}
	defer gReader.Close()

	b, err := ioutil.ReadAll(gReader)
	if nil != err {
		return nil, err
	}

	return b, nil
}

var gZipCodec = &GZipCodec{}

func NewCompressionCodecSelector(threshold int) CodecSelector {
	return &compressionCodecSelector{
		threshold: threshold,
	}
}

type compressionCodecSelector struct {
	threshold int
}

func (cs *compressionCodecSelector) Encode(message []byte) (messageType int, encMessage []byte, err error) {
	switch {
	case cs.threshold < len(message):
		messageType = css.BinaryMessage
		encMessage, err = gZipCodec.Encode(message)
	default:
		messageType = css.TextMessage
		encMessage, err = DefaultCodec.Encode(message)
	}

	return
}

func (cs *compressionCodecSelector) Decode(messageType int, message []byte) (decMessage []byte, err error) {
	switch messageType {
	case css.BinaryMessage:
		decMessage, err = gZipCodec.Decode(message)
	default:
		decMessage, err = DefaultCodec.Decode(message)
	}

	return
}