From bd0fea358f966e4a0a710b23cbdea1c3a830fb76 Mon Sep 17 00:00:00 2001
From: BinBin <hebinsquared@gmail.com>
Date: Sun, 18 Jun 2023 10:46:31 -0700
Subject: [PATCH] Update: Implemented cgroups v1 mvp for printing out specific
 metrics

---
 README.md |  2 +-
 main.go   | 44 +++++++++++++++++++++++++++++++++++---------
 2 files changed, 36 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 6150f6d..cc3c9ed 100644
--- a/README.md
+++ b/README.md
@@ -1,5 +1,5 @@
 # Cgroup Stats CLI Tool
-This repository contains a command line tool written in Go that uses the containerd/cgroups library to print out cgroup (control group) statistics for a given cgroup. You can specify the type of stats to print out using command line flags.
+This repository contains a command line tool written in Go that uses the containerd/cgroups library to print out cgroup (control group) statistics for a given cgroup. This package currently only supports cgroupsV1. You can specify the type of stats to print out using command line flags.
 
 ### Dependencies
 This tool requires the containerd/cgroups Go library. You can install it using the following command:
diff --git a/main.go b/main.go
index 0fcf053..7a59257 100644
--- a/main.go
+++ b/main.go
@@ -6,6 +6,7 @@ import (
 	"os"
 
 	"github.com/containerd/cgroups"
+	v1 "github.com/containerd/cgroups/stats/v1"
 )
 
 type cliOptions struct {
@@ -18,8 +19,7 @@ type cliOptions struct {
 }
 
 func parseFlags(flags *cliOptions) *cliOptions {
-
-	flags.cgroupPath = flag.String("path", "", "The cgroup path to read stats.")
+	flags.cgroupPath = flag.String("path", "", "The cgroup path to read stats. Path should not include `/sys/fs/cgroup/` prefix, it should start with your own cgroups name")
 	flags.onlyCPU = flag.Bool("cpu", false, "show cpu stats only")
 	flags.onlyMem = flag.Bool("mem", false, "show mem stats only")
 	flags.onlyPids = flag.Bool("pids", false, "show pids stats only")
@@ -30,9 +30,39 @@ func parseFlags(flags *cliOptions) *cliOptions {
 	return flags
 }
 
+func getMetrics(cg cgroups.Cgroup) *v1.Metrics {
+	metrics, err := cg.Stat(cgroups.IgnoreNotExist)
+	if err != nil {
+		fmt.Println("err retrieving cgroup stats:", err)
+		os.Exit(1)
+	}
 
-func main() {
+	return metrics
+}
 
+func printMetrics(metrics *v1.Metrics, flags *cliOptions) {
+	if *flags.onlyCPU {
+		fmt.Printf("CPU Metrics: %+v\n", metrics.CPU)
+	}
+	if *flags.onlyMem {
+		fmt.Printf("Memory Metrics: %+v\n", metrics.Memory)
+	}
+	if *flags.onlyPids {
+		fmt.Printf("PIDs Metrics: %+v\n", metrics.Pids)
+	}
+	if *flags.onlyBlkio {
+		fmt.Printf("BlkIO Metrics: %+v\n", metrics.Blkio)
+	}
+	if *flags.onlyHugetlb {
+		fmt.Printf("HugeTLB Metrics: %+v\n", metrics.Hugetlb)
+	}
+
+	if !*flags.onlyCPU && !*flags.onlyMem && !*flags.onlyPids && !*flags.onlyBlkio && !*flags.onlyHugetlb {
+		fmt.Printf("Cgroup metrics: %+v\n", metrics)
+	}
+}
+
+func main() {
 	flags := parseFlags(new(cliOptions))
 
 	if *flags.cgroupPath == "" {
@@ -46,11 +76,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	metrics, err := cg.Stat(cgroups.IgnoreNotExist)
-	if err != nil {
-		fmt.Println("err retrieving memory stats:", err)
-		os.Exit(1)
-	}
+	metrics := getMetrics(cg)
 
-	fmt.Printf("Cgroup metrics: %+v\n", metrics)
+	printMetrics(metrics, flags)
 }