add -h and -m flags to psi-top

tools/psi-top

@@ -1,9 +1,7 @@
 #!/usr/bin/env python3
 
 import os
-
+from argparse import ArgumentParser
-
-###############################################################################
 
 
 def psi_path_to_metrics(psi_path):
@@ -11,9 +9,9 @@ def psi_path_to_metrics(psi_path):
     """
     with open(psi_path) as f:
         psi_list = f.readlines()
-    # print(psi_list)
+
     some_list, full_list = psi_list[0].split(' '), psi_list[1].split(' ')
-    # print(some_list, full_list)
+
     some_avg10 = some_list[1].split('=')[1]
     some_avg60 = some_list[2].split('=')[1]
     some_avg300 = some_list[3].split('=')[1]
@@ -26,6 +24,21 @@ def psi_path_to_metrics(psi_path):
             full_avg10, full_avg60, full_avg300)
 
 
+def psi_path_to_metrics_cpu(psi_path):
+    """
+    """
+    with open(psi_path) as f:
+        psi_list = f.readlines()
+
+    some_list = psi_list[0].rstrip().split(' ')
+
+    some_avg10 = some_list[1].split('=')[1]
+    some_avg60 = some_list[2].split('=')[1]
+    some_avg300 = some_list[3].split('=')[1]
+
+    return (some_avg10, some_avg60, some_avg300)
+
+
 def cgroup2_root():
     """
     """
@@ -35,7 +48,7 @@ def cgroup2_root():
                 return line.partition(cgroup2_separator)[0].partition(' ')[2]
 
 
-def get_psi_mem_files(cgroup2_path):
+def get_psi_mem_files(cgroup2_path, met):
     """
     """
     path_list = []
@@ -43,7 +56,7 @@ def get_psi_mem_files(cgroup2_path):
     for root, dirs, files in os.walk(cgroup2_path):
         for file in files:
             path = os.path.join(root, file)
-            if path.endswith('/memory.pressure'):
+            if path.endswith('/{}.pressure'.format(met)):
                 path_list.append(path)
 
     return path_list
@@ -52,13 +65,43 @@ def get_psi_mem_files(cgroup2_path):
 def psi_path_to_cgroup2(path):
     """
     """
-    return path.partition(cgroup2_mountpoint)[2][:-16]
+
+    if path.endswith('/cpu.pressure'):
+        return path.partition(cgroup2_mountpoint)[
+            2].partition('/cpu.pressure')[0]
+
+    if path.endswith('/io.pressure'):
+        return path.partition(cgroup2_mountpoint)[
+            2].partition('/io.pressure')[0]
+
+    if path.endswith('/memory.pressure'):
+        return path.partition(cgroup2_mountpoint)[
+            2].partition('/memory.pressure')[0]
 
 
-###############################################################################
+parser = ArgumentParser()
+
+parser.add_argument(
+    '-m',
+    '--metrics',
+    help="""metrics (memory, io or cpu)""",
+    default='memory',
+    type=str
+)
 
 
-psi_path = '/proc/pressure/memory'
+args = parser.parse_args()
+
+
+met = args.metrics
+
+
+if not (met == 'memory' or met == 'io' or met == 'cpu'):
+    print('Invalid metrics:', met)
+    exit(1)
+
+
+psi_path = '/proc/pressure/{}'.format(met)
 mounts = '/proc/mounts'
 cgroup2_separator = ' cgroup2 rw,'
 
@@ -68,15 +111,10 @@ cgroup2_mountpoint = cgroup2_root()
 if cgroup2_mountpoint is None:
     print('ERROR: cgroup_v2 hierarchy is not mounted')
     exit(1)
-else:
-    print('cgroup_v2 mountpoint:', cgroup2_mountpoint)
-
-
-psi_support = os.path.exists(psi_path)
 
 
 try:
-    psi_path_to_metrics(psi_path)
+    psi_path_to_metrics('/proc/pressure/memory')
 except Exception as e:
     print('ERROR: {}'.format(e))
     print('PSI metrics are not provided by the kernel. Exit.')
@@ -84,19 +122,54 @@ except Exception as e:
 
 
 if cgroup2_mountpoint is not None:
-    y = get_psi_mem_files(cgroup2_mountpoint)
+    y = get_psi_mem_files(cgroup2_mountpoint, met)
-    for path in y:
-        pass  # print(psi_path_to_cgroup2(path))
 
-path_list = get_psi_mem_files(cgroup2_mountpoint)
 
-print('      avg10  avg60 avg300         avg10  avg60 avg300  cgroup2')
+path_list = get_psi_mem_files(cgroup2_mountpoint, met)
-print('      -----  ----- ------         -----  ----- ------  ---------')
+
+
+head_mem_io = '''PSI metrics: {}
+cgroup_v2 mountpoint: {}
+---------------------------------------------
+         some        |          full        |
+---------------------|--------------------- | ---------
+ avg10  avg60 avg300 |  avg10  avg60 avg300 | cgroup_v2
+ -----  ----- ------ |  -----  ----- ------ | ---------'''.format(
+    met, cgroup2_mountpoint)
+
+
+head_cpu = '''PSI metrics: {}
+cgroup_v2 mountpoint: {}
+----------------------
+         some        |
+---------------------|
+ avg10  avg60 avg300 | cgroup_v2
+ -----  ----- ------ | ---------'''.format(
+    met, cgroup2_mountpoint)
+
+
+if met == 'cpu':
+    print(head_cpu)
+else:
+    print(head_mem_io)
+
+
+if met == 'cpu':
+
+    some_avg10, some_avg60, some_avg300 = psi_path_to_metrics_cpu(psi_path)
+
+    print('{} {} {} | {}'.format(
+        some_avg10.rjust(6),
+        some_avg60.rjust(6),
+        some_avg300.rjust(6),
+        '[SYSTEM_WIDE]'))
+
+else:
 
     (some_avg10, some_avg60, some_avg300, full_avg10, full_avg60, full_avg300
      ) = psi_path_to_metrics(psi_path)
 
-print('some {} {} {} | full {} {} {}  {}'.format(
+    print('{} {} {} | {} {} {} | {}'.format(
         some_avg10.rjust(6),
         some_avg60.rjust(6),
         some_avg300.rjust(6),
@@ -106,10 +179,21 @@ print('some {} {} {} | full {} {} {}  {}'.format(
 
 
 for psi_path in path_list:
+
+    if met == 'cpu':
+        some_avg10, some_avg60, some_avg300 = psi_path_to_metrics_cpu(psi_path)
+
+        print('{} {} {} | {}'.format(
+            some_avg10.rjust(6),
+            some_avg60.rjust(6),
+            some_avg300.rjust(6),
+            psi_path_to_cgroup2(psi_path)))
+
+    else:
         (some_avg10, some_avg60, some_avg300,
          full_avg10, full_avg60, full_avg300) = psi_path_to_metrics(psi_path)
 
-    print('some {} {} {} | full {} {} {}  {}'.format(
+        print('{} {} {} | {} {} {} | {}'.format(
             some_avg10.rjust(6),
             some_avg60.rjust(6),
             some_avg300.rjust(6),