feat: [shard-distributor]Compress data when writing to ETCD

cmd/server/go.mod

@@ -120,7 +120,7 @@ require (
 	github.com/golang-jwt/jwt/v5 v5.2.0 // indirect
 	github.com/golang/groupcache v0.0.0-20210331224755-41bb18bfe9da // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
-	github.com/golang/snappy v0.0.4 // indirect
+	github.com/golang/snappy v1.0.0 // indirect
 	github.com/google/s2a-go v0.1.4 // indirect
 	github.com/googleapis/enterprise-certificate-proxy v0.2.4 // indirect
 	github.com/googleapis/gax-go/v2 v2.12.0 // indirect

cmd/server/go.sum

@@ -169,8 +169,8 @@ github.com/golang/protobuf v1.5.4 h1:i7eJL8qZTpSEXOPTxNKhASYpMn+8e5Q6AdndVa1dWek
 github.com/golang/protobuf v1.5.4/go.mod h1:lnTiLA8Wa4RWRcIUkrtSVa5nRhsEGBg48fD6rSs7xps=
 github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/golang/snappy v0.0.3/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
-github.com/golang/snappy v0.0.4 h1:yAGX7huGHXlcLOEtBnF4w7FQwA26wojNCwOYAEhLjQM=
-github.com/golang/snappy v0.0.4/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
+github.com/golang/snappy v1.0.0 h1:Oy607GVXHs7RtbggtPBnr2RmDArIsAefDwvrdWvRhGs=
+github.com/golang/snappy v1.0.0/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
 github.com/google/go-cmp v0.2.0/go.mod h1:oXzfMopK8JAjlY9xF4vHSVASa0yLyX7SntLO5aqRK0M=
 github.com/google/go-cmp v0.3.0/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=
 github.com/google/go-cmp v0.3.1/go.mod h1:8QqcDgzrUqlUb/G2PQTWiueGozuR1884gddMywk6iLU=

common/config/config.go

@@ -631,11 +631,12 @@ type (
 	// ShardDistribution is a configuration for leader election running.
 	// This configuration should be in sync with sharddistributor.
 	ShardDistribution struct {
-		LeaderStore Store         `yaml:"leaderStore"`
-		Election    Election      `yaml:"election"`
-		Namespaces  []Namespace   `yaml:"namespaces"`
-		Process     LeaderProcess `yaml:"process"`
-		Store       Store         `yaml:"store"`
+		LeaderStore     Store         `yaml:"leaderStore"`
+		Election        Election      `yaml:"election"`
+		Namespaces      []Namespace   `yaml:"namespaces"`
+		Process         LeaderProcess `yaml:"process"`
+		Store           Store         `yaml:"store"`
+		DataCompression bool          `yaml:"dataCompression" default:"true"`
 	}
 
 	// Store is a generic container for any storage configuration that should be parsed by the implementation.

config/development.yaml

@@ -186,3 +186,4 @@ shardDistribution:
   process:
     period: 1s
     heartbeatTTL: 2s
+  dataCompression: false

service/sharddistributor/config/config.go

@@ -49,11 +49,12 @@ type (
 
 	// ShardDistribution is a configuration for leader election running.
 	ShardDistribution struct {
-		LeaderStore Store         `yaml:"leaderStore"`
-		Election    Election      `yaml:"election"`
-		Namespaces  []Namespace   `yaml:"namespaces"`
-		Process     LeaderProcess `yaml:"process"`
-		Store       Store         `yaml:"store"`
+		LeaderStore     Store         `yaml:"leaderStore"`
+		Election        Election      `yaml:"election"`
+		Namespaces      []Namespace   `yaml:"namespaces"`
+		Process         LeaderProcess `yaml:"process"`
+		Store           Store         `yaml:"store"`
+		DataCompression bool          `yaml:"dataCompression" default:"true"`
 	}
 
 	// Store is a generic container for any storage configuration that should be parsed by the implementation.
@@ -132,10 +133,11 @@ func GetShardDistributionFromExternal(in config.ShardDistribution) ShardDistribu
 	}
 
 	return ShardDistribution{
-		LeaderStore: Store(in.LeaderStore),
-		Store:       Store(in.Store),
-		Election:    Election(in.Election),
-		Namespaces:  namespaces,
-		Process:     LeaderProcess(in.Process),
+		LeaderStore:     Store(in.LeaderStore),
+		Store:           Store(in.Store),
+		Election:        Election(in.Election),
+		Namespaces:      namespaces,
+		Process:         LeaderProcess(in.Process),
+		DataCompression: in.DataCompression,
 	}
 }

service/sharddistributor/store/etcd/executorstore/common/compression.go

@@ -0,0 +1,79 @@
+package common
+
+import (
+	"bytes"
+	"encoding/json"
+	"fmt"
+	"io"
+
+	"github.com/golang/snappy"
+
+	"github.com/uber/cadence/common/log"
+	"github.com/uber/cadence/common/types"
+)
+
+var (
+	// snappyMagic is a magic prefix prepended to compressed data to distinguish it from uncompressed data
+	snappyMagic = []byte{0xff, 0x06, 0x00, 0x00, 's', 'N', 'a', 'P', 'p', 'Y'}
+)
+
+// Compress compresses data using snappy's framed format if compression is enabled
+// If compressionEnabled returns false, returns the data as-is without compression
+func Compress(data []byte, compressionEnabled bool) ([]byte, error) {
+	if !compressionEnabled {
+		return data, nil
+	}
+
+	var buf bytes.Buffer
+	w := snappy.NewBufferedWriter(&buf)
+
+	if _, err := w.Write(data); err != nil {
+		return nil, err
+	}
+	if err := w.Close(); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+// Decompress decodes snappy-compressed data
+// If the snappy header is present, it will successfully decompress it or return an error
+// If the snappy header is absent, it treats data as uncompressed and returns it as-is
+func Decompress(data []byte) ([]byte, error) {
+	if len(data) == 0 {
+		return data, nil
+	}
+
+	if !hasFramedHeader(data) {
+		return data, nil
+	}
+	r := snappy.NewReader(bytes.NewReader(data))
+	decompressed, err := io.ReadAll(r)
+	if err != nil {
+		return nil, err
+	}
+	return decompressed, nil
+}
+
+// CompressedActiveStatus returns the compressed active status string.
+func CompressedActiveStatus(compressionEnabled bool) string {
+	compressed, _ := Compress([]byte(fmt.Sprintf(`"%s"`, types.ExecutorStatusACTIVE)), compressionEnabled)
+	return string(compressed)
+}
+
+// DecompressAndUnmarshal decompresses data and unmarshals it into the target
+// errorContext is used to provide meaningful error messages
+func DecompressAndUnmarshal(data []byte, target interface{}, errorContext string, logger log.Logger) error {
+	decompressed, err := Decompress(data)
+	if err != nil {
+		return fmt.Errorf("decompress %s: %w", errorContext, err)
+	}
+	if err := json.Unmarshal(decompressed, target); err != nil {
+		return fmt.Errorf("unmarshal %s: %w", errorContext, err)
+	}
+	return nil
+}
+
+func hasFramedHeader(b []byte) bool {
+	return len(b) >= len(snappyMagic) && bytes.Equal(b[:len(snappyMagic)], snappyMagic)
+}

service/sharddistributor/store/etcd/executorstore/common/compression_test.go

@@ -0,0 +1,169 @@
+package common
+
+import (
+	"encoding/json"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/uber/cadence/common/log/testlogger"
+	"github.com/uber/cadence/common/types"
+)
+
+func TestCompressDecompress(t *testing.T) {
+	original := []byte(`{"status":"ACTIVE","shards":["shard1","shard2"]}`)
+
+	compressed, err := Compress(original, true)
+	require.NoError(t, err)
+	require.NotNil(t, compressed)
+
+	assert.NotEqual(t, original, compressed)
+
+	decompressed, err := Decompress(compressed)
+	require.NoError(t, err)
+	assert.Equal(t, original, decompressed)
+}
+
+func TestCompressWithDisabledFlag(t *testing.T) {
+	original := []byte(`{"status":"ACTIVE","shards":["shard1","shard2"]}`)
+
+	result, err := Compress(original, false)
+	require.NoError(t, err)
+	assert.Equal(t, original, result, "Should return original data when compression is disabled")
+
+	var status map[string]interface{}
+	err = json.Unmarshal(result, &status)
+	require.NoError(t, err)
+	assert.Equal(t, "ACTIVE", status["status"])
+}
+
+func TestDecompress(t *testing.T) {
+	t.Run("Empty data", func(t *testing.T) {
+		decompressed, err := Decompress([]byte{})
+		require.NoError(t, err)
+		assert.Empty(t, decompressed)
+	})
+
+	t.Run("Nil data", func(t *testing.T) {
+		decompressed, err := Decompress(nil)
+		require.NoError(t, err)
+		assert.Nil(t, decompressed)
+	})
+
+	t.Run("Uncompressed data", func(t *testing.T) {
+		uncompressed := []byte(`{"status":"ACTIVE"}`)
+
+		result, err := Decompress(uncompressed)
+		require.NoError(t, err)
+		assert.Equal(t, uncompressed, result, "Uncompressed data is returned as-is")
+
+		var status map[string]string
+		err = json.Unmarshal(result, &status)
+		require.NoError(t, err)
+		assert.Equal(t, "ACTIVE", status["status"])
+	})
+
+	t.Run("Compressed data", func(t *testing.T) {
+		original := []byte(`{"status":"DRAINING"}`)
+		compressed, err := Compress(original, true)
+		require.NoError(t, err)
+
+		result, err := Decompress(compressed)
+		require.NoError(t, err)
+		assert.Equal(t, original, result)
+
+		var status map[string]string
+		err = json.Unmarshal(result, &status)
+		require.NoError(t, err)
+		assert.Equal(t, "DRAINING", status["status"])
+	})
+}
+
+func TestDecompressAndUnmarshal(t *testing.T) {
+	type testData struct {
+		Status string   `json:"status"`
+		Shards []string `json:"shards"`
+	}
+	logger := testlogger.New(t)
+
+	t.Run("Uncompressed data", func(t *testing.T) {
+		data := []byte(`{"status":"ACTIVE","shards":["shard1","shard2"]}`)
+
+		var result testData
+		err := DecompressAndUnmarshal(data, &result, "test data", logger)
+		require.NoError(t, err)
+		assert.Equal(t, "ACTIVE", result.Status)
+		assert.Equal(t, []string{"shard1", "shard2"}, result.Shards)
+	})
+
+	t.Run("Compressed data", func(t *testing.T) {
+		original := testData{
+			Status: "DRAINING",
+			Shards: []string{"shard3", "shard4"},
+		}
+		originalJSON, _ := json.Marshal(original)
+		compressed, err := Compress(originalJSON, true)
+		require.NoError(t, err)
+
+		var result testData
+		err = DecompressAndUnmarshal(compressed, &result, "test data", logger)
+		require.NoError(t, err)
+		assert.Equal(t, original.Status, result.Status)
+		assert.Equal(t, original.Shards, result.Shards)
+	})
+
+	t.Run("Invalid JSON in uncompressed data", func(t *testing.T) {
+		invalidJSON := []byte(`{invalid json}`)
+
+		var result testData
+		err := DecompressAndUnmarshal(invalidJSON, &result, "test data", logger)
+		require.Error(t, err)
+		assert.Contains(t, err.Error(), "unmarshal test data")
+	})
+}
+
+func TestCompressedActiveStatus(t *testing.T) {
+	t.Run("Compression enabled", func(t *testing.T) {
+		compressed := CompressedActiveStatus(true)
+		require.NotEmpty(t, compressed)
+
+		decompressed, err := Decompress([]byte(compressed))
+		require.NoError(t, err)
+
+		var status types.ExecutorStatus
+		err = json.Unmarshal(decompressed, &status)
+		require.NoError(t, err)
+		assert.Equal(t, types.ExecutorStatusACTIVE, status)
+	})
+
+	t.Run("Compression disabled", func(t *testing.T) {
+		uncompressed := CompressedActiveStatus(false)
+		require.NotEmpty(t, uncompressed)
+
+		var status types.ExecutorStatus
+		err := json.Unmarshal([]byte(uncompressed), &status)
+		require.NoError(t, err)
+		assert.Equal(t, types.ExecutorStatusACTIVE, status)
+	})
+}
+
+func TestHasFramedHeader(t *testing.T) {
+	t.Run("Data with header", func(t *testing.T) {
+		data := append(snappyMagic, []byte("some data")...)
+		assert.True(t, hasFramedHeader(data))
+	})
+
+	t.Run("Data without header", func(t *testing.T) {
+		data := []byte(`{"json":"data"}`)
+		assert.False(t, hasFramedHeader(data))
+	})
+
+	t.Run("Empty data", func(t *testing.T) {
+		assert.False(t, hasFramedHeader([]byte{}))
+	})
+
+	t.Run("Data shorter than header", func(t *testing.T) {
+		assert.False(t, hasFramedHeader([]byte{0xff, 0x06}))
+	})
+}