/*
 * Copyright (C) 2013-2021 Canonical, Ltd.
 * Copyright (C) 2022-2026 Colin Ian King.
 *
 * This program is free software; you can redistribute it and/or
 * modify it under the terms of the GNU General Public License
 * as published by the Free Software Foundation; either version 2
 * of the License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301, USA.
 *
 */
#include "stress-ng.h"
#include "core-builtin.h"
#include "core-capabilities.h"
#include "core-killpid.h"
#include "core-out-of-memory.h"
#include "core-resources.h"
#include "core-signal.h"

#define MIN_MEM_FREE	(16 * MB)

#define MIN_RESOURCES_PROCS	(1)
#define MAX_RESOURCES_PROCS	(4096)
#define DEFAULT_RESOURCES_PROCS	(256)

#define MIN_RESOURCES_NUM	(1)
#define MAX_RESOURCES_NUM	(4096)
#define DEFAULT_RESOURCES_NUM	(1024)

static const stress_help_t help[] = {
	{ NULL,	"resources N",	     "start N workers consuming system resources" },
	{ NULL,	"resources-mlock",   "attempt to mlock pages into memory" },
	{ NULL, "resources-num N",   "number of resources to allocate per instance" },
	{ NULL,	"resources-ops N",   "stop after N resource bogo operations" },
	{ NULL, "resources-procs N", "number of child processes per instance" },
	{ NULL,	NULL,                NULL }
};

static pid_t stress_reources_pid = -1;

/*
 *  stress_resources_alarm()
 *	send SIGALARM to all known valid running PIDs in s_pids[]
 */
static void stress_resources_alarm(
	const stress_pid_t * const s_pids,
	const size_t resources_procs)
{
	size_t i;

	for (i = 0; i < resources_procs; i++) {
		if (s_pids[i].pid != -1)
			(void)kill(s_pids[i].pid, SIGALRM);
	}
	return;
}

/*
 *  stress_resources()
 *	stress by forking and exiting
 */
