From 5ed586dafb0546a488dded23e379e4b148739e87 Mon Sep 17 00:00:00 2001
From: Lucas Jacques <contact@lucasjacques.com>
Date: Thu, 5 Dec 2024 20:15:53 +0100
Subject: [PATCH] feat: add cpu max burst support

Signed-off-by: Lucas Jacques <contact@lucasjacques.com>
---
 cgroup2/cpu.go        | 20 ++++++++++++++++----
 cgroup2/cpuv2_test.go | 11 +++++++----
 cgroup2/utils.go      |  3 +++
 3 files changed, 26 insertions(+), 8 deletions(-)

diff --git a/cgroup2/cpu.go b/cgroup2/cpu.go
index dcb253db..b951928a 100644
--- a/cgroup2/cpu.go
+++ b/cgroup2/cpu.go
@@ -32,11 +32,16 @@ func NewCPUMax(quota *int64, period *uint64) CPUMax {
 	return CPUMax(strings.Join([]string{max, strconv.FormatUint(*period, 10)}, " "))
 }
 
+func NewCPUMaxBurst(burst uint64) string {
+	return strconv.FormatUint(burst, 10)
+}
+
 type CPU struct {
-	Weight *uint64
-	Max    CPUMax
-	Cpus   string
-	Mems   string
+	Weight   *uint64
+	Max      CPUMax
+	MaxBurst string
+	Cpus     string
+	Mems     string
 }
 
 func (c CPUMax) extractQuotaAndPeriod() (int64, uint64) {
@@ -79,5 +84,12 @@ func (r *CPU) Values() (o []Value) {
 			value:    r.Mems,
 		})
 	}
+	if r.MaxBurst != "" {
+		o = append(o, Value{
+			filename: "cpu.max.burst",
+			value:    r.MaxBurst,
+		})
+	}
+
 	return o
 }
diff --git a/cgroup2/cpuv2_test.go b/cgroup2/cpuv2_test.go
index 6751d23d..669283d0 100644
--- a/cgroup2/cpuv2_test.go
+++ b/cgroup2/cpuv2_test.go
@@ -32,6 +32,7 @@ func TestCgroupv2CpuStats(t *testing.T) {
 	group := "/cpu-test-cg"
 	groupPath := fmt.Sprintf("%s-%d", group, os.Getpid())
 	var (
+		burst  uint64 = 1000
 		quota  int64  = 10000
 		period uint64 = 8000
 		weight uint64 = 100
@@ -39,10 +40,11 @@ func TestCgroupv2CpuStats(t *testing.T) {
 
 	c, err := NewManager(defaultCgroup2Path, groupPath, &Resources{
 		CPU: &CPU{
-			Weight: &weight,
-			Max:    NewCPUMax(&quota, &period),
-			Cpus:   "0",
-			Mems:   "0",
+			Weight:   &weight,
+			Max:      NewCPUMax(&quota, &period),
+			Cpus:     "0",
+			Mems:     "0",
+			MaxBurst: NewCPUMaxBurst(burst),
 		},
 	})
 	require.NoError(t, err, "failed to init new cgroup manager")
@@ -54,6 +56,7 @@ func TestCgroupv2CpuStats(t *testing.T) {
 	checkFileContent(t, c.path, "cpu.max", "10000 8000")
 	checkFileContent(t, c.path, "cpuset.cpus", "0")
 	checkFileContent(t, c.path, "cpuset.mems", "0")
+	checkFileContent(t, c.path, "cpu.max.burst", strconv.FormatUint(burst, 10))
 }
 
 func TestSystemdCgroupCpuController(t *testing.T) {
diff --git a/cgroup2/utils.go b/cgroup2/utils.go
index 0974ce6a..356bdc54 100644
--- a/cgroup2/utils.go
+++ b/cgroup2/utils.go
@@ -172,6 +172,9 @@ func ToResources(spec *specs.LinuxResources) *Resources {
 		if period := cpu.Period; period != nil {
 			resources.CPU.Max = NewCPUMax(cpu.Quota, period)
 		}
+		if burst := cpu.Burst; burst != nil {
+			resources.CPU.MaxBurst = NewCPUMaxBurst(*burst)
+		}
 	}
 	if mem := spec.Memory; mem != nil {
 		resources.Memory = &Memory{}