1#include <arpa/inet.h>
2#include <iostream>
3#include <chrono>
4#include <cutils/sockets.h>
5#include <hardware/gralloc.h>
6#include <vector>
7#include <tuple>
8#include <algorithm>
9#include <tuple>
10#include <numeric>
11#include <fcntl.h>
12#include <string>
13#include <fstream>
14#include <sys/stat.h>
15#include <sys/wait.h>
16
17using namespace std;
18
19#define ASSERT_TRUE(cond) \
20do { \
21    if (!(cond)) {\
22       cerr << __func__ << "( " << getpid() << "):" << __LINE__ << " condition:" << #cond << " failed\n" << endl; \
23       exit(EXIT_FAILURE); \
24    } \
25} while (0)
26
27class Pipe {
28    int m_readFd;
29    int m_writeFd;
30    Pipe(const Pipe &) = delete;
31    Pipe& operator=(const Pipe &) = delete;
32    Pipe& operator=(const Pipe &&) = delete;
33public:
34    Pipe(int readFd, int writeFd) : m_readFd{readFd}, m_writeFd{writeFd} {
35        fcntl(m_readFd, F_SETFD, FD_CLOEXEC);
36        fcntl(m_writeFd, F_SETFD, FD_CLOEXEC);
37    }
38    Pipe(Pipe&& rval) noexcept {
39        m_readFd = rval.m_readFd;
40        m_writeFd = rval.m_writeFd;
41        rval.m_readFd = 0;
42        rval.m_writeFd = 0;
43    }
44    ~Pipe() {
45        if (m_readFd)
46            close(m_readFd);
47        if (m_writeFd)
48            close(m_writeFd);
49    }
50    void preserveOverFork(bool preserve) {
51        if (preserve) {
52            fcntl(m_readFd, F_SETFD, 0);
53            fcntl(m_writeFd, F_SETFD,0);
54        } else {
55            fcntl(m_readFd, F_SETFD, FD_CLOEXEC);
56            fcntl(m_writeFd, F_SETFD, FD_CLOEXEC);
57        }
58    }
59    int getReadFd() {
60        return m_readFd;
61    }
62    int getWriteFd() {
63        return m_writeFd;
64    }
65    void signal() {
66        bool val = true;
67        int error = write(m_writeFd, &val, sizeof(val));
68        ASSERT_TRUE(error == sizeof(val));
69    };
70    void wait() {
71        bool val = false;
72        int error = read(m_readFd, &val, sizeof(val));
73        ASSERT_TRUE(error == sizeof(val));
74    }
75    bool wait_ret_error() {
76        bool val = false;
77        int error = read(m_readFd, &val, sizeof(val));
78        return (error != 1);
79    }
80    template <typename T> void send(const T& v) {
81        int error = write(m_writeFd, &v, sizeof(T));
82        ASSERT_TRUE(error >= 0);
83    }
84    template <typename T> void recv(T& v) {
85        int error = read(m_readFd, &v, sizeof(T));
86        ASSERT_TRUE(error >= 0);
87    }
88    static Pipe makePipeFromFds(int readFd, int writeFd) {
89        return Pipe(readFd, writeFd);
90    }
91    static tuple<Pipe, Pipe> createPipePair() {
92        int a[2];
93        int b[2];
94
95        int error1 = pipe(a);
96        int error2 = pipe(b);
97        ASSERT_TRUE(error1 >= 0);
98        ASSERT_TRUE(error2 >= 0);
99
100        return make_tuple(Pipe(a[0], b[1]), Pipe(b[0], a[1]));
101    }
102};
103
104void createProcess(Pipe pipe, const char *exName, const char *arg)
105{
106    pipe.preserveOverFork(true);
107    pid_t pid = fork();
108    // child proc
109    if (pid == 0) {
110        char readFdStr[16];
111        char writeFdStr[16];
112        snprintf(readFdStr, sizeof(readFdStr), "%d", pipe.getReadFd());
113        snprintf(writeFdStr, sizeof(writeFdStr), "%d", pipe.getWriteFd());
114        execl(exName, exName, "--worker", arg, readFdStr, writeFdStr, nullptr);
115        ASSERT_TRUE(0);
116    }
117    // parent process
118    else if (pid > 0) {
119        pipe.preserveOverFork(false);
120        return;
121    }
122    else {
123        ASSERT_TRUE(0);
124    }
125}
126
127
128static void write_oomadj_to_lmkd(int oomadj) {
129    // Connect to lmkd and store our oom_adj
130    int lmk_procprio_cmd[4];
131    int sock;
132    int tries = 10;
133    while ((sock = socket_local_client("lmkd",
134                    ANDROID_SOCKET_NAMESPACE_RESERVED,
135                    SOCK_SEQPACKET)) < 0) {
136        usleep(100000);
137        if (tries-- < 0) break;
138    }
139    if (sock < 0) {
140        cout << "Failed to connect to lmkd, errno " << errno << endl;
141        exit(1);
142    }
143    lmk_procprio_cmd[0] = htonl(1);
144    lmk_procprio_cmd[1] = htonl(getpid());
145    lmk_procprio_cmd[2] = htonl(getuid());
146    lmk_procprio_cmd[3] = htonl(oomadj);
147
148    int written = write(sock, lmk_procprio_cmd, sizeof(lmk_procprio_cmd));
149    cout << "Wrote " << written << " bytes to lmkd control socket." << endl;
150}
151
152#ifdef ENABLE_MEM_CGROUPS
153static void create_memcg() {
154    char buf[256];
155    pid_t pid = getpid();
156    snprintf(buf, sizeof(buf), "/dev/memctl/apps/%u", pid);
157
158    int tasks = mkdir(buf, S_IRWXU | S_IRWXG | S_IROTH | S_IXOTH);
159    if (tasks < 0) {
160        cout << "Failed to create memory cgroup" << endl;
161        return;
162    }
163    snprintf(buf, sizeof(buf), "/dev/memctl/apps/%u/tasks", pid);
164    tasks = open(buf, O_WRONLY);
165    if (tasks < 0) {
166        cout << "Unable to add process to memory cgroup" << endl;
167        return;
168    }
169    snprintf(buf, sizeof(buf), "%u", pid);
170    write(tasks, buf, strlen(buf));
171    close(tasks);
172}
173#endif
174
175size_t s = 4 * (1 << 20);
176void *gptr;
177int main(int argc, char *argv[])
178{
179    if ((argc > 1) && (std::string(argv[1]) == "--worker")) {
180#ifdef ENABLE_MEM_CGROUPS
181        create_memcg();
182#endif
183        write_oomadj_to_lmkd(atoi(argv[2]));
184        Pipe p{atoi(argv[3]), atoi(argv[4])};
185
186        long long allocCount = 0;
187        while (1) {
188            p.wait();
189            char *ptr = (char*)malloc(s);
190            memset(ptr, (int)allocCount >> 10, s);
191            for (int i = 0; i < s; i+= 4096) {
192                *((long long*)&ptr[i]) = allocCount + i;
193            }
194            usleep(10 * 1000);
195            gptr = ptr;
196            //cout << "total alloc: " << allocCount / (1<<20)<< " adj: " << argv[2]<< endl;;
197            //cout << "ptr: " << (long long)(void*)ptr << endl;;
198            p.signal();
199            allocCount += s;
200        }
201    } else {
202        cout << "parent:" << argc << endl;
203
204        write_oomadj_to_lmkd(-1000);
205        for (int i = 1000; i >= 0; i -= 100) {
206            auto pipes = Pipe::createPipePair();
207            char arg[16];
208            snprintf(arg, sizeof(arg), "%d", i);
209            createProcess(std::move(std::get<1>(pipes)), argv[0], arg);
210            Pipe &p = std::get<0>(pipes);
211
212            size_t t = 0;
213            while (1) {
214                //;cout << getpid() << ":" << "parent signal" << endl;
215                p.signal();
216                if (p.wait_ret_error()) {
217                    int status;
218                    waitpid(0, &status, 0);
219                    break;
220                }
221                t += s;
222            }
223            cout << "adj: " << i << " sz: " << t / (1 << 20) << endl;
224        }
225    }
226    return 0;
227}
228