Adding support for cgroup v2 (#499)

* Adding support for cgroup v2 * Address comments --------- Co-authored-by: Sundaram Ananthanarayanan <[email protected]>
Netflix · Jul 28, 2023 · ea8ed4f · ea8ed4f
1 parent d5129e8
commit ea8ed4f
Show file tree

Hide file tree

Showing 12 changed files with 202 additions and 12 deletions.
diff --git a/...erver/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutorStarter.java b/...erver/mantis-server-agent/src/main/java/io/mantisrx/server/agent/TaskExecutorStarter.java
@@ -33,9 +33,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Optional;
-import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.Executor;
-import java.util.concurrent.TimeUnit;
 import javax.annotation.Nullable;
 import lombok.AccessLevel;
 import lombok.RequiredArgsConstructor;

diff --git a/...er/mantis-server-agent/src/main/java/io/mantisrx/server/agent/metrics/cgroups/Cgroup.java b/...er/mantis-server-agent/src/main/java/io/mantisrx/server/agent/metrics/cgroups/Cgroup.java
@@ -17,9 +17,18 @@
 package io.mantisrx.server.agent.metrics.cgroups;
 
 import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 
 interface Cgroup {
+
+    /**
+     * Represents if this cgroup is a v1 or v2 cgroup
+     * @return
+     */
+    Boolean isV1();
+
+    List<Long> getMetrics(String subsystem, String metricName) throws IOException;
     Long getMetric(String subsystem, String metricName) throws IOException;
     Map<String, Long> getStats(String subsystem, String stat) throws IOException;
 }
diff --git a/...antis-server-agent/src/main/java/io/mantisrx/server/agent/metrics/cgroups/CgroupImpl.java b/...antis-server-agent/src/main/java/io/mantisrx/server/agent/metrics/cgroups/CgroupImpl.java
@@ -22,31 +22,69 @@
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.nio.file.Paths;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
 import java.util.Map;
+import java.util.Objects;
 import java.util.stream.Collectors;
+import lombok.AccessLevel;
+import lombok.Getter;
 import lombok.RequiredArgsConstructor;
 
 @RequiredArgsConstructor
 public class CgroupImpl implements Cgroup {
+
     private final String path;
 
+    /**
+     * Maybe change this to the below command in the future.
+     * <p>
+     * ``` stat -fc %T /sys/fs/cgroup/ ``` This should return cgroup2fs for cgroupv2 and tmpfs for
+     * cgroupv1.
+     */
+    @Getter(lazy = true, value = AccessLevel.PRIVATE)
+    private final boolean old = getSubsystems().size() > 0;
+
+    @Override
+    public Boolean isV1() {
+        return isOld();
+    }
+
+    @Override
+    public List<Long> getMetrics(String subsystem, String metricName) throws IOException {
+        Path metricPath = Paths.get(path, subsystem, metricName);
+        try {
+            return
+                Files.readAllLines(metricPath)
+                    .stream()
+                    .findFirst()
+                    .map(s -> Arrays.asList(s.split(" ")))
+                    .orElse(Collections.emptyList())
+                    .stream()
+                    .map(CgroupImpl::convertStringToLong)
+                    .collect(Collectors.toList());
+        } catch (Exception e) {
+            throw new IOException(e);
+        }
+    }
+
     @Override
     public Long getMetric(String subsystem, String metricName) throws IOException {
         Path metricPath = Paths.get(path, subsystem, metricName);
         try {
-            return Files.readAllLines(metricPath).stream().findFirst().map(CgroupImpl::convertStringToLong).get();
+            return Files.readAllLines(metricPath).stream().findFirst()
+                .map(CgroupImpl::convertStringToLong).get();
         } catch (Exception e) {
             throw new IOException(e);
         }
     }
 
     /**
-     * Example usage:
-     * user 43873627
-     * system 4185541
+     * Example usage: user 43873627 system 4185541
      *
      * @param subsystem subsystem the stat file is part of.
-     * @param stat name of the stat file
+     * @param stat      name of the stat file
      * @return map of metrics to their corresponding values
      * @throws IOException
      */
@@ -58,7 +96,8 @@ public Map<String, Long> getStats(String subsystem, String stat) throws IOExcept
                 .stream()
                 .map(l -> {
                     String[] parts = l.split("\\s+");
-                    Preconditions.checkArgument(parts.length == 2, "Expected two parts only but was {} parts", parts.length);
+                    Preconditions.checkArgument(parts.length == 2,
+                        "Expected two parts only but was {} parts", parts.length);
                     return new Tuple2<>(parts[0], convertStringToLong(parts[1]));
                 })
                 .collect(Collectors.toMap(t -> t._1, t -> t._2));
@@ -68,8 +107,8 @@ public Map<String, Long> getStats(String subsystem, String stat) throws IOExcept
      * Convert a number from its string representation to a long.
      *
      * @param strval: value to convert
-     * @return The converted long value. Long max value is returned if the
-     *         string representation exceeds the range of type long.
+     * @return The converted long value. Long max value is returned if the string representation
+     * exceeds the range of type long.
      */
     private static long convertStringToLong(String strval) {
         try {
@@ -80,4 +119,13 @@ private static long convertStringToLong(String strval) {
             return Long.MAX_VALUE;
         }
     }
+
+    private List<String> getSubsystems() {
+        return
+            Arrays.asList(Objects.requireNonNull(Paths.get(path).toFile().listFiles()))
+                .stream()
+                .filter(s -> s.isDirectory())
+                .map(s -> s.getName())
+                .collect(Collectors.toList());
+    }
 }
diff --git a/...gent/src/main/java/io/mantisrx/server/agent/metrics/cgroups/CpuAcctsSubsystemProcess.java b/...gent/src/main/java/io/mantisrx/server/agent/metrics/cgroups/CpuAcctsSubsystemProcess.java
@@ -18,6 +18,7 @@
 
 import io.mantisrx.runtime.loader.config.Usage;
 import java.io.IOException;
+import java.util.List;
 import java.util.Map;
 import java.util.Optional;
 import lombok.AccessLevel;
@@ -35,6 +36,14 @@ class CpuAcctsSubsystemProcess implements SubsystemProcess {
 
     @Override
     public void getUsage(Usage.UsageBuilder resourceUsageBuilder) throws IOException {
+        if (cgroup.isV1()) {
+            handleV1(resourceUsageBuilder);
+        } else {
+            handleV2(resourceUsageBuilder);
+        }
+    }
+
+    private void handleV1(Usage.UsageBuilder resourceUsageBuilder) throws IOException {
         final Map<String, Long> stat = cgroup.getStats("cpuacct", "cpuacct.stat");
         Optional<Long> user = Optional.ofNullable(stat.get("user"));
         Optional<Long> system = Optional.ofNullable(stat.get("system"));
@@ -61,4 +70,20 @@ public void getUsage(Usage.UsageBuilder resourceUsageBuilder) throws IOException
             log.warn("quota={} & period={} are not configured correctly", quota, period);
         }
     }
+
+    private void handleV2(Usage.UsageBuilder resourceUsageBuilder) throws IOException {
+        Map<String, Long> cpuStats = cgroup.getStats("", "cpu.stat");
+        resourceUsageBuilder
+            .cpusUserTimeSecs(cpuStats.getOrDefault("user_usec", 0L) / 1000.0)
+            .cpusSystemTimeSecs(cpuStats.getOrDefault("system_usec", 0L) / 1000.0);
+
+        List<Long> metrics = cgroup.getMetrics("", "cpu.max");
+        if (metrics.size() != 2) {
+            log.warn("cpu.max metrics={} are not configured correctly", metrics);
+        } else {
+            Long quota = metrics.get(0);
+            Long period = metrics.get(1);
+            resourceUsageBuilder.cpusLimit(quota / (float) period);
+        }
+    }
 }
diff --git a/...-agent/src/main/java/io/mantisrx/server/agent/metrics/cgroups/MemorySubsystemProcess.java b/...-agent/src/main/java/io/mantisrx/server/agent/metrics/cgroups/MemorySubsystemProcess.java
@@ -35,6 +35,14 @@ class MemorySubsystemProcess implements SubsystemProcess {
 
     @Override
     public void getUsage(UsageBuilder usageBuilder) throws IOException {
+        if (cgroup.isV1()) {
+            handleV1(usageBuilder);
+        } else {
+            handleV2(usageBuilder);
+        }
+    }
+
+    private void handleV1(UsageBuilder usageBuilder) throws IOException {
         Long memLimit = cgroup.getMetric("memory", "memory.limit_in_bytes");
         usageBuilder.memLimit(memLimit);
 
@@ -47,4 +55,20 @@ public void getUsage(UsageBuilder usageBuilder) throws IOException {
             log.warn("stats for memory not found; stats={}", stats);
         }
     }
+
+    private void handleV2(UsageBuilder usageBuilder) throws IOException {
+        Long memLimit = cgroup.getMetric("", "memory.max");
+        usageBuilder.memLimit(memLimit);
+
+        Long memCurrent = cgroup.getMetric("", "memory.current");
+        usageBuilder.memRssBytes(memCurrent);
+
+        Map<String, Long> stats = cgroup.getStats("", "memory.stat");
+        Optional<Long> anon = Optional.ofNullable(stats.get("anon"));
+        if (anon.isPresent()) {
+            usageBuilder.memAnonBytes(anon.get());
+        } else {
+            log.warn("stats for memory not found; stats={}", stats);
+        }
+    }
 }
diff --git a/.../src/test/java/io/mantisrx/server/agent/metrics/cgroups/TestCpuAcctsSubsystemProcess.java b/.../src/test/java/io/mantisrx/server/agent/metrics/cgroups/TestCpuAcctsSubsystemProcess.java
@@ -22,14 +22,20 @@
 
 import io.mantisrx.runtime.loader.config.Usage;
 import io.mantisrx.shaded.com.google.common.collect.ImmutableMap;
+import io.mantisrx.shaded.com.google.common.io.Resources;
+import java.io.IOException;
 import org.junit.Test;
 
 public class TestCpuAcctsSubsystemProcess {
-    private final Cgroup cgroup = mock(Cgroup.class);
-    private final CpuAcctsSubsystemProcess process = new CpuAcctsSubsystemProcess(cgroup);
 
     @Test
     public void testWhenCgroupsReturnsCorrectData() throws Exception {
+        final Cgroup cgroup = mock(Cgroup.class);
+        final CpuAcctsSubsystemProcess process = new CpuAcctsSubsystemProcess(cgroup);
+
+        when(cgroup.isV1()).thenReturn(true);
+        when(cgroup.getStats("cpuacct", "cpuacct.stat"))
+            .thenReturn(ImmutableMap.<String, Long>of("user", 43873627L, "system", 4185541L));
         when(cgroup.getStats("cpuacct", "cpuacct.stat"))
             .thenReturn(ImmutableMap.<String, Long>of("user", 43873627L, "system", 4185541L));
         when(cgroup.getMetric("cpuacct", "cpu.cfs_quota_us"))
@@ -44,4 +50,18 @@ public void testWhenCgroupsReturnsCorrectData() throws Exception {
         assertEquals(438736L, (long) usage.getCpusUserTimeSecs());
         assertEquals(41855L, (long) usage.getCpusSystemTimeSecs());
     }
+
+    @Test
+    public void testCgroupsV2() throws IOException {
+        final Cgroup cgroupv2 =
+            new CgroupImpl(Resources.getResource("example2").getPath());
+
+        final CpuAcctsSubsystemProcess process = new CpuAcctsSubsystemProcess(cgroupv2);
+        final Usage.UsageBuilder usageBuilder = Usage.builder();
+        process.getUsage(usageBuilder);
+        final Usage usage = usageBuilder.build();
+        assertEquals(2L, (long) usage.getCpusLimit());
+        assertEquals(4231313L, (long) usage.getCpusUserTimeSecs());
+        assertEquals(1277084L, (long) usage.getCpusSystemTimeSecs());
+    }
 }
diff --git a/...nt/src/test/java/io/mantisrx/server/agent/metrics/cgroups/TestMemorySubsystemProcess.java b/...nt/src/test/java/io/mantisrx/server/agent/metrics/cgroups/TestMemorySubsystemProcess.java
@@ -22,6 +22,8 @@
 
 import io.mantisrx.runtime.loader.config.Usage;
 import io.mantisrx.shaded.com.google.common.collect.ImmutableMap;
+import io.mantisrx.shaded.com.google.common.io.Resources;
+import java.io.IOException;
 import org.junit.Test;
 
 public class TestMemorySubsystemProcess {
@@ -31,6 +33,7 @@ public class TestMemorySubsystemProcess {
 
     @Test
     public void testHappyPath() throws Exception {
+        when(cgroup.isV1()).thenReturn(true);
         when(cgroup.getMetric("memory", "memory.limit_in_bytes"))
             .thenReturn(17179869184L);
         when(cgroup.getStats("memory", "memory.stat"))
@@ -81,4 +84,18 @@ public void testHappyPath() throws Exception {
         assertEquals(14828109824L, (long) usage.getMemRssBytes());
         assertEquals(14828109824L, (long) usage.getMemAnonBytes());
     }
+
+    @Test
+    public void testCgroupv2() throws IOException {
+        final Cgroup cgroupv2 =
+            new CgroupImpl(Resources.getResource("example2").getPath());
+
+        final MemorySubsystemProcess process = new MemorySubsystemProcess(cgroupv2);
+        final Usage.UsageBuilder usageBuilder = Usage.builder();
+        process.getUsage(usageBuilder);
+        final Usage usage = usageBuilder.build();
+        assertEquals(2147483648L, (long) usage.getMemLimit());
+        assertEquals(1693843456L, (long) usage.getMemRssBytes());
+        assertEquals(945483776L, (long) usage.getMemAnonBytes());
+    }
 }
diff --git a/mantis-server/mantis-server-agent/src/test/resources/example2/cpu.max b/mantis-server/mantis-server-agent/src/test/resources/example2/cpu.max
@@ -0,0 +1 @@
+200000 100000
diff --git a/mantis-server/mantis-server-agent/src/test/resources/example2/cpu.stat b/mantis-server/mantis-server-agent/src/test/resources/example2/cpu.stat
@@ -0,0 +1,6 @@
+usage_usec 5508398803
+user_usec 4231313976
+system_usec 1277084827
+nr_periods 2379234
+nr_throttled 18
+throttled_usec 13579
diff --git a/mantis-server/mantis-server-agent/src/test/resources/example2/memory.current b/mantis-server/mantis-server-agent/src/test/resources/example2/memory.current
@@ -0,0 +1 @@
+1693843456
diff --git a/mantis-server/mantis-server-agent/src/test/resources/example2/memory.max b/mantis-server/mantis-server-agent/src/test/resources/example2/memory.max
@@ -0,0 +1 @@
+2147483648
diff --git a/mantis-server/mantis-server-agent/src/test/resources/example2/memory.stat b/mantis-server/mantis-server-agent/src/test/resources/example2/memory.stat
@@ -0,0 +1,40 @@
+anon 945483776
+file 159186944
+kernel_stack 4210688
+pagetables 5844992
+percpu 14941440
+sock 4096
+shmem 3862528
+file_mapped 13840384
+file_dirty 0
+file_writeback 0
+swapcached 0
+anon_thp 0
+file_thp 0
+shmem_thp 0
+inactive_anon 945917952
+active_anon 200704
+inactive_file 143499264
+active_file 15040512
+unevictable 0
+slab_reclaimable 558328568
+slab_unreclaimable 4480536
+slab 562809104
+workingset_refault_anon 0
+workingset_refault_file 0
+workingset_activate_anon 0
+workingset_activate_file 0
+workingset_restore_anon 0
+workingset_restore_file 0
+workingset_nodereclaim 0
+pgfault 121278150
+pgmajfault 19
+pgrefill 0
+pgscan 0
+pgsteal 0
+pgactivate 48641
+pgdeactivate 0
+pglazyfree 785
+pglazyfreed 0
+thp_fault_alloc 0
+thp_collapse_alloc 0