1 #include <errno.h>
2 #include <fcntl.h>
3 #include <getopt.h>
4 #include <stdio.h>
5 #include <stdlib.h>
6 #include <string.h>
7 #include <sys/mman.h>
8 #include <sys/wait.h>
9 #include <unistd.h>
10 
alloc_set(size_t size)11 void* alloc_set(size_t size) {
12     void* addr = NULL;
13 
14     addr = malloc(size);
15     if (!addr) {
16         printf("Allocating %zd MB failed\n", size / 1024 / 1024);
17     } else {
18         memset(addr, 0, size);
19     }
20     return addr;
21 }
22 
add_pressure(size_t * shared,size_t size,size_t step_size,size_t duration,const char * oom_score)23 void add_pressure(size_t* shared, size_t size, size_t step_size, size_t duration,
24                   const char* oom_score) {
25     int fd, ret;
26 
27     fd = open("/proc/self/oom_score_adj", O_WRONLY);
28     ret = write(fd, oom_score, strlen(oom_score));
29     if (ret < 0) {
30         printf("Writing oom_score_adj failed with err %s\n", strerror(errno));
31     }
32     close(fd);
33 
34     if (alloc_set(size)) {
35         *shared = size;
36     }
37 
38     while (alloc_set(step_size)) {
39         size += step_size;
40         *shared = size;
41         usleep(duration);
42     }
43 }
44 
usage()45 void usage() {
46     printf("Usage: [OPTIONS]\n\n"
47            "  -d N: Duration in microsecond to sleep between each allocation.\n"
48            "  -i N: Number of iterations to run the alloc process.\n"
49            "  -o N: The oom_score to set the child process to before alloc.\n"
50            "  -s N: Number of bytes to allocate in an alloc process loop.\n");
51 }
52 
main(int argc,char * argv[])53 int main(int argc, char* argv[]) {
54     pid_t pid;
55     size_t* shared;
56     int c, i = 0;
57 
58     size_t duration = 1000;
59     int iterations = 0;
60     const char* oom_score = "899";
61     size_t step_size = 2 * 1024 * 1024;  // 2 MB
62     size_t size = step_size;
63 
64     while ((c = getopt(argc, argv, "hi:d:o:s:")) != -1) {
65         switch (c) {
66             case 'i':
67                 iterations = atoi(optarg);
68                 break;
69             case 'd':
70                 duration = atoi(optarg);
71                 break;
72             case 'o':
73                 oom_score = optarg;
74                 break;
75             case 's':
76                 step_size = atoi(optarg);
77                 break;
78             case 'h':
79                 usage();
80                 abort();
81             default:
82                 abort();
83         }
84     }
85 
86     shared = (size_t*)mmap(NULL, sizeof(size_t), PROT_READ | PROT_WRITE, MAP_ANONYMOUS | MAP_SHARED,
87                            0, 0);
88 
89     while (iterations == 0 || i < iterations) {
90         *shared = 0;
91         pid = fork();
92         if (!pid) {
93             /* Child */
94             add_pressure(shared, size, step_size, duration, oom_score);
95             /* Shoud not get here */
96             exit(0);
97         } else {
98             wait(NULL);
99             printf("Child %d allocated %zd MB\n", i, *shared / 1024 / 1024);
100             size = *shared / 2;
101         }
102         i++;
103     }
104 }
105