From 1fa3473911791fa2121538d0223fdc01fc058d26 Mon Sep 17 00:00:00 2001
From: Trey Blancher <trey@blancher.net>
Date: Mon, 7 Aug 2023 21:36:51 -0400
Subject: [PATCH] Added memory PSI info, also flush stdount for printf

---
 monitor.c | 51 +++++++++++++++++++++++++++++++++++++++------------
 1 file changed, 39 insertions(+), 12 deletions(-)

diff --git a/monitor.c b/monitor.c
index 6355fc2..e7e1e24 100644
--- a/monitor.c
+++ b/monitor.c
@@ -8,14 +8,18 @@
 
 #define CPU_TRACKING_WINDOW_SECS    1
 #define IO_TRACKING_WINDOW_SECS     1
+#define MEM_TRACKING_WINDOW_SECS    1
 #define CPU_TRIGGER_THRESHOLD_MS    100
 #define IO_TRIGGER_THRESHOLD_MS     100
+#define MEM_TRIGGER_THRESHOLD_MS    100
 #define CPU_PRESSURE_FILE           "/proc/pressure/cpu"
 #define IO_PRESSURE_FILE            "/proc/pressure/io"
+#define MEM_PRESSURE_FILE           "/proc/pressure/memory"
 #define FD_CPU_IDX                  0
 #define FD_IO_IDX                   1
+#define FD_MEM_IDX                  2
 
-struct pollfd fds[2];
+struct pollfd fds[3];
 
 /*
  One function that prints the system call and the error details
@@ -47,38 +51,53 @@ void setup_polling() {
     fds[FD_IO_IDX].fd = open(IO_PRESSURE_FILE, O_RDWR | O_NONBLOCK);
     if (fds[FD_IO_IDX].fd < 0)
         fatal_error("open(): " IO_PRESSURE_FILE);
+    /* Let's setup our MEM PSI trigger */
+    fds[FD_MEM_IDX].fd = open(MEM_PRESSURE_FILE, O_RDWR | O_NONBLOCK);
+    if (fds[FD_MEM_IDX].fd < 0)
+        fatal_error("open(): " MEM_PRESSURE_FILE);
 
-    fds[FD_CPU_IDX].events = fds[FD_IO_IDX].events = POLLPRI;
+    fds[FD_CPU_IDX].events = fds[FD_MEM_IDX].events = fds[FD_IO_IDX].events = POLLPRI;
 
     char trigger[128];
+
     snprintf(trigger, 128, "some %d %d", CPU_TRIGGER_THRESHOLD_MS * 1000, CPU_TRACKING_WINDOW_SECS * 1000000);
     printf("Trigger: %s\n", trigger);
+    fflush(stdout);
     if (write(fds[FD_CPU_IDX].fd, trigger, strlen(trigger) + 1) < 0)
         fatal_error("write(): " CPU_PRESSURE_FILE);
+
+    snprintf(trigger, 128, "some %d %d", MEM_TRIGGER_THRESHOLD_MS * 1000, MEM_TRACKING_WINDOW_SECS * 1000000);
+    printf("Trigger: %s\n", trigger);
+    fflush(stdout);
+    if (write(fds[FD_MEM_IDX].fd, trigger, strlen(trigger) + 1) < 0)
+        fatal_error("write(): " MEM_PRESSURE_FILE);
+
     snprintf(trigger, 128, "some %d %d", IO_TRIGGER_THRESHOLD_MS * 1000, IO_TRACKING_WINDOW_SECS * 1000000);
     printf("Trigger: %s\n", trigger);
+    fflush(stdout);
     if (write(fds[FD_IO_IDX].fd, trigger, strlen(trigger) + 1) < 0)
         fatal_error("write(): " IO_PRESSURE_FILE);
 }
 
 
 /*
- * This is the main function where we wait for notifications from
- * PSI. We increment 2 separate variables that track CPU and I/O
- * notification counts separately and print them.
- * */
+ * This is the main function where we wait for notifications from PSI. We
+ * increment 2 separate variables that track CPU, MEM and I/O notification
+ * counts separately and print them.  
+*/
 
 void wait_for_notification() {
     int cpu_event_counter = 1;
+    int mem_event_counter = 1;
     int io_event_counter = 1;
 
     while (1) {
-        int n = poll(fds, 2, -1);
+        int n = poll(fds, 3, -1);
         if (n < 0) {
             fatal_error("poll()");
         }
 
-        for (int i = 0; i < 2; i++) {
+        for (int i = 0; i < 3; i++) {
 
             /* If the fd of the current iteration does not have any
              * events, move on to the next fd.
@@ -91,10 +110,18 @@ void wait_for_notification() {
                 exit(1);
             }
             if (fds[i].revents & POLLPRI) {
-                if (i == FD_CPU_IDX)
+                if (i == FD_CPU_IDX) {
                     printf("CPU PSI event %d triggered.\n", cpu_event_counter++);
-                else
-                    printf("I/O PSI event %d triggered.\n", io_event_counter++);
+                    fflush(stdout);
+                }
+                else if (i == FD_MEM_IDX) {
+                    printf("MEM PSI event %d triggered.\n", mem_event_counter++);
+                    fflush(stdout);
+                }
+                else {
+                    printf("IO PSI event %d triggered.\n", io_event_counter++);
+                    fflush(stdout);
+                }
             } else {
                 fprintf(stderr, "Unrecognized event: 0x%x.\n", fds[i].revents);
                 exit(1);
@@ -123,4 +150,4 @@ int main() {
     setup_polling();
     wait_for_notification();
     return 0;
-}
\ No newline at end of file
+}