// SPDX-License-Identifier: GPL-2.0
/*
 * RMAP functional tests
 *
 * Author(s): Wei Yang <richard.weiyang@gmail.com>
 */

#include "kselftest_harness.h"
#include <strings.h>
#include <pthread.h>
#include <numa.h>
#include <numaif.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include <sys/types.h>
#include <signal.h>
#include <time.h>
#include <sys/sem.h>
#include <unistd.h>
#include <fcntl.h>

#include "vm_util.h"

#define TOTAL_LEVEL 5
#define MAX_CHILDREN 3

#define FAIL_ON_CHECK	(1 << 0)
#define FAIL_ON_WORK	(1 << 1)

struct sembuf sem_wait = {0, -1, 0};
struct sembuf sem_signal = {0, 1, 0};

enum backend_type {
	ANON,
	SHM,
	NORM_FILE,
};

#define PREFIX "kst_rmap"
#define MAX_FILENAME_LEN 256
const char *suffixes[] = {
	"",
	"_shm",
	"_file",
};

struct global_data;
typedef int (*work_fn)(struct global_data *data);
typedef int (*check_fn)(struct global_data *data);
typedef void (*prepare_fn)(struct global_data *data);

struct global_data {
	int worker_level;

	int semid;
	int pipefd[2];

	unsigned int mapsize;
	unsigned int rand_seed;
	char *region;

	prepare_fn do_prepare;
	work_fn do_work;
	check_fn do_check;

	enum backend_type backend;
	char filename[MAX_FILENAME_LEN];

	unsigned long *expected_pfn;
};

/*
 * Create a process tree with TOTAL_LEVEL height and at most MAX_CHILDREN
 * children for each.
 *
 * It will randomly select one process as 'worker' process which will
 * 'do_work' until all processes are created. And all other processes will
 * wait until 'worker' finish its work.
 */
void propagate_children(struct __test_metadata *_metadata, struct global_data *data)
{
	pid_t root_pid, pid;
	unsigned int num_child;
	int status;
	int ret = 0;
	int curr_child, worker_child;
	int curr_level = 1;
	bool is_worker = true;

	root_pid = getpid();
repeat:
	num_child = rand_r(&data->rand_seed) % MAX_CHILDREN + 1;
	worker_child = is_worker ? rand_r(&data->rand_seed) % num_child : -1;

	for (curr_child = 0; curr_child < num_child; curr_child++) {
		pid = fork();

		if (pid < 0) {
			perror("Error: fork\n");
		} else if (pid == 0) {
			curr_level++;

			if (curr_child != worker_child)
				is_worker = false;

			if (curr_level == TOTAL_LEVEL)
				break;

			data->rand_seed += curr_child;
			goto repeat;
		}
	}

	if (data->do_prepare)
		data->do_prepare(data);

	close(data->pipefd[1]);

	if (is_worker && curr_level == data->worker_level) {
		/* This is the worker process, first wait last process created */
		char buf;

		while (read(data->pipefd[0], &buf, 1) > 0)
			;

		if (data->do_work)
			ret = data->do_work(data);

		/* Kick others */
		semctl(data->semid, 0, IPC_RMID);
	} else {
		/* Wait worker finish */
		semop(data->semid, &sem_wait, 1);
		if (data->do_check)
			ret = data->do_check(data);
	}

	/* Wait all child to quit */
	while (wait(&status) > 0) {
		if (WIFEXITED(status))
			ret |= WEXITSTATUS(status);
	}

	if (getpid() == root_pid) {
		if (ret & FAIL_ON_WORK)
			SKIP(return, "Failed in worker");

		ASSERT_EQ(ret, 0);
	} else {
		exit(ret);
	}
}

FIXTURE(migrate)
{
	struct global_data data;
};

