c_cpp 内存cgroup中的内存使用情况监视

Posted 2021-05-14

#define __GNU_SOURCE
#include <assert.h>
#include <ctype.h>
#include <errno.h>
#include <fcntl.h>
#include <inttypes.h>
#include <stdint.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

#include <sys/epoll.h>
#include <sys/eventfd.h>
#include <sys/types.h>

static ssize_t
cgroup_memory_limit(void);
static size_t
compute_threshold(size_t limit, double pct);
static ssize_t
current_memory_usage(void);

#define KB (1 * 1024)
#define MB (KB * KB)

#define MAX_EVENTS                                                             \
  5 /* Maximum number of events to be returned from                            \
       a single epoll_wait() call */

#define errExit(msg)                                                           \
  do {                                                                         \
    perror(msg);                                                               \
    exit(EXIT_FAILURE);                                                        \
  } while (0)

#define checkCall(num, msg)                                                    \
  if (num == -1) {                                                             \
    errExit(msg);                                                              \
  }

static char* levels[] = {
  [0] = "yellow",
  [1] = "orange",
  [2] = "red",
};

typedef struct
{
  int low, mid, high;
} limits_t;

typedef struct
{
  int efd; /* event fd */
  int fd;  /* file fd */
  unsigned short level;
  size_t lim;
} cg_ctrl_t;

