1 // Copyright (c) 2013 The Chromium OS Authors. All rights reserved.
2 // Use of this source code is governed by a BSD-style license that can be
3 // found in the LICENSE file.
4
5 #define CMD_POKE 1
6 #define CMD_BALLOON 2
7 #define CMD_EXIT 3
8
9 #define TOUCH_LIMIT 1000
10 #define WRITE_MOD 10
11
12 // Allocate memory in 1 MiB chunks
13 #define CHUNK_SIZE (1 << 20)
14
15 #include <fcntl.h>
16 #include <stdbool.h>
17 #include <stdint.h>
18 #include <stdio.h>
19 #include <stdlib.h>
20 #include <string.h>
21 #include <time.h>
22 #include <unistd.h>
23 #include <sys/resource.h>
24 #include <sys/socket.h>
25 #include <sys/stat.h>
26 #include <sys/time.h>
27 #include <sys/types.h>
28 #include <sys/un.h>
29
30 // Hog's main buffer
31 char *global_buf = NULL;
32 size_t buf_size = 0;
33
34 // Stores a chunk of fake data that will give a target compression ratio
35 char *fake_data;
36
37 // Dummy global that forces compiler to perform the read
38 volatile char dummy;
39
40 struct PokeResult {
41 uint64_t real_time;
42 uint64_t user_time;
43 uint64_t sys_time;
44 uint64_t faults;
45 } __attribute__((packed));
46
47 // Reads and writes random pages in global_buf.
TouchMemory()48 static void TouchMemory() {
49 for (int i = 0; i < TOUCH_LIMIT; i++) {
50 unsigned int index = (unsigned int)rand();
51
52 // Randomly do a write instead of a read.
53 if (rand() % WRITE_MOD == 0) {
54 global_buf[index % buf_size] = 0x00;
55 } else {
56 dummy = global_buf[index % buf_size];
57 }
58 }
59 }
60
61 // Allocates memory and copies fake data in to ensure there's no copy-on-write
62 // business going on.
BalloonMemory(size_t balloon_size)63 static void BalloonMemory(size_t balloon_size) {
64 size_t new_buf_size = buf_size + balloon_size * CHUNK_SIZE;
65 global_buf = realloc(global_buf, new_buf_size);
66
67 // Copy fake data into every chunk that we allocate.
68 for (unsigned int chunk = 0; chunk < balloon_size; chunk++) {
69 char *new_chunk = global_buf + buf_size + chunk * CHUNK_SIZE;
70 memcpy(new_chunk, fake_data, CHUNK_SIZE);
71 }
72
73 buf_size = new_buf_size;
74 }
75
76 // Calculates the difference between two timespecs in milliseconds.
DiffTimespec(struct timespec start,struct timespec end)77 static uint64_t DiffTimespec(struct timespec start, struct timespec end) {
78 return (end.tv_sec - start.tv_sec) * 1000 +
79 (end.tv_nsec - start.tv_nsec) / 1000000;
80 }
81
82 // Calculates the difference between two timevals in milliseconds.
DiffTimeval(struct timeval start,struct timeval end)83 static uint64_t DiffTimeval(struct timeval start, struct timeval end) {
84 return (end.tv_sec - start.tv_sec) * 1000 +
85 (end.tv_usec - start.tv_usec) / 1000;
86 }
87
main(int argc,char * argv[])88 int main(int argc, char *argv[]) {
89 int sockfd;
90 struct sockaddr_un test_sock_addr;
91 int compression_factor = 3;
92 int random_fd = open("/dev/urandom", O_RDONLY);
93
94 if (argc < 2) {
95 fprintf(stderr, "Usage: %s SOCKETNAME COMPRESSION_FACTOR\n", argv[0]);
96 return 1;
97 }
98
99 if (argc == 3) {
100 compression_factor = atoi(argv[2]);
101 }
102
103 srand(getpid());
104
105 test_sock_addr.sun_family = AF_UNIX;
106 strncpy(test_sock_addr.sun_path, argv[1], strlen(argv[1]) + 1);
107
108 sockfd = socket(AF_UNIX, SOCK_STREAM, 0);
109
110 if (sockfd < 0) {
111 perror("could not open socket");
112 return 1;
113 }
114
115 // Unlink any existing socket with this name.
116 struct stat file_stat;
117 if (stat(argv[1], &file_stat) == 0) {
118 if (S_ISSOCK(file_stat.st_mode)) {
119 unlink(argv[1]);
120 } else {
121 fprintf(stderr,
122 "there is a file with the given socket name already; aborting\n");
123 return 1;
124 }
125 }
126
127 if (bind(sockfd, (struct sockaddr *)&test_sock_addr, sizeof test_sock_addr)) {
128 perror("could not bind to socket");
129 return 1;
130 }
131
132 if (listen(sockfd, 1)) {
133 perror("could not listen to socket");
134 return 1;
135 }
136
137 int connfd;
138 if ((connfd = accept(sockfd, NULL, NULL)) < 0) {
139 perror("could not accept connection");
140 return 1;
141 }
142
143 // Fill fake_data with fake data so that it compresses to roughly the desired
144 // compression factor. Random data should be uncompressible, while long
145 // sequences of ones are highly compressible.
146 fake_data = malloc(CHUNK_SIZE);
147 read(random_fd, fake_data, CHUNK_SIZE / compression_factor);
148
149 memset(fake_data + CHUNK_SIZE / compression_factor, 1,
150 CHUNK_SIZE - (CHUNK_SIZE / compression_factor));
151
152 // Allocate one chunk worth of data to start with.
153 BalloonMemory(1);
154
155 while (true) {
156 uint32_t command;
157 uint32_t balloon_size;
158 struct sockaddr src_addr;
159 struct timespec time_start;
160 struct timespec time_end;
161 struct rusage usage_start;
162 struct rusage usage_end;
163 struct PokeResult result;
164
165 ssize_t bytes_read = recv(connfd, &command, sizeof(command), 0);
166
167 if (bytes_read < 0) {
168 perror("error while reading from socket");
169 return 1;
170 } else if (bytes_read == 0) {
171 // Remote socket closed early; clean up this hog.
172 fprintf(stderr, "read 0 bytes from socket; terminating\n");
173 return 0;
174 } else if (bytes_read != sizeof(command)) {
175 fprintf(stderr, "read %li bytes (expected %lu); aborting\n",
176 bytes_read, sizeof(command));
177 return 1;
178 }
179
180 switch(command) {
181 case CMD_POKE:
182 // Touch pages of memory while monitoring time and resource usage.
183 getrusage(RUSAGE_SELF, &usage_start);
184 clock_gettime(CLOCK_REALTIME, &time_start);
185
186 TouchMemory();
187
188 clock_gettime(CLOCK_REALTIME, &time_end);
189 getrusage(RUSAGE_SELF, &usage_end);
190
191 // Send stats back to monitor script.
192 result.real_time = DiffTimespec(time_start, time_end);
193 result.user_time = DiffTimeval(usage_start.ru_utime,
194 usage_end.ru_utime);
195 result.sys_time = DiffTimeval(usage_start.ru_stime,
196 usage_end.ru_stime);
197 result.faults = usage_end.ru_majflt - usage_start.ru_majflt;
198
199 send(connfd, &result, sizeof(result), 0);
200 break;
201 case CMD_BALLOON:
202 bytes_read = recv(connfd, &balloon_size, sizeof(balloon_size), 0);
203
204 if (bytes_read < 0) {
205 perror("error while reading from socket");
206 return 1;
207 } else if (bytes_read == 0) {
208 fprintf(stderr, "read 0 bytes from socket; terminating\n");
209 return 0;
210 }
211
212 BalloonMemory(balloon_size);
213 send(connfd, &balloon_size, sizeof(balloon_size), 0);
214 break;
215 case CMD_EXIT:
216 fprintf(stderr, "exiting\n");
217 return 0;
218 default:
219 fprintf(stderr, "unexpected command: %d\n", command);
220 }
221 }
222
223 return 0;
224 }
225