// SPDX-License-Identifier: GPL-2.0

#include <linux/compiler.h>
#include <stdio.h>
#include <stdlib.h>
#include <sys/prctl.h>
#include <sys/wait.h>
#include <unistd.h>

#include "../tests.h"

static int loops = 100;
static char buf;
int context_switch_loop_work = 1234;

#define write_block(fd) \
	do { \
		if (write(fd, &buf, 1) <= 0) \
			return 1; \
	} while (0)

#define read_block(fd) \
	do { \
		if (read(fd, &buf, 1) <= 0) \
			return 1; \
	} while (0)

/* Not static to avoid LTO clobbering the function name */
int context_switch_loop_proc1(int in_fd, int out_fd);
int context_switch_loop_proc1(int in_fd, int out_fd)
{
	for (int i = 0; i < loops; i++) {
		read_block(in_fd);
		context_switch_loop_work += i * 3;
		write_block(out_fd);
	}
	return 0;
}

int context_switch_loop_proc2(int in_fd, int out_fd);
int context_switch_loop_proc2(int in_fd, int out_fd)
{
	for (int i = 0; i < loops; i++) {
		write_block(out_fd);
		context_switch_loop_work += i * 7;
		read_block(in_fd);
	}
	return 0;
}

/*
 * Launches two processes that take turns to execute a multiplication N times
 */
static int context_switch_loop(int argc, const char **argv)
{
	int a_to_b[2], b_to_a[2];
	pid_t proc1_pid;
	int status;
	int ret;

	if (argc > 0) {
		loops = atoi(argv[0]);
		if (loops < 0) {
			fprintf(stderr, "Invalid number of loops: %s\n", argv[0]);
			return 1;
		}
	}

	if (pipe(a_to_b) || pipe(b_to_a)) {
		perror("Pipe error");
		return 1;
	}

	proc1_pid = fork();
	if (proc1_pid < 0) {
		perror("Fork error");
		return 1;
	}

	if (!proc1_pid) {
		close(a_to_b[0]);
		close(b_to_a[1]);
		prctl(PR_SET_NAME, "proc1", 0, 0, 0);
		ret = context_switch_loop_proc1(b_to_a[0], a_to_b[1]);
		close(a_to_b[1]);
		close(b_to_a[0]);
		exit(ret);
	}

	close(a_to_b[1]);
	close(b_to_a[0]);
	prctl(PR_SET_NAME, "proc2", 0, 0, 0);
	ret = context_switch_loop_proc2(a_to_b[0], b_to_a[1]);
	close(a_to_b[0]);
	close(b_to_a[1]);

	if (ret) {
		kill(proc1_pid, SIGKILL);
		return ret;
	}

	if (waitpid(proc1_pid, &status, 0) != proc1_pid || !WIFEXITED(status) ||
	    WEXITSTATUS(status))
		return 1;

	return 0;
}

DEFINE_WORKLOAD(context_switch_loop);