/*****************************************************************************\
 *  srun_cr.c - Checkpoint/Restart wrapper for srun
 *****************************************************************************
 *  Copyright (C) 2009 National University of Defense Technology, China.
 *  Written by Hongia Cao.
 *  CODE-OCEC-09-009. All rights reserved.
 *
 *  This file is part of Slurm, a resource management program.
 *  For details, see <https://slurm.schedmd.com/>.
 *  Please also read the included file: DISCLAIMER.
 *
 *  Slurm 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.
 *
 *  In addition, as a special exception, the copyright holders give permission
 *  to link the code of portions of this program with the OpenSSL library under
 *  certain conditions as described in each individual source file, and
 *  distribute linked combinations including the two. You must obey the GNU
 *  General Public License in all respects for all of the code used other than
 *  OpenSSL. If you modify file(s) with this exception, you may extend this
 *  exception to your version of the file(s), but you are not obligated to do
 *  so. If you do not wish to do so, delete this exception statement from your
 *  version.  If you delete this exception statement from all source files in
 *  the program, then also delete it here.
 *
 *  Slurm 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 Slurm; if not, write to the Free Software Foundation, Inc.,
 *  51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA.
\*****************************************************************************/

#include "config.h"

#include <errno.h>
#include <inttypes.h>
#include <libcr.h>
#include <poll.h>
#include <pthread.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <time.h>
#include <sys/resource.h>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/types.h>
#include <sys/un.h>
#include <sys/wait.h>
#include <unistd.h>

#include "slurm/slurm.h"

#include "src/common/fd.h"
#include "src/common/log.h"
#include "src/common/macros.h"
#include "src/common/read_config.h"
#include "src/common/xmalloc.h"
#include "src/common/xstring.h"

static char *cr_run_path = BLCR_HOME "/bin/cr_run";
static char *srun_path = SLURM_PREFIX "/bin/srun";

/* global variables */
static char **srun_argv = NULL;
static pid_t srun_pid = 0;

static uint32_t jobid = 0;
static uint32_t stepid = 0xFFFFFFFF;
static char *nodelist = NULL;

static char cr_sock_addr[32];
static int listen_fd = -1;

static int step_launched = 0;
static pthread_mutex_t step_launch_mutex = PTHREAD_MUTEX_INITIALIZER;
static pthread_cond_t step_launch_cond = PTHREAD_COND_INITIALIZER;

static cr_client_id_t cr_id = -1;

static void remove_listen_socket(void);
static int  _wait_for_srun_connect(void);
static void _read_info_from_srun(int srun_fd);

/**************** copied and modified from cr_restart of BLCR ****************/
static void signal_child (int, siginfo_t *, void *);

static void
signal_self(int sig)
{
	struct sigaction sa;

	/* restore default (in kernel) handler */
	sa.sa_handler = SIG_DFL;
	sa.sa_flags = SA_RESTART | SA_NOMASK;
	sigemptyset(&sa.sa_mask);
	(void)sigaction(sig, &sa, NULL);

	/* send to self */
	raise(sig);

	/* restore self as handler */
	sa.sa_sigaction = &signal_child;
	sa.sa_flags = SA_RESTART | SA_NOMASK | SA_SIGINFO;
	(void)sigaction(sig, &sa, NULL);
}

static void
signal_child (int sig, siginfo_t *siginfo, void *context)
{
	if (srun_pid == 0) {	/* srun not forked yet */
		signal_self(sig);
		return;
	}

	if ((siginfo->si_code > 0) &&	/* si_code > 0 indicates sent by kernel */
	    ((sig == SIGILL) || (sig == SIGFPE) ||
	     (sig == SIGBUS) || (sig == SIGSEGV) )) {
		/* This signal is OUR error, so we don't forward */
		signal_self(sig);
	} else if ((sig == SIGTSTP) || (sig == SIGTTIN) || (sig == SIGTTOU)) {
		/* The catchable stop signals go to child AND self */
		(void)kill(srun_pid, sig);
		signal_self(sig);
	} else {
		/* default case */
		kill(srun_pid, sig);
	}
}

