Support reading and updating io.max cgroup in node-admin

5 files changed, 203 insertions, 23 deletions

diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/Cgroup.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/Cgroup.java
index e40e3c6c003..9079aa6fc3f 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/Cgroup.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/Cgroup.java
@@ -24,7 +24,7 @@ import java.util.logging.Logger;
 public class Cgroup {
     private static final Logger logger = Logger.getLogger(Cgroup.class.getName());
 
-    private static Map<String, Consumer<UnixPath>> cgroupDirectoryCallbacks = new HashMap<>();
+    private static final Map<String, Consumer<UnixPath>> cgroupDirectoryCallbacks = new HashMap<>();
 
     private final Path root;
     private final Path relativePath;
@@ -135,6 +135,9 @@ public class Cgroup {
     /** Returns the memory controller of this cgroup (memory.* files). */
     public MemoryController memory() { return new MemoryController(this); }
 
+    /** Returns the IO controller of this cgroup (io.* files). */
+    public IoController io() { return new IoController(this); }
+
     /**
      * Wraps {@code command} to ensure it is executed in this cgroup.
      *
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/IoController.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/IoController.java
new file mode 100644
index 00000000000..849e1ddf1a4
--- /dev/null
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/IoController.java
@@ -0,0 +1,114 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.node.admin.cgroup;
+
+import ai.vespa.validation.Validation;
+import com.yahoo.vespa.hosted.node.admin.component.TaskContext;
+import com.yahoo.vespa.hosted.node.admin.task.util.file.UnixPath;
+
+import java.util.Map;
+import java.util.Optional;
+import java.util.logging.Logger;
+import java.util.stream.Collectors;
+
+import static java.lang.Integer.parseInt;
+
+/**
+ * Represents a cgroup v2 IO controller, i.e. all io.* files.
+ *
+ * @author freva
+ */
+public class IoController {
+    private static final Logger logger = Logger.getLogger(IoController.class.getName());
+    private final Cgroup cgroup;
+
+    IoController(Cgroup cgroup) {
+        this.cgroup = cgroup;
+    }
+
+    public record Device(int major, int minor) implements Comparable<Device> {
+        public Device {
+            // https://www.halolinux.us/kernel-architecture/representation-of-major-and-minor-numbers.html
+            Validation.requireInRange(major, "device major", 0, 0xFFF);
+            Validation.requireInRange(minor, "device minor", 0, 0xFFFFF);
+        }
+
+        private String toFileContent() { return major + ":" + minor; }
+        private static Device fromString(String device) {
+            String[] parts = device.split(":");
+            return new Device(parseInt(parts[0]), parseInt(parts[1]));
+        }
+        public static Device fromDeviceNumber(int deviceNumber) {
+            return new Device(deviceNumber >>> 8, deviceNumber & 0xFF);
+        }
+
+        @Override
+        public int compareTo(Device o) {
+            return major != o.major ? Integer.compare(major, o.major) : Integer.compare(minor, o.minor);
+        }
+    }
+
+    /**
+     * Defines max allowed IO:
+     * <ul>
+     *     <ol><b>rbps</b>: Read bytes per seconds</ol>
+     *     <ol><b>riops</b>: Read IO operations per seconds</ol>
+     *     <ol><b>wbps</b>: Write bytes per seconds</ol>
+     *     <ol><b>wiops</b>: Write IO operations per seconds</ol>
+     * </ul>.
+     */
+    public record Max(Size rbps, Size wbps, Size riops, Size wiops) {
+        public static Max UNLIMITED = new Max(Size.max(), Size.max(), Size.max(), Size.max());
+
+        // Keys can be specified in any order, this is the order they are outputted in from io.max
+        // https://github.com/torvalds/linux/blob/c1a515d3c0270628df8ae5f5118ba859b85464a2/block/blk-throttle.c#L1541
+        private String toFileContent() { return "rbps=%s wbps=%s riops=%s wiops=%s".formatted(rbps, wbps, riops, wiops); }
+
+        public static Max fromString(String max) {
+            String[] parts = max.split(" ");
+            Size rbps = Size.max(), riops = Size.max(), wbps = Size.max(), wiops = Size.max();
+            for (String part : parts) {
+                if (part.isEmpty()) continue;
+                String[] kv = part.split("=");
+                if (kv.length != 2) throw new IllegalArgumentException("Invalid io.max format: " + max);
+                switch (kv[0]) {
+                    case "rbps" -> rbps = Size.from(kv[1]);
+                    case "riops" -> riops = Size.from(kv[1]);
+                    case "wbps" -> wbps = Size.from(kv[1]);
+                    case "wiops" -> wiops = Size.from(kv[1]);
+                    default -> throw new IllegalArgumentException("Unknown key " + kv[0]);
+                }
+            }
+            return new Max(rbps, wbps, riops, wiops);
+        }
+    }
+
+    /**
+     * Returns the maximum allowed IO usage, by device, or empty if cgroup is not found.
+     *
+     * @see Max
+     */
+    public Optional<Map<Device, Max>> readMax() {
+        return cgroup.readIfExists("io.max")
+                     .map(content -> content
+                             .lines()
+                             .map(line -> {
+                                 String[] parts = line.strip().split(" ", 2);
+                                 return Map.entry(Device.fromString(parts[0]), Max.fromString(parts[1]));
+                             })
+                             .collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue)));
+    }
+
+    public boolean updateMax(TaskContext context, Device device, Max max) {
+        Max prevMax = readMax()
+                .map(maxByDevice -> maxByDevice.get(device))
+                .orElse(Max.UNLIMITED);
+        if (prevMax.equals(max)) return false;
+
+        UnixPath path = cgroup.unixPath().resolve("io.max");
+        context.recordSystemModification(logger, "Updating %s for device %s from '%s' to '%s'",
+                path, device.toFileContent(), prevMax.toFileContent(), max.toFileContent());
+        path.writeUtf8File(device.toFileContent() + ' ' + max.toFileContent() + '\n');
+        return true;
+    }
+
+}
diff --git a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/Size.java b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/Size.java
index 5e6ca7de8bd..a8cbe2e8afe 100644
--- a/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/Size.java
+++ b/node-admin/src/main/java/com/yahoo/vespa/hosted/node/admin/cgroup/Size.java
@@ -10,12 +10,13 @@ import java.util.Objects;
  */
 public class Size {
     private static final String MAX = "max";
+    private static final Size MAX_SIZE = new Size(true, 0);
 
     private final boolean max;
     private final long value;
 
     public static Size max() {
-        return new Size(true, 0);
+        return MAX_SIZE;
     }
 
     public static Size from(long value) {
@@ -23,7 +24,7 @@ public class Size {
     }
 
     public static Size from(String value) {
-        return value.equals(MAX) ? new Size(true, 0) : new Size(false, Long.parseLong(value));
+        return value.equals(MAX) ? MAX_SIZE : new Size(false, Long.parseLong(value));
     }
 
     private Size(boolean max, long value) {
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/cgroup/CgroupTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/cgroup/CgroupTest.java
index 27580082020..c93d90329f7 100644
--- a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/cgroup/CgroupTest.java
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/cgroup/CgroupTest.java
@@ -1,6 +1,4 @@
 // Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
-
-// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 package com.yahoo.vespa.hosted.node.admin.cgroup;
 
 import com.yahoo.vespa.hosted.node.admin.container.ContainerId;
@@ -10,7 +8,6 @@ import com.yahoo.vespa.hosted.node.admin.task.util.file.UnixPath;
 import com.yahoo.vespa.test.file.TestFileSystem;
 import org.junit.jupiter.api.Test;
 
-import java.io.IOException;
 import java.nio.file.FileSystem;
 import java.util.Map;
 import java.util.Optional;
@@ -23,6 +20,8 @@ import static com.yahoo.vespa.hosted.node.admin.cgroup.CpuController.StatField.T
 import static com.yahoo.vespa.hosted.node.admin.cgroup.CpuController.StatField.USER_USAGE_USEC;
 import static com.yahoo.vespa.hosted.node.admin.cgroup.CpuController.sharesToWeight;
 import static com.yahoo.vespa.hosted.node.admin.cgroup.CpuController.weightToShares;
+import static com.yahoo.vespa.hosted.node.admin.cgroup.IoController.Device;
+import static com.yahoo.vespa.hosted.node.admin.cgroup.IoController.Max;
 import static org.junit.jupiter.api.Assertions.assertEquals;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
@@ -75,35 +74,39 @@ public class CgroupTest {
     }
 
     @Test
-    public void reads_cpu_stats() throws IOException {
-        cgroupRoot.resolve("cpu.stat").writeUtf8File("usage_usec 17794243\n" +
-                "user_usec 16099205\n" +
-                "system_usec 1695038\n" +
-                "nr_periods 12465\n" +
-                "nr_throttled 25\n" +
-                "throttled_usec 14256\n");
+    public void reads_cpu_stats() {
+        cgroupRoot.resolve("cpu.stat").writeUtf8File("""
+                usage_usec 17794243
+                user_usec 16099205
+                system_usec 1695038
+                nr_periods 12465
+                nr_throttled 25
+                throttled_usec 14256
+                """);
 
         assertEquals(Map.of(TOTAL_USAGE_USEC, 17794243L, USER_USAGE_USEC, 16099205L, SYSTEM_USAGE_USEC, 1695038L,
                 TOTAL_PERIODS, 12465L, THROTTLED_PERIODS, 25L, THROTTLED_TIME_USEC, 14256L), containerCgroup.cpu().readStats());
     }
 
     @Test
-    public void reads_memory_metrics() throws IOException {
+    public void reads_memory_metrics() {
         cgroupRoot.resolve("memory.current").writeUtf8File("2525093888\n");
         assertEquals(2525093888L, containerCgroup.memory().readCurrent().value());
 
         cgroupRoot.resolve("memory.max").writeUtf8File("4322885632\n");
         assertEquals(4322885632L, containerCgroup.memory().readMax().value());
 
-        cgroupRoot.resolve("memory.stat").writeUtf8File("anon 3481600\n" +
-                "file 69206016\n" +
-                "kernel_stack 73728\n" +
-                "slab 3552304\n" +
-                "percpu 262336\n" +
-                "sock 73728\n" +
-                "shmem 8380416\n" +
-                "file_mapped 1081344\n" +
-                "file_dirty 135168\n");
+        cgroupRoot.resolve("memory.stat").writeUtf8File("""
+                anon 3481600
+                file 69206016
+                kernel_stack 73728
+                slab 3552304
+                percpu 262336
+                sock 73728
+                shmem 8380416
+                file_mapped 1081344
+                file_dirty 135168
+                """);
         assertEquals(69206016L, containerCgroup.memory().readFileSystemCache().value());
     }
 
@@ -117,4 +120,37 @@ public class CgroupTest {
                     () -> "Original shares: " + originalShares + ", round trip shares: " + roundTripShares + ", diff: " + diff);
         }
     }
+
+    @Test
+    void reads_io_max() {
+        assertEquals(Optional.empty(), containerCgroup.io().readMax());
+
+        cgroupRoot.resolve("io.max").writeUtf8File("");
+        assertEquals(Optional.of(Map.of()), containerCgroup.io().readMax());
+
+        cgroupRoot.resolve("io.max").writeUtf8File("""
+                253:1 rbps=11 wbps=max riops=22 wiops=33
+                253:0 rbps=max wbps=44 riops=max wiops=55
+                """);
+        assertEquals(Map.of(new Device(253, 1), new Max(Size.from(11), Size.max(), Size.from(22), Size.from(33)),
+                            new Device(253, 0), new Max(Size.max(), Size.from(44), Size.max(), Size.from(55))),
+                     containerCgroup.io().readMax().orElseThrow());
+    }
+
+    @Test
+    void writes_io_max() {
+        Device device = new Device(253, 0);
+        Max initial = new Max(Size.max(), Size.from(44), Size.max(), Size.from(55));
+        assertTrue(containerCgroup.io().updateMax(context, device, initial));
+        assertEquals("253:0 rbps=max wbps=44 riops=max wiops=55\n", cgroupRoot.resolve("io.max").readUtf8File());
+
+        cgroupRoot.resolve("io.max").writeUtf8File("""
+                253:1 rbps=11 wbps=max riops=22 wiops=33
+                253:0 rbps=max wbps=44 riops=max wiops=55
+                """);
+        assertFalse(containerCgroup.io().updateMax(context, device, initial));
+
+        cgroupRoot.resolve("io.max").writeUtf8File("");
+        assertFalse(containerCgroup.io().updateMax(context, device, Max.UNLIMITED));
+    }
 }
diff --git a/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/cgroup/IoControllerTest.java b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/cgroup/IoControllerTest.java
new file mode 100644
index 00000000000..d2a4ebbfbbd
--- /dev/null
+++ b/node-admin/src/test/java/com/yahoo/vespa/hosted/node/admin/cgroup/IoControllerTest.java
@@ -0,0 +1,26 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+package com.yahoo.vespa.hosted.node.admin.cgroup;
+
+import org.junit.jupiter.api.Test;
+
+import static com.yahoo.vespa.hosted.node.admin.cgroup.IoController.Device;
+import static com.yahoo.vespa.hosted.node.admin.cgroup.IoController.Max;
+import static org.junit.jupiter.api.Assertions.assertEquals;
+
+/**
+ * @author freva
+ */
+class IoControllerTest {
+
+    @Test
+    void device_number_parsing() {
+        assertEquals(new Device(253, 15), Device.fromDeviceNumber(253 * 256 + 15));
+        assertEquals(new Device(345, 123), Device.fromDeviceNumber(345 * 256 + 123));
+    }
+
+    @Test
+    void parse_io_max() {
+        assertEquals(Max.UNLIMITED, Max.fromString(""));
+        assertEquals(new Max(Size.from(1), Size.max(), Size.max(), Size.max()), Max.fromString("rbps=1 wiops=max"));
+    }
+}