From ee14a29cbf239303e56d9ce26f8c9ca64507a432 Mon Sep 17 00:00:00 2001
From: Krisztian Litkey <krisztian.litkey@intel.com>
Date: Sun, 3 Nov 2024 16:05:53 +0200
Subject: [PATCH] libmem: add ZoneAvailable.

Add ZoneAvailable to return the amount of available/allocatable
memory in a zone, capped by the amount of free memory in any of
the ancestors of a zone.

Signed-off-by: Krisztian Litkey <krisztian.litkey@intel.com>
---
 pkg/resmgr/lib/memory/zones.go | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/pkg/resmgr/lib/memory/zones.go b/pkg/resmgr/lib/memory/zones.go
index 67068183d..6306b9fa1 100644
--- a/pkg/resmgr/lib/memory/zones.go
+++ b/pkg/resmgr/lib/memory/zones.go
@@ -45,6 +45,11 @@ func (a *Allocator) ZoneFree(zone NodeMask) int64 {
 	return a.zoneFree(zone & a.masks.nodes.hasMemory)
 }
 
+// ZoneAvailable returns the amount of available memory in the zone.
+func (a *Allocator) ZoneAvailable(zone NodeMask) int64 {
+	return a.zoneAvailable(zone & a.masks.nodes.hasMemory)
+}
+
 // ZoneNumUsers returns the number of requests assigned to the zone.
 func (a *Allocator) ZoneNumUsers(zone NodeMask) int {
 	if z, ok := a.zones[zone]; ok {
@@ -154,6 +159,28 @@ func (a *Allocator) zoneFree(zone NodeMask) int64 {
 	return a.zoneCapacity(zone) - a.zoneUsage(zone)
 }
 
+func (a *Allocator) zoneAvailable(zone NodeMask) int64 {
+	available := a.zoneFree(zone)
+	if available <= 0 {
+		return 0
+	}
+
+	// Cap available amount to the smallest available free in any of our ancestors.
+	for _, z := range a.zones {
+		if z.nodes != zone && (z.nodes&zone) == zone {
+			if free := a.zoneFree(z.nodes); free < available {
+				available = free
+			}
+		}
+		if available <= 0 {
+			available = 0
+			break
+		}
+	}
+
+	return available
+}
+
 func (a *Allocator) zoneAssign(zone NodeMask, req *Request) {
 	z, ok := a.zones[zone]
 	if !ok {