static void
mimic_exit(int status)
{
	if (WIFEXITED(status)) {
		/* easy to mimic normal return */
		exit(WEXITSTATUS(status));
	} else if (WIFSIGNALED(status)) {
		/* disable generation of a 'core' */
		struct rlimit r;
		r.rlim_cur = r.rlim_max = 0;
		(void)setrlimit(RLIMIT_CORE, &r);

		/* now raise the signal */
		signal_self(WTERMSIG(status));
	} else {
		error("%s: Unexpected status from child", __func__);
		exit(-1);
	}
}
/****************************************************************/
static void
on_child_exit(int signum, siginfo_t *siginfo, void *arg)
{
	int status;

	/*
  	 * if srun_cr is checkpoint/restart-ed after srun exited,
  	 * srun_pid will be the pid of the new srun.
	 */
	cr_enter_cs(cr_id);
	if (waitpid(srun_pid, &status, WNOHANG) == srun_pid) {
		verbose("%s: srun(%d) exited, status: %d",
			__func__, srun_pid, status);
		mimic_exit(status);
	}
	kill(srun_pid, SIGKILL);
	cr_leave_cs(cr_id);
}

static int
_slurm_debug_env_val (void)
{
	long int level = 0;
	const char *val;

	if ((val = getenv ("SLURM_DEBUG"))) {
		char *p;
		if ((level = strtol (val, &p, 10)) < -LOG_LEVEL_INFO)
			level = -LOG_LEVEL_INFO;
		if (p && *p != '\0')
			level = 0;
	}
	return ((int) level);
}


static void
update_env(char *name, char *val)
{
	char *buf = NULL;

	xstrfmtcat (buf, "%s=%s", name, val);
	if (putenv(buf))
		fatal("%s: failed to update env: %m", __func__);
}

static int
init_srun_argv(int argc, char **argv)
{
	int i;

	srun_argv = (char **)xmalloc(sizeof(char *) * (argc + 3));

	srun_argv[0] = cr_run_path;
	srun_argv[1] = "--omit";
	srun_argv[2] = srun_path;
	for (i = 1; i < argc; i ++) {
		srun_argv[i + 2] = argv[i];
	}
	srun_argv[argc + 2] = NULL;

	return  0;
}

/* remove the listen socket file */
static void
remove_listen_socket(void)
{
	unlink(cr_sock_addr);
}

/*
 * create_listen_socket - create a listening UNIX domain socket
 *     for srun to connect
 * RETURN: the socket fd on success, -1 on error
 */
static int
create_listen_socket(void)
{
	struct sockaddr_un sa;
	unsigned int sa_len;
	int re_use_addr = 1;


	close (listen_fd);	/* close possible old socket */

	sprintf(cr_sock_addr, "/tmp/sock.srun_cr.%u", (unsigned int)getpid());

	listen_fd = socket(AF_UNIX, SOCK_STREAM, 0);
	if (listen_fd < 0) {
		error("%s: failed to create listen socket: %m", __func__);
		return -1;
	}

	sa.sun_family = AF_UNIX;
	strcpy(sa.sun_path, cr_sock_addr);
	sa_len = strlen(sa.sun_path) + sizeof(sa.sun_family);

	unlink(sa.sun_path);	/* remove possible old socket */

	if (setsockopt(listen_fd, SOL_SOCKET, SO_REUSEADDR,
		       (void*)&re_use_addr, sizeof(int)) == -1) {
		error("%s: setsockopt: %m", __func__);
	}

	if (bind(listen_fd, (struct sockaddr *)&sa, sa_len) < 0) {
		error("%s: failed to bind listen socket: %m", __func__);
		unlink(sa.sun_path);
		return -1;
	}

	if (listen(listen_fd, 2) < 0) {
		error("%s: failed to listen: %m", __func__);
		unlink(sa.sun_path);
		return -1;
	}

	fd_set_nonblocking(listen_fd);

	return listen_fd;
}

