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
31char *global_buf = NULL;
32size_t buf_size = 0;
33
34// Stores a chunk of fake data that will give a target compression ratio
35char *fake_data;
36
37// Dummy global that forces compiler to perform the read
38volatile char dummy;
39
40struct 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.
48static 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.
63static 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.
77static 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.
83static 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
88int 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