From ba11b2a5d097933492685bad3a370d15846703b7 Mon Sep 17 00:00:00 2001
From: Alessandro Patti <ale812@yahoo.it>
Date: Wed, 18 Dec 2024 21:43:31 -0800
Subject: [PATCH] Fix parsing cpu.max file in cgroup v2 implementation

Fixes https://github.com/bazelbuild/bazel/issues/24680

Closes #24714.

PiperOrigin-RevId: 707775390
Change-Id: I1a34c0917174aca3bbcfdec8668aef2d7e68ae6e
---
 .../cgroups/controller/v2/UnifiedCpu.java     |  6 ++++-
 .../build/lib/sandbox/cgroups/CpuTest.java    | 26 +++++++++++++++++--
 2 files changed, 29 insertions(+), 3 deletions(-)

diff --git a/src/main/java/com/google/devtools/build/lib/sandbox/cgroups/controller/v2/UnifiedCpu.java b/src/main/java/com/google/devtools/build/lib/sandbox/cgroups/controller/v2/UnifiedCpu.java
index bb7343d57efb75..449c5a30289352 100644
--- a/src/main/java/com/google/devtools/build/lib/sandbox/cgroups/controller/v2/UnifiedCpu.java
+++ b/src/main/java/com/google/devtools/build/lib/sandbox/cgroups/controller/v2/UnifiedCpu.java
@@ -42,7 +42,7 @@ public Path getPath() {
   public void setCpus(double cpus) throws IOException {
     long period;
     try (Scanner scanner = new Scanner(Files.newBufferedReader(path.resolve("cpu.max")))) {
-      period = scanner.skip(".*\\s").nextInt();
+      period = scanner.skip("\\S+\\s").nextInt();
     }
     long quota = Math.round(period * cpus);
     String limit = String.format("%d %d", quota, period);
@@ -52,6 +52,10 @@ public void setCpus(double cpus) throws IOException {
   @Override
   public long getCpus() throws IOException {
     try (Scanner scanner = new Scanner(Files.newBufferedReader(path.resolve("cpu.max")))) {
+      if (!scanner.hasNextLong()) {
+        // Not a number, assume no limit and all processors available
+        return Runtime.getRuntime().availableProcessors();
+      }
       long quota = scanner.nextLong();
       long period = scanner.nextLong();
       return quota / period;
diff --git a/src/test/java/com/google/devtools/build/lib/sandbox/cgroups/CpuTest.java b/src/test/java/com/google/devtools/build/lib/sandbox/cgroups/CpuTest.java
index 9f26b16b6218b4..29242d0757cb95 100644
--- a/src/test/java/com/google/devtools/build/lib/sandbox/cgroups/CpuTest.java
+++ b/src/test/java/com/google/devtools/build/lib/sandbox/cgroups/CpuTest.java
@@ -58,10 +58,32 @@ public void setCpuLimit_v2() throws IOException {
     assertThat(Files.asCharSource(limit, UTF_8).read()).isEqualTo("500000 100000");
   }
 
+  @Test
+  public void setCpuLimitNewLine_v2() throws IOException {
+    File limit = scratch.file("cgroup/cpu/cpu.max", "-1 100000\n").getPathFile();
+    Cpu cpu = new UnifiedCpu(scratch.path("cgroup/cpu").getPathFile().toPath());
+    cpu.setCpus(5);
+    assertThat(Files.asCharSource(limit, UTF_8).read()).isEqualTo("500000 100000");
+  }
+
   @Test
   public void getCpuLimit_v2() throws IOException {
     scratch.file("cgroup/cpu/cpu.max", "6000 1000");
-    Cpu memory = new UnifiedCpu(scratch.path("cgroup/cpu").getPathFile().toPath());
-    assertThat(memory.getCpus()).isEqualTo(6);
+    Cpu cpu = new UnifiedCpu(scratch.path("cgroup/cpu").getPathFile().toPath());
+    assertThat(cpu.getCpus()).isEqualTo(6);
+  }
+
+  @Test
+  public void getCpuLimitNewLine_v2() throws IOException {
+    scratch.file("cgroup/cpu/cpu.max", "6000 1000\n");
+    Cpu cpu = new UnifiedCpu(scratch.path("cgroup/cpu").getPathFile().toPath());
+    assertThat(cpu.getCpus()).isEqualTo(6);
+  }
+
+  @Test
+  public void getCpuLimitMax_v2() throws IOException {
+    scratch.file("cgroup/cpu/cpu.max", "max 1000\n");
+    Cpu cpu = new UnifiedCpu(scratch.path("cgroup/cpu").getPathFile().toPath());
+    assertThat(cpu.getCpus()).isEqualTo(Runtime.getRuntime().availableProcessors());
   }
 }