Compression tests

Signed-off-by: Rohit Nayak <rohit@planetscale.com>

go.mod

@@ -93,6 +93,7 @@ require (
 require (
 	github.com/DataDog/datadog-go/v5 v5.5.0
 	github.com/Shopify/toxiproxy/v2 v2.7.0
+	github.com/andybalholm/brotli v1.0.6
 	github.com/bndr/gotabulate v1.1.2
 	github.com/gammazero/deque v0.2.1
 	github.com/google/safehtml v0.1.0

go.sum

@@ -61,6 +61,8 @@ github.com/alecthomas/template v0.0.0-20160405071501-a0175ee3bccc/go.mod h1:LOuy
 github.com/alecthomas/template v0.0.0-20190718012654-fb15b899a751/go.mod h1:LOuyumcjzFXgccqObfd/Ljyb9UuFJ6TxHnclSeseNhc=
 github.com/alecthomas/units v0.0.0-20151022065526-2efee857e7cf/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
 github.com/alecthomas/units v0.0.0-20190717042225-c3de453c63f4/go.mod h1:ybxpYRFXyAe+OPACYpWeL0wqObRcbAqCMya13uyzqw0=
+github.com/andybalholm/brotli v1.0.6 h1:Yf9fFpf49Zrxb9NlQaluyE92/+X7UVHlhMNJN2sxfOI=
+github.com/andybalholm/brotli v1.0.6/go.mod h1:fO7iG3H7G2nSZ7m0zPUDn85XEX2GTukHGRSepvi9Eig=
 github.com/aquarapid/vaultlib v0.5.1 h1:vuLWR6bZzLHybjJBSUYPgZlIp6KZ+SXeHLRRYTuk6d4=
 github.com/aquarapid/vaultlib v0.5.1/go.mod h1:yT7AlEXtuabkxylOc/+Ulyp18tff1+QjgNLTnFWTlOs=
 github.com/armon/circbuf v0.0.0-20150827004946-bbbad097214e/go.mod h1:3U/XgcO3hCbHZ8TKRvWD2dDTCfh9M9ya+I9JpbB7O8o=

go/vt/topotools/routing_rules_test.go

@@ -17,13 +17,28 @@ limitations under the License.
 package topotools
 
 import (
+	"bytes"
+	"compress/gzip"
 	"context"
+	"encoding/gob"
 	"errors"
+	"fmt"
+	"io"
+	"io/ioutil"
+	"math/rand"
+	"sort"
+	"strconv"
+	"strings"
 	"testing"
+	"time"
 
+	"github.com/andybalholm/brotli"
+	"github.com/klauspost/compress/zstd"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"vitess.io/vitess/go/vt/log"
+
 	"vitess.io/vitess/go/vt/topo/memorytopo"
 )
 
@@ -109,3 +124,213 @@ func TestKeyspaceRoutingRulesRoundTrip(t *testing.T) {
 	require.NoError(t, err, "could not fetch keyspace routing rules from topo")
 	assert.EqualValues(t, rulesMap, roundtripRulesMap)
 }
+
+// TestLotsOfKeyspaceRoutingRules checks the size of the rules map for a huge number of rules and the effectiveness
+// of storing the map compressed in the topo.
+func TestLotsOfKeyspaceRoutingRules(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	ts := memorytopo.NewServer(ctx, "zone1")
+	defer ts.Close()
+	const numRules = 10
+	maxVal := int64(9999999999)
+	minVal := int64(9999999999) - 1*numRules
+
+	keyPrefix := ""   //"original_keyspace_"
+	valuePrefix := "" //"target_keyspace_"
+	getKey := func(i int64) string {
+		return fmt.Sprintf("%s%d", keyPrefix, i)
+	}
+	getValue := func(i int64) string {
+		return fmt.Sprintf("%s%d", valuePrefix, i)
+	}
+	rulesMap := make(map[string]string)
+	concatted := strings.Builder{}
+	var keyList []string
+	for j := 0; j < numRules; j++ {
+		i := int64(rand.Intn(int(maxVal-minVal))) + minVal
+		rulesMap[getKey(i)] = getValue(i)
+		concatted.WriteString(fmt.Sprintf("%s:%s,", getKey(i), getValue(i)))
+		keyList = append(keyList, getKey(i))
+	}
+	log.Infof("Key list: %v:%s", keyList, compressRanges(strings.Join(keyList, ",")))
+	t1 := time.Now()
+	err := SaveKeyspaceRoutingRules(ctx, ts, rulesMap)
+	log.Infof("SaveKeyspaceRoutingRules took %v", time.Since(t1))
+	require.NoError(t, err, "could not save keyspace routing rules to topo %v", rulesMap)
+
+	t1 = time.Now()
+	roundtripRulesMap, err := GetKeyspaceRoutingRules(ctx, ts)
+	log.Infof("GetKeyspaceRoutingRules took %v", time.Since(t1))
+	require.NoError(t, err, "could not fetch keyspace routing rules from topo")
+
+	assert.EqualValues(t, rulesMap, roundtripRulesMap)
+	log.Infof("Number of rules: %d", len(rulesMap))
+
+	serialized, err := serializeMap(rulesMap)
+	require.NoError(t, err, "could not serialize keyspace routing rules to topo %v", rulesMap)
+	log.Infof("Size of serialized rules map: %d KB", len(serialized)/1024)
+	compressed, err := compressData(serialized)
+	require.NoError(t, err, "could not compress keyspace routing rules to topo %v", rulesMap)
+	log.Infof("Size of compressed rules map: %d KB", len(compressed)/1024)
+	decompressed, err := decompressData(compressed)
+	require.NoError(t, err, "could not decompress keyspace routing rules to topo %v", rulesMap)
+	log.Infof("Size of decompressed rules map: %d", len(decompressed))
+	deserializedMap, err := deserializeMap(decompressed)
+	log.Infof("Compression ratio percentage: %f, reverse ratio %f",
+		float64(len(compressed))/float64(len(serialized))*100, float64(len(serialized))/float64(len(compressed)))
+	require.NoError(t, err, "could not deserialize keyspace routing rules to topo %v", rulesMap)
+	log.Infof("Size of deserialized rules map: %d", len(deserializedMap))
+	assert.EqualValues(t, rulesMap, deserializedMap)
+	compressed2 := compressed
+	compressed, err = compressData([]byte(concatted.String()))
+	compressed3 := compressed
+	_ = compressed3
+	lst := strings.Join(keyList, ",")
+	compressed, err = compressDataGzip([]byte(lst))
+	require.NoError(t, err)
+	log.Infof("Size of  concatted string: %d KB", len(concatted.String())/1024)
+	log.Infof("Compression ratio percentage for concatted string: %f", float64(concatted.Len())/float64(len(compressed)))
+	log.Infof("Sizes of compressed map: %d KB, compressed concatted : %d KB", len(compressed2)/1024, len(compressed)/1024)
+	log.Infof("Sizes of compressed map: %d KB, compressed list : %d KB", len(compressed2)/1024, len(compressed)/1024)
+
+	lst2 := compressRanges(lst)
+	log.Infof("Size of 'compressed' list %d, original list %d", len(lst2), len(lst))
+	compressed, err = compressDataGzip([]byte(lst))
+	log.Infof("Compression ratio percentage for list: %f", float64(len(lst))/float64(len(lst2)))
+}
+
+// serializeMap serializes a map into a byte slice using gob.
+func serializeMap(data map[string]string) ([]byte, error) {
+	var buf bytes.Buffer
+	enc := gob.NewEncoder(&buf)
+	if err := enc.Encode(data); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+// compressData compresses a byte slice using gzip.
+func compressDataGzip(data []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	gz := gzip.NewWriter(&buf)
+	if _, err := gz.Write(data); err != nil {
+		return nil, err
+	}
+	if err := gz.Close(); err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+// decompressData decompresses a gzip-compressed byte slice.
+func decompressDataGzip(data []byte) ([]byte, error) {
+	buf := bytes.NewBuffer(data)
+	gz, err := gzip.NewReader(buf)
+	if err != nil {
+		return nil, err
+	}
+	defer gz.Close()
+	return ioutil.ReadAll(gz)
+}
+
+// deserializeMap deserializes a byte slice into a map using gob.
+func deserializeMap(data []byte) (map[string]string, error) {
+	var m map[string]string
+	buf := bytes.NewBuffer(data)
+	dec := gob.NewDecoder(buf)
+	if err := dec.Decode(&m); err != nil {
+		return nil, err
+	}
+	return m, nil
+}
+
+// compressData compresses data using Zstandard.
+func compressDataZstd(data []byte) ([]byte, error) {
+	encoder, err := zstd.NewWriter(nil)
+	if err != nil {
+		return nil, err
+	}
+	return encoder.EncodeAll(data, make([]byte, 0, len(data))), nil
+}
+
+// decompressData decompresses data using Zstandard.
+func decompressDataZstd(data []byte) ([]byte, error) {
+	decoder, err := zstd.NewReader(nil)
+	if err != nil {
+		return nil, err
+	}
+	defer decoder.Close()
+	return decoder.DecodeAll(data, nil)
+}
+
+// compressData compresses data using Brotli.
+func compressData(data []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	writer := brotli.NewWriter(&buf)
+	_, err := writer.Write(data)
+	if err != nil {
+		return nil, err
+	}
+	err = writer.Close()
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+// decompressData decompresses Brotli compressed data.
+func decompressData(compressedData []byte) ([]byte, error) {
+	var buf bytes.Buffer
+	reader := brotli.NewReader(bytes.NewReader(compressedData))
+	_, err := io.Copy(&buf, reader)
+	if err != nil {
+		return nil, err
+	}
+	return buf.Bytes(), nil
+}
+
+// compressRanges takes a string of integers separated by commas,
+// compresses consecutive integers into ranges, and returns a single string.
+func compressRanges(input string) string {
+	// Split the input string into a slice of strings.
+	strNumbers := strings.Split(input, ",")
+	numbers := make([]int, len(strNumbers))
+
+	// Convert strings to integers.
+	for i, str := range strNumbers {
+		num, err := strconv.Atoi(str)
+		if err != nil {
+			fmt.Printf("Error converting string to int: %v\n", err)
+			return ""
+		}
+		numbers[i] = num
+	}
+
+	// Sort the slice of integers.
+	sort.Ints(numbers)
+
+	// Iterate through numbers to find ranges.
+	var ranges []string
+	for i := 0; i < len(numbers); {
+		start := numbers[i]
+		end := start
+
+		// Find the end of the current range.
+		for i+1 < len(numbers) && numbers[i+1] == numbers[i]+1 {
+			i++
+			end = numbers[i]
+		}
+
+		// Add the range to the result slice.
+		if start == end {
+			ranges = append(ranges, strconv.Itoa(start))
+		} else {
+			ranges = append(ranges, fmt.Sprintf("%d-%d", start, end))
+		}
+
+		i++
+	}
+	// Join the ranges with commas and return.
+	return strings.Join(ranges, ",")
+}