Merge pull request #1420 from databus23/compression

technosophos · web-flow · commit 2eed3f0464ff · 2016-10-20T14:41:51.000-06:00
Add compression to configmap storage driver
diff --git a/pkg/storage/driver/cfgmaps.go b/pkg/storage/driver/cfgmaps.go
@@ -17,8 +17,11 @@ limitations under the License.
 package driver // import "k8s.io/helm/pkg/storage/driver"
 
 import (
+	"bytes"
+	"compress/gzip"
 	"encoding/base64"
 	"fmt"
+	"io/ioutil"
 	"log"
 	"strconv"
 	"time"
@@ -40,6 +43,8 @@ const ConfigMapsDriverName = "ConfigMap"
 
 var b64 = base64.StdEncoding
 
+var magicGzip = []byte{0x1f, 0x8b, 0x08}
+
 // ConfigMaps is a wrapper around an implementation of a kubernetes
 // ConfigMapsInterface.
 type ConfigMaps struct {
@@ -254,13 +259,23 @@ func newConfigMapsObject(key string, rls *rspb.Release, lbs labels) (*api.Config
 }
 
 // encodeRelease encodes a release returning a base64 encoded
-// binary protobuf encoding representation, or error.
+// gzipped binary protobuf encoding representation, or error.
 func encodeRelease(rls *rspb.Release) (string, error) {
 	b, err := proto.Marshal(rls)
 	if err != nil {
 		return "", err
 	}
-	return b64.EncodeToString(b), nil
+	var buf bytes.Buffer
+	w, err := gzip.NewWriterLevel(&buf, gzip.BestCompression)
+	if err != nil {
+		return "", err
+	}
+	if _, err = w.Write(b); err != nil {
+		return "", err
+	}
+	w.Close()
+
+	return b64.EncodeToString(buf.Bytes()), nil
 }
 
 // decodeRelease decodes the bytes in data into a release
@@ -274,6 +289,21 @@ func decodeRelease(data string) (*rspb.Release, error) {
 		return nil, err
 	}
 
+	// For backwards compatibility with releases that were stored before
+	// compression was introduced we skip decompression if the
+	// gzip magic header is not found
+	if bytes.Equal(b[0:3], magicGzip) {
+		r, err := gzip.NewReader(bytes.NewReader(b))
+		if err != nil {
+			return nil, err
+		}
+		b2, err := ioutil.ReadAll(r)
+		if err != nil {
+			return nil, err
+		}
+		b = b2
+	}
+
 	var rls rspb.Release
 	// unmarshal protobuf bytes
 	if err := proto.Unmarshal(b, &rls); err != nil {
diff --git a/pkg/storage/driver/cfgmaps_test.go b/pkg/storage/driver/cfgmaps_test.go
@@ -14,10 +14,14 @@ limitations under the License.
 package driver
 
 import (
+	"encoding/base64"
 	"reflect"
 	"testing"
 
+	"github.com/gogo/protobuf/proto"
+
 	rspb "k8s.io/helm/pkg/proto/hapi/release"
+	"k8s.io/kubernetes/pkg/api"
 )
 
 func TestConfigMapName(t *testing.T) {
@@ -46,6 +50,37 @@ func TestConfigMapGet(t *testing.T) {
 	}
 }
 
+func TestUNcompressedConfigMapGet(t *testing.T) {
+	vers := int32(1)
+	name := "smug-pigeon"
+	key := testKey(name, vers)
+	rel := releaseStub(name, vers, rspb.Status_DEPLOYED)
+
+	// Create a test fixture which contains an uncompressed release
+	cfgmap, err := newConfigMapsObject(key, rel, nil)
+	if err != nil {
+		t.Fatalf("Failed to create configmap: %s", err)
+	}
+	b, err := proto.Marshal(rel)
+	if err != nil {
+		t.Fatalf("Failed to marshal release: %s", err)
+	}
+	cfgmap.Data["release"] = base64.StdEncoding.EncodeToString(b)
+	var mock MockConfigMapsInterface
+	mock.objects = map[string]*api.ConfigMap{key: cfgmap}
+	cfgmaps := NewConfigMaps(&mock)
+
+	// get release with key
+	got, err := cfgmaps.Get(key)
+	if err != nil {
+		t.Fatalf("Failed to get release: %s", err)
+	}
+	// compare fetched release with original
+	if !reflect.DeepEqual(rel, got) {
+		t.Errorf("Expected {%q}, got {%q}", rel, got)
+	}
+}
+
 func TestConfigMapList(t *testing.T) {
 	cfgmaps := newTestFixtureCfgMaps(t, []*rspb.Release{
 		releaseStub("key-1", 1, rspb.Status_DELETED),