FIXTURE_SETUP(migrate)
{
	struct global_data *data = &self->data;

	if (numa_available() < 0)
		SKIP(return, "NUMA not available");
	if (numa_bitmask_weight(numa_all_nodes_ptr) <= 1)
		SKIP(return, "Not enough NUMA nodes available");

	data->mapsize = getpagesize();

	data->expected_pfn = mmap(0, sizeof(unsigned long),
				PROT_READ | PROT_WRITE,
				MAP_SHARED | MAP_ANONYMOUS, -1, 0);
	ASSERT_NE(data->expected_pfn, MAP_FAILED);

	/* Prepare semaphore */
	data->semid = semget(IPC_PRIVATE, 1, 0666 | IPC_CREAT);
	ASSERT_NE(data->semid, -1);
	ASSERT_NE(semctl(data->semid, 0, SETVAL, 0), -1);

	/* Prepare pipe */
	ASSERT_NE(pipe(data->pipefd), -1);

	data->rand_seed = time(NULL);
	srand(data->rand_seed);

	data->worker_level = rand() % TOTAL_LEVEL + 1;

	data->do_prepare = NULL;
	data->do_work = NULL;
	data->do_check = NULL;

	data->backend = ANON;
};

FIXTURE_TEARDOWN(migrate)
{
	struct global_data *data = &self->data;

	if (data->region != MAP_FAILED)
		munmap(data->region, data->mapsize);
	data->region = MAP_FAILED;
	if (data->expected_pfn != MAP_FAILED)
		munmap(data->expected_pfn, sizeof(unsigned long));
	data->expected_pfn = MAP_FAILED;
	semctl(data->semid, 0, IPC_RMID);
	data->semid = -1;

	close(data->pipefd[0]);

	switch (data->backend) {
	case ANON:
		break;
	case SHM:
		shm_unlink(data->filename);
		break;
	case NORM_FILE:
		unlink(data->filename);
		break;
	}
}

void access_region(struct global_data *data)
{
	/*
	 * Force read "region" to make sure page fault in.
	 */
	FORCE_READ(*data->region);
}

int try_to_move_page(char *region)
{
	int ret;
	int node;
	int status = 0;
	int failures = 0;

	ret = move_pages(0, 1, (void **)&region, NULL, &status, MPOL_MF_MOVE_ALL);
	if (ret != 0) {
		perror("Failed to get original numa");
		return FAIL_ON_WORK;
	}

	/* Pick up a different target node */
	for (node = 0; node <= numa_max_node(); node++) {
		if (numa_bitmask_isbitset(numa_all_nodes_ptr, node) && node != status)
			break;
	}

	if (node > numa_max_node()) {
		ksft_print_msg("Couldn't find available numa node for testing\n");
		return FAIL_ON_WORK;
	}

	while (1) {
		ret = move_pages(0, 1, (void **)&region, &node, &status, MPOL_MF_MOVE_ALL);

		/* migrate successfully */
		if (!ret)
			break;

		/* error happened */
		if (ret < 0) {
			ksft_perror("Failed to move pages");
			return FAIL_ON_WORK;
		}

		/* migration is best effort; try again */
		if (++failures >= 100)
			return FAIL_ON_WORK;
	}

	return 0;
}

int move_region(struct global_data *data)
{
	int ret;
	int pagemap_fd;

	ret = try_to_move_page(data->region);
	if (ret != 0)
		return ret;

	pagemap_fd = open("/proc/self/pagemap", O_RDONLY);
	if (pagemap_fd == -1)
		return FAIL_ON_WORK;
	*data->expected_pfn = pagemap_get_pfn(pagemap_fd, data->region);

	return 0;
}

int has_same_pfn(struct global_data *data)
{
	unsigned long pfn;
	int pagemap_fd;

	if (data->region == MAP_FAILED)
		return 0;

	pagemap_fd = open("/proc/self/pagemap", O_RDONLY);
	if (pagemap_fd == -1)
		return FAIL_ON_CHECK;

	pfn = pagemap_get_pfn(pagemap_fd, data->region);
	if (pfn != *data->expected_pfn)
		return FAIL_ON_CHECK;

	return 0;
}

