Skip to content

Commit

Permalink
[YUNIKORN-2287] Decompress function doesn't need to decode base64
Browse files Browse the repository at this point in the history
Signed-off-by: PoAn Yang <payang@apache.org>
  • Loading branch information
FrankYang0529 committed Dec 26, 2023
1 parent 24c699f commit 60037a6
Show file tree
Hide file tree
Showing 10 changed files with 289 additions and 18 deletions.
5 changes: 1 addition & 4 deletions pkg/common/utils/utils_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ package utils
import (
"bytes"
"compress/gzip"
"encoding/base64"
"fmt"
"strings"
"testing"
Expand Down Expand Up @@ -1059,11 +1058,9 @@ func TestGzipCompressedConfigMap(t *testing.T) {
if err := gzWriter.Close(); err != nil {
t.Fatal("expected nil, got error")
}
encodedConfigString := make([]byte, base64.StdEncoding.EncodedLen(len(b.Bytes())))
base64.StdEncoding.Encode(encodedConfigString, b.Bytes())
confMap := conf.FlattenConfigMaps([]*v1.ConfigMap{
{Data: map[string]string{}},
{Data: map[string]string{conf.CMSvcClusterID: "new"}, BinaryData: map[string][]byte{"queues.yaml.gz": encodedConfigString}},
{Data: map[string]string{conf.CMSvcClusterID: "new"}, BinaryData: map[string][]byte{"queues.yaml.gz": b.Bytes()}},
})
config := GetCoreSchedulerConfigFromConfigMap(confMap)
assert.Equal(t, configs.DefaultSchedulerConfig, config)
Expand Down
10 changes: 1 addition & 9 deletions pkg/conf/schedulerconf.go
Original file line number Diff line number Diff line change
Expand Up @@ -21,7 +21,6 @@ package conf
import (
"bytes"
"compress/gzip"
"encoding/base64"
"encoding/json"
"errors"
"flag"
Expand Down Expand Up @@ -452,17 +451,10 @@ func DumpConfiguration() {

func Decompress(key string, value []byte) (string, string) {
var uncompressedData string
decodedValue := make([]byte, base64.StdEncoding.DecodedLen(len(value)))
n, err := base64.StdEncoding.Decode(decodedValue, value)
if err != nil {
log.Log(log.ShimConfig).Error("failed to decode schedulerConfig entry", zap.Error(err))
return "", ""
}
decodedValue = decodedValue[:n]
splitKey := strings.Split(key, ".")
compressionAlgo := splitKey[len(splitKey)-1]
if strings.EqualFold(compressionAlgo, constants.GzipSuffix) {
reader := bytes.NewReader(decodedValue)
reader := bytes.NewReader(value)
gzReader, err := gzip.NewReader(reader)
if err != nil {
log.Log(log.ShimConfig).Error("failed to decompress decoded schedulerConfig entry", zap.Error(err))
Expand Down
6 changes: 2 additions & 4 deletions pkg/conf/schedulerconf_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -87,17 +87,15 @@ func TestDecompress(t *testing.T) {
assert.NilError(t, err, "expected nil, got error")
t.Fatal("expected nil, got error")
}
encodedConfigString := make([]byte, base64.StdEncoding.EncodedLen(len(b.Bytes())))
base64.StdEncoding.Encode(encodedConfigString, b.Bytes())
key, decodedConfigString := Decompress("queues.yaml."+constants.GzipSuffix, encodedConfigString)
key, decodedConfigString := Decompress("queues.yaml."+constants.GzipSuffix, b.Bytes())
assert.Equal(t, "queues.yaml", key)
assert.Equal(t, configs.DefaultSchedulerConfig, decodedConfigString)
}

func TestDecompressUnknownKey(t *testing.T) {
encodedConfigString := make([]byte, base64.StdEncoding.EncodedLen(len([]byte(configs.DefaultSchedulerConfig))))
base64.StdEncoding.Encode(encodedConfigString, []byte(configs.DefaultSchedulerConfig))
key, decodedConfigString := Decompress("queues.yaml.bin", encodedConfigString)
key, decodedConfigString := Decompress("queues.yaml.bin", []byte(configs.DefaultSchedulerConfig))
assert.Equal(t, "queues.yaml", key)
assert.Assert(t, len(decodedConfigString) == 0, "expected decodedConfigString to be nil")
}
Expand Down
3 changes: 2 additions & 1 deletion test/e2e/bin_packing/bin_packing_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -198,8 +198,9 @@ var _ = Describe("", func() {
tests.LogTestClusterInfoWrapper(testDescription.FailureMessage(), []string{ns})
tests.LogYunikornContainer(testDescription.FailureMessage())
}

By("Tear down namespace: " + ns)
err := kClient.DeleteNamespace(ns)
err := kClient.TearDownNamespace(ns)
Ω(err).NotTo(HaveOccurred())
})
})
94 changes: 94 additions & 0 deletions test/e2e/configmap/configmap_suite_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package configmap

import (
"path/filepath"
"testing"

"github.com/onsi/ginkgo/v2"
"github.com/onsi/ginkgo/v2/reporters"
"github.com/onsi/gomega"

"github.com/apache/yunikorn-k8shim/test/e2e/framework/configmanager"
"github.com/apache/yunikorn-k8shim/test/e2e/framework/helpers/common"
"github.com/apache/yunikorn-k8shim/test/e2e/framework/helpers/k8s"
"github.com/apache/yunikorn-k8shim/test/e2e/framework/helpers/yunikorn"
)

func init() {
configmanager.YuniKornTestConfig.ParseFlags()
}

func TestConfigMap(t *testing.T) {
ginkgo.ReportAfterSuite("TestConfigMap", func(report ginkgo.Report) {
err := common.CreateJUnitReportDir()
Ω(err).NotTo(gomega.HaveOccurred())
err = reporters.GenerateJUnitReportWithConfig(
report,
filepath.Join(configmanager.YuniKornTestConfig.LogDir, "TEST-configmap_junit.xml"),
reporters.JunitReportConfig{OmitSpecLabels: true},
)
Ω(err).NotTo(HaveOccurred())
})
gomega.RegisterFailHandler(ginkgo.Fail)
ginkgo.RunSpecs(t, "TestConfigMap", ginkgo.Label("TestConfigMap"))
}

var (
kClient k8s.KubeCtl
restClient yunikorn.RClient
)

var _ = BeforeSuite(func() {

kClient = k8s.KubeCtl{}
Ω(kClient.SetClient()).To(BeNil())

restClient = yunikorn.RClient{}
Ω(restClient).NotTo(BeNil())

yunikorn.EnsureYuniKornConfigsPresent()

// Restart yunikorn and port-forward
// Required to change node sort policy.
ginkgo.By("Restart the scheduler pod")
yunikorn.RestartYunikorn(&kClient)

ginkgo.By("Port-forward scheduler pod after restart")
yunikorn.RestorePortForwarding(&kClient)
})

var _ = AfterSuite(func() {})

var Describe = ginkgo.Describe
var It = ginkgo.It
var By = ginkgo.By
var BeforeEach = ginkgo.BeforeEach
var AfterEach = ginkgo.AfterEach
var BeforeSuite = ginkgo.BeforeSuite
var AfterSuite = ginkgo.AfterSuite

// Declarations for Gomega Matchers
var Equal = gomega.Equal
var Ω = gomega.Expect
var BeNil = gomega.BeNil
var BeNumerically = gomega.BeNumerically
var HaveOccurred = gomega.HaveOccurred
var BeTrue = gomega.BeTrue
144 changes: 144 additions & 0 deletions test/e2e/configmap/configmap_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,144 @@
/*
Licensed to the Apache Software Foundation (ASF) under one
or more contributor license agreements. See the NOTICE file
distributed with this work for additional information
regarding copyright ownership. The ASF licenses this file
to you under the Apache License, Version 2.0 (the
"License"); you may not use this file except in compliance
with the License. You may obtain a copy of the License at
http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package configmap

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

v1 "k8s.io/api/core/v1"

"github.com/apache/yunikorn-core/pkg/common/configs"
"github.com/apache/yunikorn-k8shim/test/e2e/framework/configmanager"
"github.com/apache/yunikorn-k8shim/test/e2e/framework/helpers/common"
"github.com/apache/yunikorn-k8shim/test/e2e/framework/helpers/k8s"
"github.com/apache/yunikorn-k8shim/test/e2e/framework/helpers/yunikorn"
)

var oldConfigMap = new(v1.ConfigMap)

var _ = Describe("ConfigMap", func() {
BeforeEach(func() {
By("Get previous config")
var err error
oldConfigMap, err = kClient.GetConfigMaps(configmanager.YuniKornTestConfig.YkNamespace,
configmanager.DefaultYuniKornConfigMap)
Ω(err).NotTo(HaveOccurred())
Ω(oldConfigMap).NotTo(BeNil())
})

It("Verify_Compressed_ConfigMap", func() {
schedulerConfig := &configs.SchedulerConfig{
Partitions: []configs.PartitionConfig{{
Name: "default",
Queues: []configs.QueueConfig{{
Name: "root",
SubmitACL: "*",
Parent: true,
Queues: []configs.QueueConfig{{
Name: "queue1",
SubmitACL: "*",
}},
}},
}},
}

// Wait for 1 second to set a new timestamp. If we don't wait for it, we may get a same timestamp.
time.Sleep(1 * time.Second)
ts, err := common.SetQueueTimestamp(schedulerConfig, "default", "root")
Ω(err).NotTo(HaveOccurred())

schedulerConfigStr, err := common.ToYAML(schedulerConfig)
Ω(err).NotTo(HaveOccurred())

// use gzip to compress
var b bytes.Buffer
gzWriter := gzip.NewWriter(&b)
_, err = gzWriter.Write([]byte(schedulerConfigStr))
Ω(err).NotTo(HaveOccurred())
err = gzWriter.Close()
Ω(err).NotTo(HaveOccurred())

var configMap *v1.ConfigMap
configMap, err = kClient.GetConfigMaps(configmanager.YuniKornTestConfig.YkNamespace,
configmanager.DefaultYuniKornConfigMap)
Ω(err).NotTo(HaveOccurred())
Ω(configMap).NotTo(BeNil())

configMap.Data = map[string]string{}
if configMap.BinaryData == nil {
configMap.BinaryData = map[string][]byte{}
}
configMap.BinaryData[configmanager.DefaultPolicyGroup+".gz"] = b.Bytes()

By("Updating the config map with BinaryData")
_, err = kClient.UpdateConfigMap(configMap, configmanager.YuniKornTestConfig.YkNamespace)
Ω(err).NotTo(HaveOccurred())

err = yunikorn.WaitForQueueTS("root", ts, 30*time.Second)
Ω(err).NotTo(HaveOccurred())

checkSchedulerConfig(schedulerConfig)
})

It("Verify_Compressed_ConfigMap_From_File", func() {
configMap, err := k8s.GetConfigMapObj("../testdata/compressed_yunikorn-configs.yaml")
Ω(err).NotTo(HaveOccurred())
Ω(configMap).NotTo(BeNil())

By("Updating the config map with BinaryData")
configMap.Namespace = configmanager.YuniKornTestConfig.YkNamespace
_, err = kClient.UpdateConfigMap(configMap, configmanager.YuniKornTestConfig.YkNamespace)
Ω(err).NotTo(HaveOccurred())

queuesGz := configMap.BinaryData[configmanager.DefaultPolicyGroup+".gz"]
Ω(len(queuesGz)).NotTo(Equal(0))
reader := bytes.NewReader(queuesGz)
gzReader, err := gzip.NewReader(reader)
Ω(err).NotTo(HaveOccurred())
err = gzReader.Close()
Ω(err).NotTo(HaveOccurred())
decompressedBytes, err := io.ReadAll(gzReader)
Ω(err).NotTo(HaveOccurred())

schedulerConfig, err := configs.LoadSchedulerConfigFromByteArray(decompressedBytes)
Ω(err).NotTo(HaveOccurred())
Ω(len(schedulerConfig.Partitions)).To(Equal(1))
Ω(len(schedulerConfig.Partitions[0].Queues)).To(Equal(1))

ts := schedulerConfig.Partitions[0].Queues[0].Properties["timestamp"]
err = yunikorn.WaitForQueueTS("root", ts, 30*time.Second)
Ω(err).NotTo(HaveOccurred())

checkSchedulerConfig(schedulerConfig)
})

AfterEach(func() {
yunikorn.RestoreConfigMapWrapper(oldConfigMap, "")
})
})

func checkSchedulerConfig(schedulerConfig *configs.SchedulerConfig) {
configDAOInfo, err := restClient.GetConfig()
Ω(err).NotTo(HaveOccurred())
Ω(configDAOInfo).NotTo(BeNil())
Ω(configDAOInfo.Partitions).To(Equal(schedulerConfig.Partitions))
}
1 change: 1 addition & 0 deletions test/e2e/framework/configmanager/constants.go
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ const (
QueuesPath = "ws/v1/partition/%s/queues"
AppsPath = "ws/v1/partition/%s/queue/%s/applications"
AppPath = "ws/v1/partition/%s/queue/%s/application/%s"
ConfigPath = "ws/v1/config"
ClustersPath = "ws/v1/clusters"
NodesPath = "ws/v1/partition/%s/nodes"
UserUsagePath = "ws/v1/partition/%s/usage/user/%s"
Expand Down
10 changes: 10 additions & 0 deletions test/e2e/framework/helpers/yunikorn/rest_api_utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -97,6 +97,16 @@ func (c *RClient) GetQueues(partition string) (*dao.PartitionQueueDAOInfo, error
return queues, err
}

func (c *RClient) GetConfig() (*dao.ConfigDAOInfo, error) {
req, err := c.newRequest("GET", configmanager.ConfigPath, nil)
if err != nil {
return nil, err
}
var config *dao.ConfigDAOInfo
_, err = c.do(req, &config)
return config, err
}

func (c *RClient) GetHealthCheck() (dao.SchedulerHealthDAOInfo, error) {
req, err := c.newRequest("GET", configmanager.HealthCheckPath, nil)
if err != nil {
Expand Down
1 change: 1 addition & 0 deletions test/e2e/framework/helpers/yunikorn/wrappers.go
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,7 @@ func RestoreConfigMapWrapper(oldConfigMap *v1.ConfigMap, annotation string) {
ts, tsErr := common.SetQueueTimestamp(oldSC, "default", "root")
Ω(tsErr).NotTo(HaveOccurred())
c.Data = oldConfigMap.Data
c.BinaryData = oldConfigMap.BinaryData
c.Data[configmanager.DefaultPolicyGroup], err = common.ToYAML(oldSC)
Ω(err).NotTo(HaveOccurred())

Expand Down
33 changes: 33 additions & 0 deletions test/e2e/testdata/compressed_yunikorn-configs.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
#Licensed to the Apache Software Foundation (ASF) under one
#or more contributor license agreements. See the NOTICE file
#distributed with this work for additional information
#regarding copyright ownership. The ASF licenses this file
#to you under the Apache License, Version 2.0 (the
#"License"); you may not use this file except in compliance
#with the License. You may obtain a copy of the License at
#
#http://www.apache.org/licenses/LICENSE-2.0
#
#Unless required by applicable law or agreed to in writing, software
#distributed under the License is distributed on an "AS IS" BASIS,
#WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#See the License for the specific language governing permissions and
#limitations under the License.
apiVersion: v1
kind: ConfigMap
metadata:
name: yunikorn-configs
binaryData:
# echo "
# partitions:
# - name: default
# queues:
# - name: root
# submitacl: '*'
# parent: true
# properties:
# timestamp: '1257894000'
# queues:
# - name: parent
# submitacl: '*'" | gzip | base64
queues.yaml.gz: H4sIANGBiWUAA2WO2wrCMBBE3/sV81YQhCiKmr9ZdQuB5mKy+/+maptK5/Hscma6RFmcuBiK7YA9Anm2ePJAOkolwEtZ+XOdMn/kGOWHgKJ374Qeo0W/6xdc3RzEQrJygzkmrp1NOUWc5yLkUzUcjufL9XYyxjTV/4j1kG/J6rCZ8wavmg4s5AAAAA==

0 comments on commit 60037a6

Please sign in to comment.