pss-total/src/main.c

#include <assert.h>
#include <ctype.h>
#include <dirent.h>
#include <errno.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include "lib/include/ns_fget_line.h"
#include "lib/include/ns_str.h"

double read_file(const char *filename) {
  FILE *fp;
  int status;
  Str line = {0};
  char *pss;

  double sum = 0;
  size_t line_pss = 0;

  errno = 0;
  fp = fopen(filename, "r");

  if (fp == NULL) {
    return 0;
  }

  while (1) {
    ns_str_init_empty(&line, 256);

    errno = 0;
    status = ns_fget_line(fp, &line);

    if (status == EOF) {
      if (errno != 0) {
        perror("error reading smaps file");
        ns_str_destroy(&line);
        return 0;
      } else {
        ns_str_destroy(&line);
        break;
      }
    }

    status = sscanf(line.arr, "Pss: %lu", &line_pss);

    if (status != 1) {
      ns_str_destroy(&line);
      continue;
    }

    sum += line_pss;

    ns_str_destroy(&line);
  }

  fclose(fp);

  return sum;
}

double get_total_KiB(void) {
  DIR *proc_dp;
  struct dirent *proc_ent;
  Str smaps_path = {0};
  double total_KiB = 0;

  errno = 0;
  proc_dp = opendir("/proc");

  if (proc_dp == NULL) {
    perror("error opening /proc for enumeration");
    assert(proc_dp != NULL);
  }

  while (1) {
    errno = 0;
    proc_ent = readdir(proc_dp);

    if (proc_ent == NULL) {
      if (errno != 0) {
        perror("error enumerating /proc");
        assert(errno == 0);
      } else {
        break;
      }
    }

    // skip if not a PID
    if (!isdigit(proc_ent->d_name[0])) continue;

    ns_str_set(&smaps_path, "/proc/");
    ns_str_push(&smaps_path, proc_ent->d_name);
    ns_str_push(&smaps_path, "/smaps");

    total_KiB += read_file(smaps_path.arr);

    ns_str_destroy(&smaps_path);
  }

  closedir(proc_dp);

  return total_KiB;
}

size_t get_memtotal(void) {
  FILE *meminfo_fp;
  Str line = {0};
  int status;
  size_t memtotal;

  errno = 0;
  meminfo_fp = fopen("/proc/meminfo", "r");

  if (meminfo_fp == NULL) {
    perror("cannot open file pointer to /proc/meminfo");
    assert(meminfo_fp != NULL);
  }

  ns_str_init_empty(&line, 256);

  errno = 0;
  status = ns_fget_line(meminfo_fp, &line);

  if (status == EOF) {
    if (errno != 0) {
      perror("error reading meminfo");
      fclose(meminfo_fp);
      ns_str_destroy(&line);
      assert(errno == 0);
    } else {
      fprintf(stderr, "/proc/meminfo is empty... something is wrong.\n");
      fclose(meminfo_fp);
      ns_str_destroy(&line);
      exit(1);
    }
  }

  status = sscanf(line.arr, "MemTotal: %lu", &memtotal);

  if (status != 1) {
    fprintf(stderr, "MemTotal is not on the first line of /proc/meminfo??? \n");
    fclose(meminfo_fp);
    ns_str_destroy(&line);
    exit(1);
  }

  fclose(meminfo_fp);
  ns_str_destroy(&line);

  return memtotal;
}

void get_json(char *json, size_t json_len) {
  const char *prefixes[] = {"KiB", "MiB", "GiB"};
  const char *classes[] = {"low", "quarter", "half", "three_quarters",
                           "nine_tenths"};
  char mem_text[16] = {0};
  size_t prefix_index = 0, classes_index = 0;
  double total_KiB = 0, pretty_total = 0, percentage_use = 0;

  total_KiB = get_total_KiB();
  pretty_total = total_KiB;

  while (pretty_total >= 1024 && prefix_index < 2) {
    pretty_total /= 1024;
    prefix_index++;
  }

  snprintf((char *restrict)&mem_text, 16, "%0.2f %s", pretty_total,
           prefixes[prefix_index]);

  percentage_use = total_KiB / get_memtotal();

  classes_index = 0;
  if (percentage_use >= 0.9) {
    classes_index = 4;
  } else if (percentage_use >= 0.75) {
    classes_index = 3;
  } else if (percentage_use >= 0.50) {
    classes_index = 2;
  } else if (percentage_use >= 0.25) {
    classes_index = 1;
  }

  snprintf(json, json_len,
           "{\"text\":\"%s\",\"class\":\"%s\",\"percentage\":%0.2f}", mem_text,
           classes[classes_index], percentage_use * 100);
}

// signal stuff
static volatile sig_atomic_t run = 1;
static void stop(int _) { run = 0; }

#define JSON_LEN 64
int main(int argc, char **argv) {
  char json[JSON_LEN] = {0};

  if (argc == 2) {
    if ((strcmp(argv[1], "--help") == 0) || (strcmp(argv[1], "-h") == 0)) {
      printf("DESCRIPTION\n");
      printf(
          "  pss-total-waybar: sum proportional set size of every process for "
          "accurate memory\n");
      printf(
          "  usage statistics. This branch outputs json for use in waybar.\n");
      printf("\n");
      printf("SUMMARY\n");
      printf("  pss-total-waybar [OPTIONS]\n");
      printf("\n");
      printf("OPTIONS\n");
      printf("  -h, --help\n");
      printf("      print this help message.\n");
      exit(0);
    }
  }

  signal(SIGINT, stop);
  signal(SIGTERM, stop);

  while (run) {
    memset(json, 0, JSON_LEN);
    get_json((char *)&json, JSON_LEN);
    printf("%s\n", json);
    fflush(stdout);
    if (run) {
      sleep(3);
    }
  }

  return 0;
}