add memory limit plugin for memory advisor

pkg/agent/sysadvisor/plugin/qosaware/resource/memory/advisor.go

@@ -51,6 +51,7 @@ func init() {
 	memadvisorplugin.RegisterInitializer(memadvisorplugin.MemsetBinder, memadvisorplugin.NewMemsetBinder)
 	memadvisorplugin.RegisterInitializer(memadvisorplugin.NumaMemoryBalancer, memadvisorplugin.NewMemoryBalancer)
 	memadvisorplugin.RegisterInitializer(memadvisorplugin.TransparentMemoryOffloading, memadvisorplugin.NewTransparentMemoryOffloading)
+	memadvisorplugin.RegisterInitializer(memadvisorplugin.MemoryLimit, memadvisorplugin.NewMemoryLimit)
 	memadvisorplugin.RegisterInitializer(provisioner.MemoryProvisioner, provisioner.NewMemoryProvisioner)
 }
 

pkg/agent/sysadvisor/plugin/qosaware/resource/memory/plugin/memory_limit.go

@@ -0,0 +1,125 @@
+/*
+Copyright 2022 The Katalyst Authors.
+
+Licensed 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 plugin
+
+import (
+	"strconv"
+
+	"go.uber.org/atomic"
+
+	"github.com/kubewharf/katalyst-api/pkg/consts"
+	"github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/memory/dynamicpolicy/memoryadvisor"
+	"github.com/kubewharf/katalyst-core/pkg/agent/sysadvisor/metacache"
+	"github.com/kubewharf/katalyst-core/pkg/agent/sysadvisor/types"
+	"github.com/kubewharf/katalyst-core/pkg/config"
+	"github.com/kubewharf/katalyst-core/pkg/metaserver"
+	"github.com/kubewharf/katalyst-core/pkg/metrics"
+	"github.com/kubewharf/katalyst-core/pkg/util/cgroup/common"
+	"github.com/kubewharf/katalyst-core/pkg/util/general"
+	"github.com/kubewharf/katalyst-core/pkg/util/native"
+)
+
+const (
+	MemoryLimit = "memory-limit"
+
+	memoryLimitStatusSucceeded = "succeeded"
+	memoryLimitStatusFailed    = "failed"
+)
+
+type memoryLimit struct {
+	metaReader            metacache.MetaReader
+	metaServer            *metaserver.MetaServer
+	emitter               metrics.MetricEmitter
+	reconcileStatus       *atomic.String
+	reclaimMemoryLimitMap map[string]map[memoryadvisor.MemoryControlKnobName]string // cgroupPath -> cgroupFile(memory.limit_in_bytes) -> target value
+	getCgroupPathFunc     func(string, string) (string, error)
+}
+
+func NewMemoryLimit(conf *config.Configuration, extraConfig interface{}, metaReader metacache.MetaReader, metaServer *metaserver.MetaServer, emitter metrics.MetricEmitter) MemoryAdvisorPlugin {
+	return &memoryLimit{
+		metaReader:        metaReader,
+		metaServer:        metaServer,
+		emitter:           emitter,
+		reconcileStatus:   atomic.NewString(memoryLimitStatusFailed),
+		getCgroupPathFunc: common.GetContainerRelativeCgroupPath,
+	}
+}
+
+func (ml *memoryLimit) Reconcile(status *types.MemoryPressureStatus) error {
+	ml.reconcileStatus.Store(memoryLimitStatusFailed)
+
+	mp := make(map[string]map[memoryadvisor.MemoryControlKnobName]string)
+	ml.metaReader.RangeContainer(func(podUID string, containerName string, containerInfo *types.ContainerInfo) bool {
+		if containerInfo.QoSLevel != consts.PodAnnotationQoSLevelReclaimedCores {
+			return true
+		}
+
+		containerID, err := ml.metaServer.GetContainerID(podUID, containerName)
+		if err != nil {
+			general.ErrorS(err, "failed to get container id", "podUID", podUID, "containerName", containerName)
+			return false
+		}
+		cgroupPath, err := ml.getCgroupPathFunc(podUID, containerID)
+		if err != nil {
+			general.ErrorS(err, "failed to get container cgroup path", "podUID", podUID, "containerName", containerName)
+			return false
+		}
+		container, err := ml.metaServer.GetContainerSpec(podUID, containerName)
+		if err != nil {
+			general.ErrorS(err, "failed to get container spec", "podUID", podUID, "containerName", containerName)
+			return false
+		}
+
+		var limit int64 = -1
+		if q := native.DefaultMemoryQuantityGetter(container.Resources.Limits); q.Value() > 0 {
+			limit = q.Value()
+		}
+		mp[cgroupPath] = map[memoryadvisor.MemoryControlKnobName]string{
+			memoryadvisor.ControlKnobKeyMemoryLimitInBytes: strconv.FormatInt(limit, 10),
+		}
+		return true
+	})
+	ml.reclaimMemoryLimitMap = mp
+	general.InfoS("reclaimed pod memory limit", "target", mp)
+
+	ml.reconcileStatus.Store(memoryLimitStatusSucceeded)
+	return nil
+}
+
+func (ml *memoryLimit) GetAdvices() types.InternalMemoryCalculationResult {
+	if ml.reconcileStatus.Load() == memoryLimitStatusFailed {
+		general.Errorf("failed to get last reconcile result")
+		return types.InternalMemoryCalculationResult{}
+	}
+
+	var extraEntries []types.ExtraMemoryAdvices
+	for cgroupPath, value := range ml.reclaimMemoryLimitMap {
+		if value == nil || value[memoryadvisor.ControlKnobKeyMemoryLimitInBytes] == "" {
+			continue
+		}
+		extraEntries = append(extraEntries, types.ExtraMemoryAdvices{
+			CgroupPath: cgroupPath,
+			Values: map[string]string{
+				string(memoryadvisor.ControlKnobKeyMemoryLimitInBytes): value[memoryadvisor.ControlKnobKeyMemoryLimitInBytes],
+			},
+		})
+	}
+
+	return types.InternalMemoryCalculationResult{
+		ExtraEntries: extraEntries,
+	}
+}

pkg/agent/sysadvisor/plugin/qosaware/resource/memory/plugin/memory_limit_test.go

@@ -0,0 +1,179 @@
+/*
+Copyright 2022 The Katalyst Authors.
+
+Licensed 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 plugin
+
+import (
+	"io/ioutil"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+	v1 "k8s.io/api/core/v1"
+	"k8s.io/apimachinery/pkg/api/resource"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	apitypes "k8s.io/apimachinery/pkg/types"
+
+	apiconsts "github.com/kubewharf/katalyst-api/pkg/consts"
+	"github.com/kubewharf/katalyst-core/cmd/katalyst-agent/app/options"
+	"github.com/kubewharf/katalyst-core/pkg/agent/qrm-plugins/memory/dynamicpolicy/memoryadvisor"
+	"github.com/kubewharf/katalyst-core/pkg/agent/sysadvisor/metacache"
+	"github.com/kubewharf/katalyst-core/pkg/agent/sysadvisor/types"
+	"github.com/kubewharf/katalyst-core/pkg/config"
+	"github.com/kubewharf/katalyst-core/pkg/metaserver"
+	"github.com/kubewharf/katalyst-core/pkg/metaserver/agent"
+	"github.com/kubewharf/katalyst-core/pkg/metaserver/agent/metric"
+	"github.com/kubewharf/katalyst-core/pkg/metaserver/agent/pod"
+	"github.com/kubewharf/katalyst-core/pkg/metrics"
+	metricspool "github.com/kubewharf/katalyst-core/pkg/metrics/metrics-pool"
+)
+
+func generateTestMemoryAdvisorConfiguration(t *testing.T) *config.Configuration {
+	conf, err := options.NewOptions().Config()
+	require.NoError(t, err)
+	require.NotNil(t, conf)
+
+	tmpStateDir, err := ioutil.TempDir("", "sys-advisor-test")
+	require.NoError(t, err)
+	tmpMemoryAdvisorSocketDir, err := ioutil.TempDir("", "sys-advisor-test")
+	require.NoError(t, err)
+
+	conf.GenericSysAdvisorConfiguration.StateFileDirectory = tmpStateDir
+	conf.QRMAdvisorConfiguration.MemoryAdvisorSocketAbsPath = tmpMemoryAdvisorSocketDir + "-memory_advisor.sock"
+	conf.QRMAdvisorConfiguration.MemoryPluginSocketAbsPath = tmpMemoryAdvisorSocketDir + "-memory_plugin.sock"
+
+	return conf
+}
+
+func getFakeCgroupPath(podUID, containerID string) string {
+	return podUID + "/" + containerID
+}
+
+func generatePod(uid, containerName, containerID string, limits v1.ResourceList) *v1.Pod {
+	return &v1.Pod{
+		ObjectMeta: metav1.ObjectMeta{
+			UID: apitypes.UID(uid),
+		},
+		Spec: v1.PodSpec{
+			Containers: []v1.Container{
+				{
+					Name: containerName,
+					Resources: v1.ResourceRequirements{
+						Limits: limits,
+					},
+				},
+			},
+		},
+		Status: v1.PodStatus{
+			ContainerStatuses: []v1.ContainerStatus{
+				{
+					Name:        containerName,
+					ContainerID: containerID,
+				},
+			},
+		},
+	}
+}
+
+func TestReconcile(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name             string
+		podUID           string
+		containerID      string
+		containerName    string
+		memoryUpperLimit resource.Quantity
+		isReclaimedPod   bool
+	}{
+		{
+			name:             "reclaimed_pod_has_memory_limit",
+			podUID:           "pod1",
+			containerID:      "container1",
+			containerName:    "container1",
+			memoryUpperLimit: resource.MustParse("100"),
+			isReclaimedPod:   true,
+		},
+		{
+			name:             "shared_pod_has_memory_limit",
+			podUID:           "pod2",
+			containerID:      "container2",
+			containerName:    "container2",
+			memoryUpperLimit: resource.MustParse("100"),
+			isReclaimedPod:   false,
+		},
+	}
+
+	for _, tc := range tests {
+		tc := tc
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			conf := generateTestMemoryAdvisorConfiguration(t)
+			metricsFetcher := metric.NewFakeMetricsFetcher(metrics.DummyMetrics{})
+			metaCache, err := metacache.NewMetaCacheImp(conf, metricspool.DummyMetricsEmitterPool{}, metricsFetcher)
+			assert.NoError(t, err)
+
+			fetcher := pod.PodFetcherStub{}
+			metaServer := &metaserver.MetaServer{
+				MetaAgent: &agent.MetaAgent{
+					PodFetcher: &fetcher,
+				},
+			}
+
+			ml := NewMemoryLimit(nil, nil, metaCache, metaServer, nil)
+			ml.(*memoryLimit).getCgroupPathFunc = func(podUID string, containerID string) (string, error) {
+				return getFakeCgroupPath(podUID, containerID), nil
+			}
+			containerInfo := types.ContainerInfo{
+				PodUID:        tc.podUID,
+				ContainerName: tc.containerName,
+			}
+			if tc.isReclaimedPod {
+				containerInfo.QoSLevel = apiconsts.PodAnnotationQoSLevelReclaimedCores
+			} else {
+				containerInfo.QoSLevel = apiconsts.PodAnnotationQoSLevelSharedCores
+			}
+
+			err = metaCache.SetContainerInfo(tc.podUID, tc.containerName, &containerInfo)
+			assert.NoError(t, err)
+
+			testPod := generatePod(tc.podUID, tc.containerName, tc.containerID, v1.ResourceList{
+				v1.ResourceMemory: tc.memoryUpperLimit,
+			})
+			fetcher.PodList = append(fetcher.PodList, testPod)
+
+			err = ml.Reconcile(nil)
+			assert.NoError(t, err)
+
+			advices := ml.GetAdvices()
+			found := false
+			for _, entry := range advices.ExtraEntries {
+				if strings.Contains(entry.CgroupPath, tc.podUID) && strings.Contains(entry.CgroupPath, tc.containerID) && tc.isReclaimedPod {
+					assert.Equal(t, entry.Values[string(memoryadvisor.ControlKnobKeyMemoryLimitInBytes)], tc.memoryUpperLimit.String())
+					found = true
+				}
+			}
+			if !found && tc.isReclaimedPod {
+				t.Errorf("Expected to find memory limit for reclaimed pod %s", tc.podUID)
+			}
+			if found && !tc.isReclaimedPod {
+				t.Errorf("Unexpected to find memory limit for shared pod %s", tc.podUID)
+			}
+		})
+	}
+}