/*
 * fork_exec_srun - fork and exec srun
 * GLOBALS cr_argv: arguments for running srun
 * RETURN: 0 on success, otherwise on error
 */
static int
fork_exec_srun(void)
{
	int rc = 0;
	sigset_t sigset;

	listen_fd = create_listen_socket();
	if (listen_fd < 0) {
		return -1;
	}

	srun_pid = fork();
	if (srun_pid < 0) {
		error("failed to fork child process: %m");
		return -1;
	} else if (srun_pid == 0) {	/* child */
		/*
		 * remove srun from the foreground process group,
		 * or Ctrl-C will cause SIGINT duplicated
		 */
		setpgid(0, 0);

		update_env("SLURM_SRUN_CR_SOCKET", cr_sock_addr);

		/*
		 * BLCR blocks all signals in thread-context callback functions
		 */
		sigemptyset(&sigset);
		pthread_sigmask(SIG_SETMASK, &sigset, NULL);

		execv(srun_argv[0], srun_argv);
		perror("failed execv srun");
		exit(-1);
	}

	return rc;
}

/*
 * get_step_image_dir - get the dir to store step task images
 * IN cr: checkpoint/restart
 * RET image dir on success, NULL on error
 *
 * NOTE: only can be called in callbak
 */
static char *
get_step_image_dir(int cr)
{
	const struct cr_checkpoint_info *ckpt_info;
	const struct cr_restart_info *rstrt_info;
	const char *dest;
	char *rchar, *dir;

	if (cr) {		/* checkpoint */
		ckpt_info = cr_get_checkpoint_info();
		if (!ckpt_info) {
			error("failed to get checkpoint info: %s",
			      cr_strerror(errno));
			return NULL;
		}
		dest = ckpt_info->dest;
	} else {		/* retart */
		rstrt_info = cr_get_restart_info();
		if (!rstrt_info) {
			error("failed to get restart info: %s",
			      cr_strerror(errno));
			return NULL;
		}
		dest = rstrt_info->src;
	}

	rchar = strrchr(dest, '/');
	if (rchar) {
		dir = xstrndup(dest, rchar - dest + 1);
	}
	xstrfmtcat(dir, "%u.%u", jobid, stepid);

	return dir;
}

static int
cr_callback(void *unused)
{
	int rc;
	char *step_image_dir = NULL;

	rc = CR_CHECKPOINT_READY;
	if (step_launched) {
		step_image_dir = get_step_image_dir(1);
		if (step_image_dir == NULL) {
			error ("failed to get step image directory");
			rc = CR_CHECKPOINT_PERM_FAILURE;
		} else if (slurm_checkpoint_tasks(jobid,
						  stepid,
						  time(NULL), /* timestamp */
						  step_image_dir,
						  60, /* wait */
						  nodelist) != SLURM_SUCCESS) {
			error ("failed to checkpoint step tasks");
			rc = CR_CHECKPOINT_PERM_FAILURE;
		}
		xfree(step_image_dir);
	}
	rc = cr_checkpoint(rc);	/* dump */

	if (rc < 0) {
		fatal("checkpoint failed: %s", cr_strerror(errno));
	} else if (rc == 0) {
		/* continue, nothing to do */
	} else {
		/* restarted */
		if (srun_pid) { /* srun forked */
			if (step_launched) {
				step_image_dir = get_step_image_dir(0);
				if (step_image_dir == NULL) {
					fatal("failed to get step image directory");
				}
				update_env("SLURM_RESTART_DIR", step_image_dir);
				xfree(step_image_dir);
			}

			if (fork_exec_srun()) {
				fatal("failed fork/exec srun");
			}
		}

		/* XXX: step_launched => listen_fd valid */
		step_launched = 0;

		debug2("step not launched.");

		slurm_cond_broadcast(&step_launch_cond);
	}

	return 0;
}

