/*
 * Test program to ensure the mlock memory limit is enforced for io_setup,
 * which pins kernel memory for the AIO event completion ring.
 *
 * Compile with:  gcc -o ioctx_alloc ioctx_alloc.c -laio -lcap
 *
 * Author: Jeff Moyer <jmoyer@redhat.com>
 * Copyright: Red Hat, Inc.  (2009)
 * License: GPLv2
 */
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <errno.h>
#include <fcntl.h>
#include <assert.h>
#include <sys/capability.h>
#include <sys/time.h>
#include <sys/resource.h>

#include <libaio.h>

int
main(int argc, char **argv)
{
	int cnt = 0;
	int ret;
	long nr_events = 0;
	cap_t caps;
	cap_value_t cap_ipc_memlock = CAP_IPC_LOCK;
	struct rlimit rlim_memlock;
	io_context_t aioctx;

	if (getrlimit(RLIMIT_MEMLOCK, &rlim_memlock) == -1) {
		perror("getrlimit");
		exit(1);
	}

	if (argc != 1) {
		if (argc < 2) {
			printf("usage: ioctx_alloc [memlock limit in MB]\n");
			exit(1);
		}

		rlim_memlock.rlim_max = atoi(argv[1]) * 1024 * 1024;
		rlim_memlock.rlim_cur = rlim_memlock.rlim_max;
		if (setrlimit(RLIMIT_MEMLOCK, &rlim_memlock) == -1) {
			perror("setrlimit");
			exit(1);
		}
	}
	printf("Maximum permissible locked memory: %ld bytes\n",
	       rlim_memlock.rlim_max);

	/* drop CAP_IPC_MEMLOCK so the memlock rlimit is enforced */
	caps = cap_get_proc();
	assert(caps);
	if (cap_set_flag(caps, CAP_EFFECTIVE,
			 1, &cap_ipc_memlock, CAP_CLEAR) == -1) {
		perror("cap_set_flag");
		exit(1);
	}
	if (cap_set_proc(caps) == -1) {
		perror("cap_set_proc");
		exit(1);
	}

	/*
	 * First, let's try the sane case of 1 io context, and see how
	 * many events we can setup.
	 */
	nr_events = 1;
	while (1) {
		memset(&aioctx, 0, sizeof(aioctx));
		ret = io_setup(nr_events, &aioctx);
		if (ret == 0) {
			io_destroy(aioctx);
			nr_events <<= 1;
			continue;
		}
		nr_events >>= 1;
		break;
	}
	while (1) {
		memset(&aioctx, 0, sizeof(aioctx));
		ret = io_setup(nr_events, &aioctx);
		if (ret == 0) {
			io_destroy(aioctx);
			nr_events++;
			continue;
		}
		nr_events--;
		break;
	}
	printf("successfully reserved %d events in a single io context\n",
	       nr_events);
	/*
	 * Now let's see how many io cotexts we can allocate.  The kernel
	 * should pin 1 page per io_setup call.
	 */
	do {
		memset(&aioctx, 0, sizeof(aioctx));
		ret = io_setup(1, &aioctx);
		cnt++;
	} while (ret == 0);
	cnt--;

	printf("Successfully allocated %d IO contexts\n", cnt);
	exit(0);
}
/*
 * Local variables:
 *   c-basic-offset: 8
 *   compile-command: "gcc -o ioctx_alloc ioctx_alloc.c -laio -lcap"
 * End:
 */