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 Jan 12, 2024
1 parent ce8ac51 commit 5a11535
Show file tree
Hide file tree
Showing 9 changed files with 284 additions and 17 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
88 changes: 88 additions & 0 deletions test/e2e/configmap/configmap_suite_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,88 @@
/*
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()

By("Port-forward the scheduler pod")
err := kClient.PortForwardYkSchedulerPod()
Ω(err).NotTo(HaveOccurred())
})

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 HaveOccurred = gomega.HaveOccurred
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 @@ -39,6 +39,7 @@ const (
AppsPath = "ws/v1/partition/%s/queue/%s/applications"
AppPath = "ws/v1/partition/%s/queue/%s/application/%s"
CompletedAppsPath = "ws/v1/partition/%s/applications/completed"
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 @@ -110,6 +110,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
2 changes: 2 additions & 0 deletions test/e2e/framework/helpers/yunikorn/wrappers.go
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,7 @@ func UpdateCustomConfigMapWrapperWithMap(oldConfigMap *v1.ConfigMap, schedPolicy
for k, v := range customMap {
c.Data[k] = v
}
c.BinaryData = nil
var d, err3 = k.UpdateConfigMap(c, configmanager.YuniKornTestConfig.YkNamespace)
Ω(err3).NotTo(HaveOccurred())
Ω(d).NotTo(BeNil())
Expand All @@ -134,6 +135,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
35 changes: 35 additions & 0 deletions test/e2e/testdata/compressed_yunikorn-configs.yaml
Original file line number Diff line number Diff line change
@@ -0,0 +1,35 @@
# 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:
# Following command is used to generate queues.yaml which is compressed by gzip and encoded by base64.
# The result is like the value in queues.yaml.gz. If updating the queues.yaml, please also update the queues.yaml.gz.
# 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 5a11535

Please sign in to comment.