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

1 files changed, 114 insertions, 0 deletions

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;
+    }
+
+}