1/*
2 * Copyright 2014, NICTA
3 *
4 * This software may be distributed and modified according to the terms of
5 * the BSD 2-Clause license. Note that NO WARRANTY is provided.
6 * See "LICENSE_BSD2.txt" for details.
7 *
8 * @TAG(NICTA_BSD)
9 */
10
11/*
12 * KILL processes that consume more than X bytes of RAM.
13 *
14 * When it triggers, it will write to syslog stating the process stopped.
15 *
16 * 2012 David Greenaway
17 */
18
19#include <stdio.h>
20#include <stdlib.h>
21#include <stdint.h>
22
23#include <sys/types.h>
24#include <sys/sysinfo.h>
25#include <sys/time.h>
26#include <sys/resource.h>
27
28#include <dirent.h>
29#include <unistd.h>
30#include <signal.h>
31#include <syslog.h>
32#include <string.h>
33
34#define MAX_PATH_SIZE 1024
35
36/* Scheduling priority we should run at. */
37#define SCHED_PRIO (-10)
38
39void
40fatal(const char *str)
41{
42    printf("%s\n", str);
43    exit(1);
44}
45
46/* Iterate through processes in the system. */
47void
48iterate_processes(void (*proc_fn)(int, void *), void *data)
49{
50    /* Open /proc */
51    DIR *proc_dir = opendir("/proc");
52    if (proc_dir == NULL) {
53        fprintf(stderr, "Could not open /proc.");
54        exit(1);
55    }
56
57    /* Read through processes. */
58    while (1) {
59        /* Read directory. */
60        struct dirent *e = readdir(proc_dir);
61        if (e == NULL)
62            break;
63
64        /* Skip non-directories. */
65        if ((e->d_type & DT_DIR) == 0)
66            continue;
67
68        /* Process? */
69        int p = atoi(e->d_name);
70        if (p != 0)
71            proc_fn(p, data);
72    }
73
74    /* Cleanup. */
75    closedir(proc_dir);
76}
77
78void test_process(int p, void *data)
79{
80    /*
81     * Test process with pid 'p'.
82     *
83     * We must be aware that the process might have died between the time
84     * we found out about it and the time we test it; we deal with such
85     * race conditions by skipping over the process.
86     */
87
88    uint64_t max_usage_mb = *(uint64_t *)data;
89    char buf[MAX_PATH_SIZE];
90    unsigned long vmem_usage = 0;
91    unsigned long rmem_usage = 0;
92    FILE *f;
93    int n;
94
95    /* Read memory usage of process. */
96    sprintf(buf, "/proc/%d/statm", p);
97    f = fopen(buf, "r");
98    if (f == NULL)
99        return;
100    n = fscanf(f, "%lu %lu", &vmem_usage, &rmem_usage);
101    if (n != 2) {
102        /* This may still not return anything if the process dies between us
103         * opening the file and performing this read. */
104        fclose(f);
105        return;
106    }
107    fclose(f);
108
109    /* Convert from pages to megabytes. */
110    rmem_usage = rmem_usage * (4096 / 1024) / 1024;
111
112    /* Kill the proces if it is too big. */
113    if (rmem_usage >= max_usage_mb) {
114        (void)kill(p, SIGKILL);
115        syslog(LOG_ALERT,
116                "killbig: Sending SIGKILL to pid %d (process size of %lu MB exceeds %lu MB).\n",
117                p, rmem_usage, max_usage_mb);
118    }
119}
120
121void usage(int argc, char **argv)
122{
123    printf("\n"
124        "usage: %s <max mem usage in MB>\n\n",
125        argc > 0 ? argv[0] : "killbig");
126}
127
128int main(int argc, char **argv)
129{
130    uint64_t max_mem_usage;
131
132    /* Determine which signal to send. */
133    if (argc < 2) {
134        usage(argc, argv);
135        return 1;
136    }
137    max_mem_usage = (uint64_t)atoll(argv[1]);
138    if (max_mem_usage <= 16) {
139        usage(argc, argv);
140        return 1;
141    }
142
143    /* Set our scheduling priority higher. */
144    (void)setpriority(PRIO_PROCESS, 0, SCHED_PRIO);
145
146    /* Monitor. */
147    while (1) {
148        iterate_processes(test_process, &max_mem_usage);
149        sleep(5);
150    }
151}
152