Repository: gocql/gocql
Pull request #1561 (2021-05-27 08:14:10)

Add support for LZ4 compression
The subject says it all.


Patch:
diff --git a/.travis.yml b/.travis.yml
index 5ccd938..491552a 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -21,14 +21,25 @@ env:
   matrix:
     - CASS=2.1.21
       AUTH=true
+      COMPRESSOR=snappy
     - CASS=2.2.14
       AUTH=true
+      COMPRESSOR=snappy
     - CASS=2.2.14
       AUTH=false
+      COMPRESSOR=snappy
+    - CASS=2.2.14
+      AUTH=false
+      COMPRESSOR=lz4
     - CASS=3.0.18
       AUTH=false
+      COMPRESSOR=snappy
+    - CASS=3.11.4
+      AUTH=false
+      COMPRESSOR=snappy
     - CASS=3.11.4
       AUTH=false
+      COMPRESSOR=lz4
 
 go:
   - 1.15.x
@@ -44,7 +55,7 @@ install:
 
 script:
   - set -e
-  - PATH=$PATH:$HOME/.local/bin bash integration.sh $CASS $AUTH
+  - PATH=$PATH:$HOME/.local/bin bash integration.sh $CASS $AUTH $COMPRESSOR
   - go vet .
 
 notifications:
diff --git a/AUTHORS b/AUTHORS
index 08cdd00..d77eaad 100644
--- a/AUTHORS
+++ b/AUTHORS
@@ -121,3 +121,4 @@ Jaume Marhuenda Beltran <jaumemarhuenda@gmail.com>
 Piotr Dulikowski <piodul@scylladb.com>
 Árni Dagur <arni@dagur.eu>
 Tushar Das <tushar.das5@gmail.com>