TEST_F(migrate, anon)
{
	struct global_data *data = &self->data;

	/* Map an area and fault in */
	data->region = mmap(0, data->mapsize, PROT_READ | PROT_WRITE,
				MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
	ASSERT_NE(data->region, MAP_FAILED);
	memset(data->region, 0xcf, data->mapsize);

	data->do_prepare = access_region;
	data->do_work = move_region;
	data->do_check = has_same_pfn;

	propagate_children(_metadata, data);
}

TEST_F(migrate, shm)
{
	int shm_fd;
	struct global_data *data = &self->data;

	snprintf(data->filename, MAX_FILENAME_LEN, "%s%s", PREFIX, suffixes[SHM]);
	shm_fd = shm_open(data->filename, O_CREAT | O_RDWR, 0666);
	ASSERT_NE(shm_fd, -1);
	ftruncate(shm_fd, data->mapsize);
	data->backend = SHM;

	/* Map a shared area and fault in */
	data->region = mmap(0, data->mapsize, PROT_READ | PROT_WRITE,
				MAP_SHARED, shm_fd, 0);
	ASSERT_NE(data->region, MAP_FAILED);
	memset(data->region, 0xcf, data->mapsize);
	close(shm_fd);

	data->do_prepare = access_region;
	data->do_work = move_region;
	data->do_check = has_same_pfn;

	propagate_children(_metadata, data);
}

TEST_F(migrate, file)
{
	int fd;
	struct global_data *data = &self->data;

	snprintf(data->filename, MAX_FILENAME_LEN, "%s%s", PREFIX, suffixes[NORM_FILE]);
	fd = open(data->filename, O_CREAT | O_RDWR | O_EXCL, 0666);
	ASSERT_NE(fd, -1);
	ftruncate(fd, data->mapsize);
	data->backend = NORM_FILE;

	/* Map a shared area and fault in */
	data->region = mmap(0, data->mapsize, PROT_READ | PROT_WRITE,
				MAP_SHARED, fd, 0);
	ASSERT_NE(data->region, MAP_FAILED);
	memset(data->region, 0xcf, data->mapsize);
	close(fd);

	data->do_prepare = access_region;
	data->do_work = move_region;
	data->do_check = has_same_pfn;

	propagate_children(_metadata, data);
}

void prepare_local_region(struct global_data *data)
{
	/* Allocate range and set the same data */
	data->region = mmap(NULL, data->mapsize, PROT_READ|PROT_WRITE,
			   MAP_PRIVATE|MAP_ANON, -1, 0);
	if (data->region == MAP_FAILED)
		return;

	memset(data->region, 0xcf, data->mapsize);
}

int merge_and_migrate(struct global_data *data)
{
	int pagemap_fd;
	int ret = 0;

	if (data->region == MAP_FAILED)
		return FAIL_ON_WORK;

	if (ksm_start() < 0)
		return FAIL_ON_WORK;

	ret = try_to_move_page(data->region);

	pagemap_fd = open("/proc/self/pagemap", O_RDONLY);
	if (pagemap_fd == -1)
		return FAIL_ON_WORK;
	*data->expected_pfn = pagemap_get_pfn(pagemap_fd, data->region);

	return ret;
}

TEST_F(migrate, ksm)
{
	int ret;
	struct global_data *data = &self->data;

	if (ksm_stop() < 0)
		SKIP(return, "accessing \"/sys/kernel/mm/ksm/run\") failed");
	if (ksm_get_full_scans() < 0)
		SKIP(return, "accessing \"/sys/kernel/mm/ksm/full_scan\") failed");

	ret = prctl(PR_SET_MEMORY_MERGE, 1, 0, 0, 0);
	if (ret < 0 && errno == EINVAL)
		SKIP(return, "PR_SET_MEMORY_MERGE not supported");
	else if (ret)
		ksft_exit_fail_perror("PR_SET_MEMORY_MERGE=1 failed");

	data->do_prepare = prepare_local_region;
	data->do_work = merge_and_migrate;
	data->do_check = has_same_pfn;

	propagate_children(_metadata, data);
}

TEST_HARNESS_MAIN