int
main(int argc, char **argv)
{
	int debug_level, sig, srun_fd;
	struct sigaction sa;
	log_options_t logopt = LOG_OPTS_STDERR_ONLY;
	struct sockaddr_un ca;
	unsigned int ca_len = sizeof(ca);

	atexit(remove_listen_socket);

	/* copied from srun */
	slurm_conf_init(NULL);
	debug_level = _slurm_debug_env_val();
	logopt.stderr_level += debug_level;
	log_init(xbasename(argv[0]), logopt, 0, NULL);

	if (init_srun_argv(argc, argv)) {
		fatal("failed to initialize arguments for running srun");
	}

	if ((cr_id = cr_init()) < 0) {
		fatal("failed to initialize libcr: %s", cr_strerror(errno));
	}
	(void)cr_register_callback(cr_callback, NULL, CR_THREAD_CONTEXT);

	/* forward signals. copied from cr_restart */
	sa.sa_sigaction = signal_child;
	sa.sa_flags = SA_RESTART | SA_NODEFER | SA_SIGINFO;
	sigemptyset(&sa.sa_mask);
	for (sig = 0;  sig < _NSIG; sig ++) {
		if (sig == SIGSTOP ||
		    sig == SIGKILL ||
		    sig == SIGCHLD)
			continue;
		sigaction(sig, &sa, NULL);
	}
	sa.sa_sigaction = on_child_exit;
	sa.sa_flags = SA_RESTART | SA_SIGINFO | SA_NOCLDSTOP;
	sigaction(SIGCHLD, &sa, NULL);

	cr_enter_cs(cr_id); /* BEGIN CS: avoid race condition of whether srun is forked */
	if ( fork_exec_srun() ) {
		fatal("failed fork/exec/wait srun");
	}
	cr_leave_cs(cr_id); /* END CS */

	while (1) {
		slurm_mutex_lock(&step_launch_mutex);
		while (step_launched) {
			/* just avoid busy waiting */
			slurm_cond_wait(&step_launch_cond, &step_launch_mutex);
		}
		slurm_mutex_unlock(&step_launch_mutex);

		if (_wait_for_srun_connect() < 0)
			continue;

		cr_enter_cs(cr_id); /* BEGIN CS: checkpoint(callback) will be delayed */

		srun_fd = accept(listen_fd, (struct sockaddr*)&ca, &ca_len);
		if (srun_fd < 0) {
			/* restarted before enter CS. socket will not be restored */
			if (errno == EBADF) {
				cr_leave_cs(cr_id);
				continue;
			} else {
				fatal("failed to accept socket: %m");
			}
		}

		_read_info_from_srun(srun_fd);
		close(srun_fd);

		step_launched = 1;
		debug2("step launched");

		cr_leave_cs(cr_id); /* END CS */
	}

	return 0;
}

static int
_wait_for_srun_connect(void)
{
	struct pollfd fds[1];
	int rc;

	fds[0].fd = listen_fd;
	fds[0].events = POLLIN;

	while ((rc = poll(fds, 1, -1)) < 0) {
		switch (errno) {
		case EAGAIN:
		case EINTR:
			continue;
		case EBADF:	/* restarted */
			return -1;
		case ENOMEM:
		case EINVAL:
		case EFAULT:
			fatal("poll: %m");
		default:
			error("poll: %m. Continuing...");
		}
	}
	return 0;
}

static void
_read_info_from_srun(int srun_fd)
{
	int len;

	if (read(srun_fd, &jobid, sizeof(uint32_t)) != sizeof(uint32_t))
		fatal("%s: failed to read jobid: %m", __func__);

	if (read(srun_fd, &stepid, sizeof(uint32_t)) != sizeof(uint32_t))
		fatal("%s: failed to read stepid: %m", __func__);

	if (read(srun_fd, &len, sizeof(int)) != sizeof(int))
		fatal("%s: failed to read nodelist length: %m", __func__);
	if (len > (32 * 1024))
		fatal("%s: nodelist length too large (%d)", __func__, len);

	xfree(nodelist);
	nodelist = (char *)xmalloc(len + 1);
	if (read(srun_fd, nodelist, len + 1) != len + 1)
		fatal("%s: failed to read nodelist: %m", __func__);
}