+Maxim Vladimirskiy <horkhe@gmail.com>
diff --git a/compressor.go b/compressor.go
index 26853ae..893d8f5 100644
--- a/compressor.go
+++ b/compressor.go
@@ -1,7 +1,11 @@
 package gocql
 
 import (
+	"encoding/binary"
+	"fmt"
+
 	"github.com/golang/snappy"
+	"github.com/pierrec/lz4/v4"
 )
 
 type Compressor interface {
@@ -26,3 +30,46 @@ func (s SnappyCompressor) Encode(data []byte) ([]byte, error) {
 func (s SnappyCompressor) Decode(data []byte) ([]byte, error) {
 	return snappy.Decode(nil, data)
 }
+
+// LZ4Compressor implements the gocql.Compressor interface and can be used to
+// compress incoming and outgoing frames. According to the Cassandra docs the
+// LZ4 protocol should be preferred over snappy. (For details refer to
+// https://cassandra.apache.org/doc/latest/operating/compression.html)
+//
+// Implementation note: Cassandra prefixes each compressed block with 4 bytes
+// of the uncompressed block length, written in big endian order. But the LZ4
+// compression library github.com/pierrec/lz4/v4 does not expect the length
+// field, so it needs to be added to compressed blocks sent to Cassandra, and
+// removed from ones received from Cassandra before decompression.
+type LZ4Compressor struct{}
+
+func (s LZ4Compressor) Name() string {
+	return "lz4"
+}
+
+func (s LZ4Compressor) Encode(data []byte) ([]byte, error) {
+	buf := make([]byte, lz4.CompressBlockBound(len(data)+4))
+	var compressor lz4.Compressor
+	n, err := compressor.CompressBlock(data, buf[4:])
+	// According to lz4.CompressBlock doc, it doesn't fail as long as the dst
+	// buffer length is at least lz4.CompressBlockBound(len(data))) bytes, but
+	// we check for error anyway just to be thorough.
+	if err != nil {
+		return nil, err
+	}
+	binary.BigEndian.PutUint32(buf, uint32(len(data)))
+	return buf[:n+4], nil
+}
+
+func (s LZ4Compressor) Decode(data []byte) ([]byte, error) {
+	if len(data) < 4 {
+		return nil, fmt.Errorf("cassandra lz4 block size should be >4, got=%d", len(data))
+	}
+	uncompressedLength := binary.BigEndian.Uint32(data)
+	if uncompressedLength == 0 {
+		return nil, nil
+	}
+	buf := make([]byte, uncompressedLength)
+	n, err := lz4.UncompressBlock(data[4:], buf)
+	return buf[:n], err
+}
diff --git a/go.mod b/go.mod
index 70e98c8..d8fd326 100644
--- a/go.mod
+++ b/go.mod
@@ -3,10 +3,11 @@ module github.com/gocql/gocql
 require (
 	github.com/bitly/go-hostpool v0.0.0-20171023180738-a3a6125de932 // indirect
 	github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 // indirect
-	github.com/golang/snappy v0.0.0-20170215233205-553a64147049
+	github.com/golang/snappy v0.0.3
 	github.com/hailocab/go-hostpool v0.0.0-20160125115350-e80d13ce29ed
 	github.com/kr/pretty v0.1.0 // indirect
-	github.com/stretchr/testify v1.3.0 // indirect
+	github.com/pierrec/lz4/v4 v4.1.7
+	github.com/stretchr/testify v1.3.0
 	gopkg.in/inf.v0 v0.9.1
 )
 
diff --git a/go.sum b/go.sum
index a2bcaf7..975861a 100644
--- a/go.sum
+++ b/go.sum
@@ -4,8 +4,8 @@ github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869 h1:DDGfHa7BWjL4Yn
 github.com/bmizerany/assert v0.0.0-20160611221934-b7ed37b82869/go.mod h1:Ekp36dRnpXw/yCqJaO+ZrUyxD+3VXMFFr56k5XYrpB4=
 github.com/davecgh/go-spew v1.1.0 h1:ZDRjVQ15GmhC3fiQ8ni8+OwkZQO4DARzQgrnXU1Liz8=
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/golang/snappy v0.0.0-20170215233205-553a64147049 h1:K9KHZbXKpGydfDN0aZrsoHpLJlZsBrGMFWbgLDGnPZk=
-github.com/golang/snappy v0.0.0-20170215233205-553a64147049/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
+github.com/golang/snappy v0.0.3 h1:fHPg5GQYlCeLIPB9BZqMVR5nR9A+IM5zcgeTdjMYmLA=
+github.com/golang/snappy v0.0.3/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/hailocab/go-hostpool v0.0.0-20160125115350-e80d13ce29ed h1:5upAirOpQc1Q53c0bnx2ufif5kANL7bfZWcc6VJWJd8=
 github.com/hailocab/go-hostpool v0.0.0-20160125115350-e80d13ce29ed/go.mod h1:tMWxXQ9wFIaZeTI9F+hmhFiGpFmhOHzyShyFUhRm0H4=
 github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=
@@ -13,6 +13,10 @@ github.com/kr/pretty v0.1.0/go.mod h1:dAy3ld7l9f0ibDNOQOHHMYYIIbhfbHSm3C4ZsoJORN
 github.com/kr/pty v1.1.1/go.mod h1:pFQYn66WHrOpPYNljwOMqo10TkYh1fy3cYio2l3bCsQ=
 github.com/kr/text v0.1.0 h1:45sCR5RtlFHMR4UwH9sdQ5TC8v0qDQCHnXt+kaKSTVE=
 github.com/kr/text v0.1.0/go.mod h1:4Jbv+DJW3UT/LiOwJeYQe1efqtUx/iVham/4vfdArNI=
+github.com/pierrec/lz4 v1.0.1 h1:w6GMGWSsCI04fTM8wQRdnW74MuJISakuUU0onU0TYB4=
+github.com/pierrec/lz4 v2.6.0+incompatible h1:Ix9yFKn1nSPBLFl/yZknTp8TU5G4Ps0JDmguYK6iH1A=
+github.com/pierrec/lz4/v4 v4.1.7 h1:UDV9geJWhFIufAliH7HQlz9wP3JA0t748w+RwbWMLow=
+github.com/pierrec/lz4/v4 v4.1.7/go.mod h1:gZWDp/Ze/IJXGXf23ltt2EXimqmTUXEy0GFuRQyBid4=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
diff --git a/integration.sh b/integration.sh
index 1c6311b..d23fd7e 100755
--- a/integration.sh
+++ b/integration.sh
@@ -6,6 +6,7 @@ function run_tests() {
 	local clusterSize=3
 	local version=$1
 	local auth=$2
+	local compressor=$3
 
 	if [ "$auth" = true ]; then
 		clusterSize=1
@@ -64,7 +65,7 @@ function run_tests() {
 	ccm status
 	ccm node1 nodetool status
 
-	local args="-gocql.timeout=60s -runssl -proto=$proto -rf=3 -clusterSize=$clusterSize -autowait=2000ms -compressor=snappy -gocql.cversion=$version -cluster=$(ccm liveset) ./..."
+	local args="-gocql.timeout=60s -runssl -proto=$proto -rf=3 -clusterSize=$clusterSize -autowait=2000ms -compressor=$compressor -gocql.cversion=$version -cluster=$(ccm liveset) ./..."
 
 	go test -v -tags unit -race
 
@@ -92,4 +93,4 @@ function run_tests() {
 	ccm remove
 }
 
-run_tests $1 $2
+run_tests $1 $2 $3


Test patch:
diff --git a/common_test.go b/common_test.go
index 0eb6295..e95cba2 100644
--- a/common_test.go
+++ b/common_test.go
@@ -87,6 +87,8 @@ func createCluster(opts ...func(*ClusterConfig)) *ClusterConfig {
 	switch *flagCompressTest {
 	case "snappy":
 		cluster.Compressor = &SnappyCompressor{}
+	case "lz4":
+		cluster.Compressor = &LZ4Compressor{}
 	case "":
 	default:
 		panic("invalid compressor: " + *flagCompressTest)
diff --git a/compressor_test.go b/compressor_test.go
index 8350939..06bfaf8 100644
--- a/compressor_test.go
+++ b/compressor_test.go
@@ -5,6 +5,7 @@ import (
 	"testing"
 
 	"github.com/golang/snappy"
+	"github.com/stretchr/testify/require"
 )
 
 func TestSnappyCompressor(t *testing.T) {
@@ -36,3 +37,26 @@ func TestSnappyCompressor(t *testing.T) {
 		t.Fatal("failed to match the expected decoded value with the result decoded value.")
 	}
 }
+
+func TestLZ4Compressor(t *testing.T) {
+	var c LZ4Compressor
+	require.Equal(t, "lz4", c.Name())
+
+	_, err := c.Decode([]byte{0, 1, 2})
+	require.EqualError(t, err, "cassandra lz4 block size should be >4, got=3")
+
+	_, err = c.Decode([]byte{0, 1, 2, 4, 5})
+	require.EqualError(t, err, "lz4: invalid source or destination buffer too short")
+
+	// If uncompressed size is zero then nothing is decoded even if present.
+	decoded, err := c.Decode([]byte{0, 0, 0, 0, 5, 7, 8})
+	require.NoError(t, err)
+	require.Nil(t, decoded)
+
+	original := []byte("My Test String")
+	encoded, err := c.Encode(original)
+	require.NoError(t, err)
+	decoded, err = c.Decode(encoded)
+	require.NoError(t, err)
+	require.Equal(t, original, decoded)
+}
diff --git a/install_test_deps.sh b/install_test_deps.sh
index 1484ac4..87f78c2 100755
--- a/install_test_deps.sh
+++ b/install_test_deps.sh
@@ -1,9 +1,11 @@
 #!/usr/bin/env bash
 
+set -x
+
 # This is not supposed to be an error-prone script; just a convenience.
 
 # Install CCM
-pip install --user cql PyYAML six
+pip install -i https://pypi.org/simple --user cql PyYAML six psutil
 git clone https://github.com/pcmanus/ccm.git
 pushd ccm
 ./setup.py install --user