mempressure_test.c

/*
 * mempressure shrinker test
 *
 * Copyright 2012 Linaro Ltd.
 *		  Anton Vorontsov <anton.vorontsov@linaro.org>
 *
 * It is pretty simple: we create two threads, the first one constantly
 * tries to allocate memory (more than we physically have), the second
 * thread listens to the kernel shrinker notifications and frees asked
 * amount of chunks. When we allocate more than available RAM, the two
 * threads start to fight. Idially, we should not OOM (but if we reclaim
 * slower than we allocate, things might OOM). Also, ideally we should not
 * grow swap too much.
 *
 * The test accepts no arguments, so you can just run it and observe the
 * output and memory usage (e.g. 'watch -n 0.2 free -m'). Upon ctrl+c, the
 * test prints total amount of bytes we helped to reclaim.
 *
 * Compile with -pthread.
 *
 * This program is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 as published
 * by the Free Software Foundation.
 */

#define _GNU_SOURCE
#include <stdio.h>
#include <stdlib.h>
#include <stdint.h>
#include <stdbool.h>
#include <unistd.h>
#include <string.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <pthread.h>
#include <signal.h>
#include <errno.h>
#include <sys/eventfd.h>
#include <sys/sysinfo.h>

#define CG			"/sys/fs/cgroup/mempressure"
#define CG_EVENT_CONTROL	(CG "/cgroup.event_control")
#define CG_SHRINKER		(CG "/mempressure.shrinker")

#define CHUNK_SIZE (1 * 1024 * 1024)

static size_t num_chunks;

static void **chunks;
static pthread_mutex_t *locks;
static int efd;
static int sfd;

static inline void pabort(bool f, int code, const char *str)
{
	if (!f)
		return;
	perror(str);
	printf("(%d)\n", code);
	abort();
}

static void init_shrinker(void)
{
	int cfd;
	int ret;
	char *str;

	cfd = open(CG_EVENT_CONTROL, O_WRONLY);
	pabort(cfd < 0, cfd, CG_EVENT_CONTROL);

	sfd = open(CG_SHRINKER, O_RDWR);
	pabort(sfd < 0, sfd, CG_SHRINKER);

	efd = eventfd(0, 0);
	pabort(efd < 0, efd, "eventfd()");

	ret = asprintf(&str, "%d %d %d\n", efd, sfd, CHUNK_SIZE);
	printf("%s\n", str);
	pabort(ret == -1, ret, "control string");

	ret = write(cfd, str, ret + 1);
	pabort(ret == -1, ret, "write() to event_control");
}

static void add_reclaimable(int chunks)
{
	int ret;
	char *str;

	ret = asprintf(&str, "%d %d\n", efd, CHUNK_SIZE);
	pabort(ret == -1, ret, "add_reclaimable, asprintf");

	ret = write(sfd, str, ret + 1);
	pabort(ret <= 0, ret, "add_reclaimable, write");
}

static int chunks_to_reclaim(void)
{
	uint64_t n = 0;
	int ret;

	ret = read(efd, &n, sizeof(n));
	pabort(ret <= 0, ret, "read() from eventfd");

	printf("%d chunks to reclaim\n", (int)n);

	return n;
}

static unsigned int reclaimed;

static void print_stats(int signum)
{
	printf("\nTOTAL: helped to reclaim %d chunks (%d MB)\n",
	       reclaimed, reclaimed * CHUNK_SIZE / 1024 / 1024);
	exit(0);
}

static void *shrinker_thr_fn(void *arg)
{
	puts("shrinker thread started");

	sigaction(SIGINT, &(struct sigaction){.sa_handler = print_stats}, NULL);

	while (1) {
		unsigned int i = 0;
		int n;

		n = chunks_to_reclaim();

		reclaimed += n;

		while (n) {
			pthread_mutex_lock(&locks[i]);
			if (chunks[i]) {
				free(chunks[i]);
				chunks[i] = NULL;
				n--;
			}
			pthread_mutex_unlock(&locks[i]);

			i = (i + 1) % num_chunks;
		}
	}
	return NULL;
}

static void consume_memory(void)
{
	unsigned int i = 0;
	unsigned int j = 0;

	puts("consuming memory...");

	while (1) {
		pthread_mutex_lock(&locks[i]);
		if (!chunks[i]) {
			chunks[i] = malloc(CHUNK_SIZE);
			pabort(!chunks[i], 0, "chunks alloc failed");
			memset(chunks[i], 0, CHUNK_SIZE);
			j++;
		}
		pthread_mutex_unlock(&locks[i]);

		if (j >= num_chunks / 10) {
			add_reclaimable(num_chunks / 10);
			printf("added %d reclaimable chunks\n", j);
			j = 0;
		}

		i = (i + 1) % num_chunks;
	}
}

int main(int argc, char *argv[])
{
	int ret;
	int i;
	pthread_t shrinker_thr;
	struct sysinfo si;

	ret = sysinfo(&si);
	pabort(ret != 0, ret, "sysinfo()");

	num_chunks = (si.totalram + si.totalswap) * si.mem_unit / 1024 / 1024;

	chunks = malloc(sizeof(*chunks) * num_chunks);
	locks = malloc(sizeof(*locks) * num_chunks);
	pabort(!chunks || !locks, ENOMEM, NULL);

	init_shrinker();

	for (i = 0; i < num_chunks; i++) {
		ret = pthread_mutex_init(&locks[i], NULL);
		pabort(ret != 0, ret, "pthread_mutex_init");
	}

	ret = pthread_create(&shrinker_thr, NULL, shrinker_thr_fn, NULL);
	pabort(ret != 0, ret, "pthread_create(shrinker)");

	consume_memory();

	ret = pthread_join(shrinker_thr, NULL);
	pabort(ret != 0, ret, "pthread_join(shrinker)");

	return 0;
}