int
main(int argc, char* argv[])
{
  int epfd, ready, c, j, numOpenFds;
  struct epoll_event ev;
  struct epoll_event evlist[MAX_EVENTS];

  limits_t args = {
    .low = MB * 1, // these values are fallback essentially
    .mid = MB * 3,
    .high = MB * 4,
  };

  if (strcmp(argv[1], "--help") == 0 || strcmp(argv[1], "-h") == 0) {
    fprintf(
      stderr,
      "Usage: %s -l <low pct> -m <mid pct> -h <high pct> " \
      "i.e.: %s -l 0.3 -m 0.5 -h 0.8\n",
      argv[0], argv[0]);
    exit(2);
  }

  while ((c = getopt(argc, argv, "h:l:m:")) != -1)
    switch (c) {
      case 'h':
        args.high =
          compute_threshold(cgroup_memory_limit(), strtod(optarg, NULL));
        break;
      case 'l':
        args.low =
          compute_threshold(cgroup_memory_limit(), strtod(optarg, NULL));
        break;
      case 'm':
        args.mid =
          compute_threshold(cgroup_memory_limit(), strtod(optarg, NULL));
        break;
      case '?':
        if (optopt == 'c')
          fprintf(stderr, "Option -%c requires an argument.\n", optopt);
        else if (isprint(optopt))
          fprintf(stderr, "Unknown option `-%c'.\n", optopt);
        else
          fprintf(stderr, "Unknown option character `\\x%x'.\n", optopt);
        return 1;
      default:
        abort();
    }

    assert(args.low < args.mid);
    assert(args.mid < args.high);

  checkCall((epfd = epoll_create(3)), "epoll_create() failed");

  pid_t pid = getpid();
  int tasks_fd;
  char pid_str[10] = { 0 };
  snprintf(pid_str, 10, "%d", pid);

#define USAGE_IN_BYTES "/sys/fs/cgroup/memory/0/memory.usage_in_bytes"
#define CG_EVENT_CTL "/sys/fs/cgroup/memory/0/cgroup.event_control"
#define TASKS "/sys/fs/cgroup/memory/0/tasks"

  tasks_fd = open(TASKS, O_RDWR);
  checkCall(tasks_fd, "open() failed");
  checkCall(write(tasks_fd, pid_str, 10), "write() failed");
  checkCall(close(tasks_fd), "close() failed");

  int usage_fd = open(USAGE_IN_BYTES, O_RDONLY);
  checkCall(usage_fd, "open failed()");

  int cgcntl_fd = open(CG_EVENT_CTL, O_WRONLY);
  checkCall(cgcntl_fd, "open failed()");

  char control_msg[16] = { 0 }; // Store string for initiating monitoring

  // Limits is used to setup watchpoints with cgroup.
  int limits[3] = {
    [0] = args.low,
    [1] = args.mid,
    [2] = args.high,
  };
  cg_ctrl_t watches[3] = { 0 };

  for (size_t i = 0; i < 3; i++) {
    checkCall((watches[i].efd = eventfd(0, 0)), "eventfd() failed");

    watches[i].lim = (size_t)(limits[i]);
    watches[i].level = (unsigned short)i;
    snprintf(
      control_msg, 16, "%d %d %zu", watches[i].efd, usage_fd, watches[i].lim);
    checkCall(write(cgcntl_fd, control_msg, 16), "write() failed");

    ev.events = EPOLLIN;
    ev.data.fd = watches[i].efd;
    if (epoll_ctl(epfd, EPOLL_CTL_ADD, watches[i].efd, &ev) == -1)
      errExit("epoll_ctl");
  }

  checkCall(close(cgcntl_fd), "close() failed");

  numOpenFds = 3;

#define LEAK_512K (2 << 18)
  while (numOpenFds > 0) {
    /* Fetch up to MAX_EVENTS items from the ready list of the
       epoll instance with a 1 second timeout */

    ready = epoll_wait(epfd, evlist, MAX_EVENTS, 1000);
    if (ready == -1) {
      if (errno == EINTR)
        continue; /* Restart if interrupted by signal */
      else
        errExit("epoll_wait()");
    }

    printf("Leaking...\n");
    void* p = malloc(LEAK_512K);
    if (!p) {
      errExit("malloc()");
    } else {
      memset(p, 0, LEAK_512K);
    }

    /* Deal with returned list of events */

    for (j = 0; j < ready; j++) {
      if (evlist[j].events & EPOLLIN) {
        eventfd_t value;
        checkCall(eventfd_read(evlist[j].data.fd, &value),
                  "eventfd_read() failed");
        assert(value == 1);
        ssize_t used;
        checkCall((used = current_memory_usage()),
                  "current_memory_usage() failed");
        for (size_t i = 0; i < 3; i++) {
          if (watches[i].efd == evlist[j].data.fd) {
            printf("limit reached (fd=%d) | %zu => %s | exceeded by: %zd\n",
                   evlist[j].data.fd,
                   watches[i].lim,
                   levels[watches[i].level],
                   used - watches[i].lim);
          }
        }
      } else if (evlist[j].events & (EPOLLHUP | EPOLLERR)) {
        printf("closing fd %d\n", evlist[j].data.fd);
        if (close(evlist[j].data.fd) == -1)
          errExit("close");
        numOpenFds--;
      }
    }
  }

  printf("All file descriptors closed; bye\n");
  exit(EXIT_SUCCESS);
}

static ssize_t
cgroup_memory_limit(void)
{
#define MEMORY_LIMIT_IN_BYTES "/sys/fs/cgroup/memory/0/memory.limit_in_bytes"
  FILE* fp = fopen(MEMORY_LIMIT_IN_BYTES, "rb");
  if (!fp) {
    return -1;
  }
  uint64_t n;
  fscanf(fp, "%ld", &n);
  checkCall(fclose(fp), "fclose() failed");
  return n;
}

static size_t
compute_threshold(size_t limit, double pct)
{
  assert(pct < 1.);
  return (double)limit * pct;
}

static ssize_t
current_memory_usage(void)
{
#define MEMORY_USAGE_IN_BYTES "/sys/fs/cgroup/memory/0/memory.usage_in_bytes"
  FILE* fp = fopen(MEMORY_USAGE_IN_BYTES, "rb");
  if (!fp) {
    return -1;
  }
  uint64_t n = 0;
  assert(fscanf(fp, "%ld", &n) > 0);
  checkCall(fclose(fp), "fclose() failed");
  return n;
}