static int stress_resources(stress_args_t *args)
{
	stress_memory_info_t info;
	const size_t pipe_size = stress_fs_max_pipe_size_get();
	size_t min_mem_free;
	size_t resources_num = DEFAULT_RESOURCES_NUM;
	size_t resources_procs = DEFAULT_RESOURCES_PROCS;
	stress_resources_t *resources;
	stress_pid_t *s_pids;
	bool resources_mlock = false;
	bool report_freeing = true;

	stress_reources_pid = getpid();

	if (!stress_setting_get("resources-mlock", &resources_mlock)) {
		if (g_opt_flags & OPT_FLAGS_AGGRESSIVE)
			resources_mlock = true;
	}
	if (!stress_setting_get("resources-num", &resources_num)) {
		if (g_opt_flags & OPT_FLAGS_MINIMIZE)
			resources_num = MIN_RESOURCES_NUM;
		if (g_opt_flags & OPT_FLAGS_MAXIMIZE)
			resources_num = MAX_RESOURCES_NUM;
	}
	if (!stress_setting_get("resources-procs", &resources_procs)) {
		if (g_opt_flags & OPT_FLAGS_MINIMIZE)
			resources_procs = MIN_RESOURCES_PROCS;
		if (g_opt_flags & OPT_FLAGS_MAXIMIZE)
			resources_procs = MAX_RESOURCES_PROCS;
	}

	stress_memory_info_get(&info);
	min_mem_free = (info.freemem / 100) * 2;
	if (min_mem_free < MIN_MEM_FREE)
		min_mem_free = MIN_MEM_FREE;

#if defined(MCL_FUTURE)
	if (resources_mlock)
		(void)shim_mlockall(MCL_FUTURE);
#else
	UNEXPECTED
#endif

	s_pids = stress_sync_s_pids_mmap(resources_procs);
	if (s_pids == MAP_FAILED) {
		pr_inf_skip("%s: mmap %zu PIDs failed%s, skipping stressor\n",
			args->name, resources_procs, stress_memory_free_get());
		return EXIT_NO_RESOURCE;
	}

	resources = (stress_resources_t *)calloc(resources_num, sizeof(*resources));
	if (!resources) {
		pr_inf_skip("%s: allocate %zu resource structures failed%s, skipping stressor\n",
			args->name, resources_num, stress_memory_free_get());
		(void)stress_sync_s_pids_munmap(s_pids, resources_procs);
		return EXIT_NO_RESOURCE;
	}

	if (stress_instance_zero(args))
		pr_inf("%s: using %zu resource%s and spawning %zu child process%s per instance\n",
			args->name,
			resources_num, resources_num == 1 ? "" : "s",
			resources_procs, resources_procs == 1 ? "" : "es");

	stress_proc_state_set(args->name, STRESS_STATE_SYNC_WAIT);
	stress_sync_start_wait(args);
	stress_proc_state_set(args->name, STRESS_STATE_RUN);

	do {
		size_t i;
		size_t forked;
		size_t reaped;
		stress_pid_t *s_pids_head;

		stress_sync_init_pids(s_pids, resources_procs);
		s_pids_head = NULL;

		for (forked = 0, i = 0; i < resources_procs; i++) {
			pid_t pid;

			stress_memory_info_get(&info);
			if ((info.freemem > 0) && (info.freemem < min_mem_free))
				break;
			if (!stress_continue(args))
				break;
			pid = fork();
			if (pid == 0) {
				size_t n;

				stress_proc_state_set(args->name, STRESS_STATE_RUN);
				stress_set_oom_adjustment(args, true);
				VOID_RET(int, stress_capabilities_drop(args->name));
				stress_make_it_fail_set();
				(void)stress_sched_settings_apply(true);

				if (!stress_continue(args))
					_exit(0);
				n = stress_resources_allocate(args, resources, resources_num, pipe_size, min_mem_free, true);
				if (stress_continue(args)) {
					shim_sched_yield();
					stress_resources_access(args, resources, n);
				}
				shim_sched_yield();
				stress_resources_free(args, resources, n);
				free(resources);
				_exit(0);
			} else if (pid > 0) {
				forked++;
			}
			s_pids[i].pid = pid;
			stress_sync_order_pid(&s_pids_head, &s_pids[i]);

			if (UNLIKELY(!stress_continue(args)))
				break;
			stress_bogo_inc(args);
		}

		reaped = 0;
		do {
			pid_t pid;
			int status;
			stress_pid_t *s_pid;

			if (!stress_continue(args)) {
				stress_resources_alarm(s_pids, resources_procs);

				if (report_freeing && stress_instance_zero(args)) {
					pr_inf("%s: freeing resources (make take a while)\n", args->name);
					report_freeing = false;
				}
			}
			pid = waitpid(0, &status, 0);
			/* retry on failures, such as EINTR */
			if (UNLIKELY(pid < 0))
				continue;

			s_pid = struct_sync_find_pid(s_pids_head, pid);
			if (UNLIKELY(s_pid == NULL))
				continue;
			s_pid->pid = -1;
			reaped++;
		} while (reaped < forked);
	} while (stress_continue(args));

	stress_proc_state_set(args->name, STRESS_STATE_DEINIT);
	free(resources);
	(void)stress_sync_s_pids_munmap(s_pids, resources_procs);

	return EXIT_SUCCESS;
}

static const stress_opt_t opts[] = {
	{ OPT_resources_mlock, "resources-mlock", TYPE_ID_BOOL, 0, 1, NULL },
	{ OPT_resources_num,   "resources-num",   TYPE_ID_SIZE_T, MIN_RESOURCES_NUM, MAX_RESOURCES_NUM, NULL },
	{ OPT_resources_procs, "resources-procs", TYPE_ID_SIZE_T, MIN_RESOURCES_PROCS, MAX_RESOURCES_PROCS, NULL },
	END_OPT,
};

