From 0a5471589a843025e73b30196bc1324617093f65 Mon Sep 17 00:00:00 2001
From: Krisztian Litkey <krisztian.litkey@intel.com>
Date: Mon, 17 Jun 2024 20:28:04 +0300
Subject: [PATCH] balloons: initial libmem conversion.

Plug in libmem-based memory allocation (and accounting).

Signed-off-by: Krisztian Litkey <krisztian.litkey@intel.com>
---
 .../balloons/policy/balloons-policy.go        | 91 ++++++++++++++++++-
 1 file changed, 90 insertions(+), 1 deletion(-)

diff --git a/cmd/plugins/balloons/policy/balloons-policy.go b/cmd/plugins/balloons/policy/balloons-policy.go
index 8a8068f62..a830905a5 100644
--- a/cmd/plugins/balloons/policy/balloons-policy.go
+++ b/cmd/plugins/balloons/policy/balloons-policy.go
@@ -26,6 +26,7 @@ import (
 	"github.com/containers/nri-plugins/pkg/resmgr/cache"
 	cpucontrol "github.com/containers/nri-plugins/pkg/resmgr/control/cpu"
 	"github.com/containers/nri-plugins/pkg/resmgr/events"
+	libmem "github.com/containers/nri-plugins/pkg/resmgr/lib/memory"
 	policy "github.com/containers/nri-plugins/pkg/resmgr/policy"
 	"github.com/containers/nri-plugins/pkg/utils"
 	"github.com/containers/nri-plugins/pkg/utils/cpuset"
@@ -71,6 +72,7 @@ type balloons struct {
 	balloons           []*Balloon  // balloon instances: reserved, default and user-defined
 
 	cpuAllocator cpuallocator.CPUAllocator // CPU allocator used by the policy
+	memAllocator *libmem.Allocator         // memory allocator used by the policy
 }
 
 // Balloon contains attributes of a balloon instance
@@ -162,6 +164,12 @@ func (p *balloons) Setup(policyOptions *policy.BackendOptions) error {
 	p.cch = policyOptions.Cache
 	p.cpuAllocator = cpuallocator.NewCPUAllocator(policyOptions.System)
 
+	malloc, err := libmem.NewAllocator(libmem.WithSystemNodes(policyOptions.System))
+	if err != nil {
+		return balloonsError("failed to create memory allocator: %w", err)
+	}
+	p.memAllocator = malloc
+
 	log.Info("setting up %s policy...", PolicyName)
 	if p.cpuTree, err = NewCpuTreeFromSystem(); err != nil {
 		log.Errorf("creating CPU topology tree failed: %s", err)
@@ -1421,6 +1429,9 @@ func (p *balloons) assignContainer(c cache.Container, bln *Balloon) {
 
 // dismissContainer removes a container from a balloon
 func (p *balloons) dismissContainer(c cache.Container, bln *Balloon) {
+	if err := p.memAllocator.Release(c.GetID()); err != nil {
+		log.Error("dismissContainer: failed to release memory for %s: %v", c.PrettyName(), err)
+	}
 	podID := c.GetPodID()
 	bln.PodIDs[podID] = removeString(bln.PodIDs[podID], c.GetID())
 	if len(bln.PodIDs[podID]) == 0 {
@@ -1442,13 +1453,91 @@ func (p *balloons) pinCpuMem(c cache.Container, cpus cpuset.CPUSet, mems idset.I
 	if p.bpoptions.PinMemory == nil || *p.bpoptions.PinMemory {
 		if c.PreserveMemoryResources() {
 			log.Debug("  - preserving %s pinning to memory %q", c.PrettyName, c.GetCpusetMems())
+			preserveMems, err := parseIDSet(c.GetCpusetMems())
+			if err != nil {
+				log.Error("failed to parse CpusetMems: %v", err)
+			} else {
+				zone := p.allocMem(c, preserveMems, true)
+				log.Debug("  - allocated preserved memory %s", c.PrettyName, zone)
+				c.SetCpusetMems(zone.MemsetString())
+			}
 		} else {
 			log.Debug("  - pinning %s to memory %s", c.PrettyName(), mems)
-			c.SetCpusetMems(mems.String())
+			zone := p.allocMem(c, mems, false)
+			log.Debug("  - allocated memory %s", c.PrettyName, zone)
+			c.SetCpusetMems(zone.MemsetString())
 		}
 	}
 }
 
+func (p *balloons) allocMem(c cache.Container, mems idset.IDSet, preserve bool) libmem.NodeMask {
+	var (
+		amount  = getMemoryLimit(c)
+		nodes   = libmem.NewNodeMask(mems.Members()...)
+		req     *libmem.Request
+		zone    libmem.NodeMask
+		updates map[string]libmem.NodeMask
+		err     error
+	)
+
+	if _, ok := p.memAllocator.AssignedZone(c.GetID()); !ok {
+		if c.PreserveMemoryResources() {
+			req = libmem.PreservedContainer(
+				c.GetID(),
+				c.PrettyName(),
+				amount,
+				nodes,
+			)
+		} else {
+			req = libmem.Container(
+				c.GetID(),
+				c.PrettyName(),
+				string(c.GetQOSClass()),
+				amount,
+				nodes,
+			)
+		}
+		zone, updates, err = p.memAllocator.Allocate(req)
+	} else {
+		zone, updates, err = p.memAllocator.Realloc(c.GetID(), nodes, 0)
+	}
+
+	if err != nil {
+		log.Error("allocMem: falling back to %s, failed to allocate memory for %s: %v",
+			nodes, c.PrettyName(), err)
+		return nodes
+	}
+
+	for oID, oz := range updates {
+		if oc, ok := p.cch.LookupContainer(oID); ok {
+			oc.SetCpusetMems(oz.MemsetString())
+		}
+	}
+
+	return zone
+}
+
+func parseIDSet(mems string) (idset.IDSet, error) {
+	cset, err := cpuset.Parse(mems)
+	if err != nil {
+		return idset.NewIDSet(), err
+	}
+	return idset.NewIDSet(cset.List()...), nil
+}
+
+func getMemoryLimit(c cache.Container) int64 {
+	res, ok := c.GetResourceUpdates()
+	if !ok {
+		res = c.GetResourceRequirements()
+	}
+
+	if limit, ok := res.Limits[corev1.ResourceMemory]; ok {
+		return limit.Value()
+	}
+
+	return 0
+}
+
 // balloonsError formats an error from this policy.
 func balloonsError(format string, args ...interface{}) error {
 	return fmt.Errorf(PolicyName+": "+format, args...)