| // Copyright (c) 2013 The Chromium OS Authors. All rights reserved. |
| // Use of this source code is governed by a BSD-style license that can be |
| // found in the LICENSE file. |
| |
| #define CMD_POKE 1 |
| #define CMD_BALLOON 2 |
| #define CMD_EXIT 3 |
| |
| #define TOUCH_LIMIT 1000 |
| #define WRITE_MOD 10 |
| |
| // Allocate memory in 1 MiB chunks |
| #define CHUNK_SIZE (1 << 20) |
| |
| #include <fcntl.h> |
| #include <stdbool.h> |
| #include <stdint.h> |
| #include <stdio.h> |
| #include <stdlib.h> |
| #include <string.h> |
| #include <time.h> |
| #include <unistd.h> |
| #include <sys/resource.h> |
| #include <sys/socket.h> |
| #include <sys/stat.h> |
| #include <sys/time.h> |
| #include <sys/types.h> |
| #include <sys/un.h> |
| |
| // Hog's main buffer |
| char *global_buf = NULL; |
| size_t buf_size = 0; |
| |
| // Stores a chunk of fake data that will give a target compression ratio |
| char *fake_data; |
| |
| // Dummy global that forces compiler to perform the read |
| volatile char dummy; |
| |
| struct PokeResult { |
| uint64_t real_time; |
| uint64_t user_time; |
| uint64_t sys_time; |
| uint64_t faults; |
| } __attribute__((packed)); |
| |
| // Reads and writes random pages in global_buf. |
| static void TouchMemory() { |
| for (int i = 0; i < TOUCH_LIMIT; i++) { |
| unsigned int index = (unsigned int)rand(); |
| |
| // Randomly do a write instead of a read. |
| if (rand() % WRITE_MOD == 0) { |
| global_buf[index % buf_size] = 0x00; |
| } else { |
| dummy = global_buf[index % buf_size]; |
| } |
| } |
| } |
| |
| // Allocates memory and copies fake data in to ensure there's no copy-on-write |
| // business going on. |
| static void BalloonMemory(size_t balloon_size) { |
| size_t new_buf_size = buf_size + balloon_size * CHUNK_SIZE; |
| global_buf = realloc(global_buf, new_buf_size); |
| |
| // Copy fake data into every chunk that we allocate. |
| for (unsigned int chunk = 0; chunk < balloon_size; chunk++) { |
| char *new_chunk = global_buf + buf_size + chunk * CHUNK_SIZE; |
| memcpy(new_chunk, fake_data, CHUNK_SIZE); |
| } |
| |
| buf_size = new_buf_size; |
| } |
| |
| // Calculates the difference between two timespecs in milliseconds. |
| static uint64_t DiffTimespec(struct timespec start, struct timespec end) { |
| return (end.tv_sec - start.tv_sec) * 1000 + |
| (end.tv_nsec - start.tv_nsec) / 1000000; |
| } |
| |
| // Calculates the difference between two timevals in milliseconds. |
| static uint64_t DiffTimeval(struct timeval start, struct timeval end) { |
| return (end.tv_sec - start.tv_sec) * 1000 + |
| (end.tv_usec - start.tv_usec) / 1000; |
| } |
| |
| int main(int argc, char *argv[]) { |
| int sockfd; |
| struct sockaddr_un test_sock_addr; |
| int compression_factor = 3; |
| int random_fd = open("/dev/urandom", O_RDONLY); |
| |
| if (argc < 2) { |
| fprintf(stderr, "Usage: %s SOCKETNAME COMPRESSION_FACTOR\n", argv[0]); |
| return 1; |
| } |
| |
| if (argc == 3) { |
| compression_factor = atoi(argv[2]); |
| } |
| |
| srand(getpid()); |
| |
| test_sock_addr.sun_family = AF_UNIX; |
| strncpy(test_sock_addr.sun_path, argv[1], strlen(argv[1]) + 1); |
| |
| sockfd = socket(AF_UNIX, SOCK_STREAM, 0); |
| |
| if (sockfd < 0) { |
| perror("could not open socket"); |
| return 1; |
| } |
| |
| // Unlink any existing socket with this name. |
| struct stat file_stat; |
| if (stat(argv[1], &file_stat) == 0) { |
| if (S_ISSOCK(file_stat.st_mode)) { |
| unlink(argv[1]); |
| } else { |
| fprintf(stderr, |
| "there is a file with the given socket name already; aborting\n"); |
| return 1; |
| } |
| } |
| |
| if (bind(sockfd, (struct sockaddr *)&test_sock_addr, sizeof test_sock_addr)) { |
| perror("could not bind to socket"); |
| return 1; |
| } |
| |
| if (listen(sockfd, 1)) { |
| perror("could not listen to socket"); |
| return 1; |
| } |
| |
| int connfd; |
| if ((connfd = accept(sockfd, NULL, NULL)) < 0) { |
| perror("could not accept connection"); |
| return 1; |
| } |
| |
| // Fill fake_data with fake data so that it compresses to roughly the desired |
| // compression factor. Random data should be uncompressible, while long |
| // sequences of ones are highly compressible. |
| fake_data = malloc(CHUNK_SIZE); |
| read(random_fd, fake_data, CHUNK_SIZE / compression_factor); |
| |
| memset(fake_data + CHUNK_SIZE / compression_factor, 1, |
| CHUNK_SIZE - (CHUNK_SIZE / compression_factor)); |
| |
| // Allocate one chunk worth of data to start with. |
| BalloonMemory(1); |
| |
| while (true) { |
| uint32_t command; |
| uint32_t balloon_size; |
| struct sockaddr src_addr; |
| struct timespec time_start; |
| struct timespec time_end; |
| struct rusage usage_start; |
| struct rusage usage_end; |
| struct PokeResult result; |
| |
| ssize_t bytes_read = recv(connfd, &command, sizeof(command), 0); |
| |
| if (bytes_read < 0) { |
| perror("error while reading from socket"); |
| return 1; |
| } else if (bytes_read == 0) { |
| // Remote socket closed early; clean up this hog. |
| fprintf(stderr, "read 0 bytes from socket; terminating\n"); |
| return 0; |
| } else if (bytes_read != sizeof(command)) { |
| fprintf(stderr, "read %li bytes (expected %lu); aborting\n", |
| bytes_read, sizeof(command)); |
| return 1; |
| } |
| |
| switch(command) { |
| case CMD_POKE: |
| // Touch pages of memory while monitoring time and resource usage. |
| getrusage(RUSAGE_SELF, &usage_start); |
| clock_gettime(CLOCK_REALTIME, &time_start); |
| |
| TouchMemory(); |
| |
| clock_gettime(CLOCK_REALTIME, &time_end); |
| getrusage(RUSAGE_SELF, &usage_end); |
| |
| // Send stats back to monitor script. |
| result.real_time = DiffTimespec(time_start, time_end); |
| result.user_time = DiffTimeval(usage_start.ru_utime, |
| usage_end.ru_utime); |
| result.sys_time = DiffTimeval(usage_start.ru_stime, |
| usage_end.ru_stime); |
| result.faults = usage_end.ru_majflt - usage_start.ru_majflt; |
| |
| send(connfd, &result, sizeof(result), 0); |
| break; |
| case CMD_BALLOON: |
| bytes_read = recv(connfd, &balloon_size, sizeof(balloon_size), 0); |
| |
| if (bytes_read < 0) { |
| perror("error while reading from socket"); |
| return 1; |
| } else if (bytes_read == 0) { |
| fprintf(stderr, "read 0 bytes from socket; terminating\n"); |
| return 0; |
| } |
| |
| BalloonMemory(balloon_size); |
| send(connfd, &balloon_size, sizeof(balloon_size), 0); |
| break; |
| case CMD_EXIT: |
| fprintf(stderr, "exiting\n"); |
| return 0; |
| default: |
| fprintf(stderr, "unexpected command: %d\n", command); |
| } |
| } |
| |
| return 0; |
| } |