static const stress_exercises_t exercises[] = {
	STRESS_EX_FEATURE("chaotic-load"),
	STRESS_EX_FEATURE("load-average"),
	STRESS_EX_FEATURE("memory-stalls"),
	STRESS_EX_FEATURE("vmalloc"),

	STRESS_EX_SYSCALL("close"),
#if defined(HAVE_EVENTFD)
	STRESS_EX_SYSCALL("eventfd"),
#endif
	STRESS_EX_SYSCALL("fcntl"),
	STRESS_EX_SYSCALL("fork"),
	STRESS_EX_SYSCALL("ftruncate"),
#if defined(HAVE_SYS_INOTIFY_H)
	STRESS_EX_SYSCALL("inotify_init"),
	STRESS_EX_SYSCALL("inotify_add_watch"),
	STRESS_EX_SYSCALL("inotify_rm_watch"),
#endif
	STRESS_EX_SYSCALL("kill"),
	STRESS_EX_SYSCALL("madvise"),
#if defined(HAVE_MEMFD_CREATE)
	STRESS_EX_SYSCALL("memfd_create"),
#endif
#if defined(__NR_memfd_secret)
	STRESS_EX_SYSCALL("memfd_secret"),
#endif
	STRESS_EX_SYSCALL("mlock"),
	STRESS_EX_SYSCALL("mmap"),
#if defined(HAVE_MQ_SYSV) &&	\
    defined(HAVE_SYS_IPC_H) &&	\
    defined(HAVE_SYS_MSG_H)
	STRESS_EX_SYSCALL("msgctl"),
	STRESS_EX_SYSCALL("msgget"),
#endif
#if defined(HAVE_LIB_PTHREAD) &&	\
    defined(HAVE_THREADS_H) &&		\
    defined(HAVE_MTX_T) && 		\
    defined(HAVE_MTX_DESTROY) &&	\
    defined(HAVE_MTX_INIT)
	STRESS_EX_SYSCALL("mtx_destroy"),
	STRESS_EX_SYSCALL("mtx_init"),
#endif
#if defined(HAVE_LIB_RT) &&	\
    defined(HAVE_MQ_POSIX) &&	\
    defined(HAVE_MQUEUE_H)
	STRESS_EX_SYSCALL("mq_destroy"),
	STRESS_EX_SYSCALL("mq_open"),
#endif
	STRESS_EX_SYSCALL("munlock"),
	STRESS_EX_SYSCALL("munmap"),
	STRESS_EX_SYSCALL("open"),
#if defined(HAVE_PIDFD_OPEN)
	STRESS_EX_SYSCALL("pidfd_open"),
#endif
#if defined(HAVE_PIDFD_GETFD)
	STRESS_EX_SYSCALL("pidfd_get"),
#endif
	STRESS_EX_SYSCALL("pipe"),
#if defined(HAVE_PKEY_ALLOC) && \
    defined(HAVE_PKEY_FREE)
	STRESS_EX_SYSCALL("pkey_alloc"),
	STRESS_EX_SYSCALL("pkey_free"),
#endif
	STRESS_EX_SYSCALL("sbrk"),
	STRESS_EX_SYSCALL("sched_yield"),
#if defined(HAVE_LIB_PTHREAD) &&	\
    defined(HAVE_SEM_POSIX)
	STRESS_EX_SYSCALL("sem_destroy>"),
	STRESS_EX_SYSCALL("sem_init"),
#endif
#if defined(HAVE_SEM_SYSV) &&	\
    defined(HAVE_KEY_T)
	STRESS_EX_SYSCALL("semctl"),
	STRESS_EX_SYSCALL("semget"),
#endif
	STRESS_EX_SYSCALL("socket"),
	STRESS_EX_SYSCALL("sockpair"),
#if defined(HAVE_LIB_RT) &&	\
    defined(HAVE_TIMER_CREATE)
	STRESS_EX_SYSCALL("timer_create"),
#endif
#if defined(HAVE_LIB_RT) &&	\
    defined(HAVE_TIMER_DELETE)
	STRESS_EX_SYSCALL("timer_delete"),
#endif
#if defined(HAVE_SYS_TIMERFD_H) &&	\
    defined(HAVE_TIMERFD_CREATE) &&	\
    defined(CLOCK_REALTIME)
	STRESS_EX_SYSCALL("timerfd_create"),
#endif
#if defined(HAVE_USERFAULTFD)
	STRESS_EX_SYSCALL("userfaultfd"),
#endif
	STRESS_EX_SYSCALL("waitpid"),

	STRESS_EX_END,
};

const stressor_info_t stress_resources_info = {
	.stressor = stress_resources,
	.classifier = CLASS_MEMORY | CLASS_OS,
	.opts = opts,
	.help = help,
	.exercises = exercises,
};
