/*****************************************************************************\
 *  step_ctx.c - step_ctx task functions for use by AIX/POE
 *****************************************************************************
 *  Copyright (C) 2004-2007 The Regents of the University of California.
 *  Copyright (C) 2008-2010 Lawrence Livermore National Security.
 *  Produced at Lawrence Livermore National Laboratory (cf, DISCLAIMER).
 *  Written by Morris Jette <jette1@llnl.gov>.
 *  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 <errno.h>
#include <netinet/in.h>
#include <poll.h>
#include <pthread.h>
#include <stdarg.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/param.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <unistd.h>

#include "slurm/slurm.h"

#include "src/common/bitstring.h"
#include "src/common/fd.h"
#include "src/common/hostlist.h"
#include "src/common/net.h"
#include "src/common/read_config.h"
#include "src/common/slurm_protocol_api.h"
#include "src/common/slurm_protocol_defs.h"
#include "src/common/timers.h"
#include "src/common/xmalloc.h"
#include "src/common/xstring.h"
#include "src/interfaces/cred.h"
#include "src/interfaces/switch.h"

#include "src/srun/launch.h"
#include "src/srun/signals.h"
#include "src/srun/step_ctx.h"

static void _job_fake_cred(struct slurm_step_ctx_struct *ctx)
{
	uint32_t node_cnt = ctx->step_resp->step_layout->node_cnt;
	slurm_cred_arg_t *arg = xmalloc(sizeof(*arg));

	memcpy(&arg->step_id, &ctx->step_req->step_id, sizeof(arg->step_id));
	arg->uid = getuid();

	arg->job_nhosts = node_cnt;
	arg->job_hostlist = ctx->step_resp->step_layout->node_list;
	arg->job_mem_alloc = xmalloc(sizeof(uint64_t));
	arg->job_mem_alloc[0] = 0;
	arg->job_mem_alloc_rep_count = xmalloc(sizeof(uint64_t));
	arg->job_mem_alloc_rep_count[0] = node_cnt;
	arg->job_mem_alloc_size = 1;

	arg->step_hostlist = ctx->step_req->node_list;
	arg->step_mem_alloc = xmalloc(sizeof(uint64_t));
	arg->step_mem_alloc[0] = 0;
	arg->step_mem_alloc_rep_count = xmalloc(sizeof(uint64_t));
	arg->step_mem_alloc_rep_count[0] = node_cnt;
	arg->step_mem_alloc_size = 1;

	arg->job_core_bitmap = bit_alloc(node_cnt);
	bit_set_all(arg->job_core_bitmap);
	arg->step_core_bitmap  = bit_alloc(node_cnt);
	bit_set_all(arg->step_core_bitmap);

	arg->cores_per_socket = xmalloc(sizeof(uint16_t));
	arg->cores_per_socket[0] = 1;
	arg->sockets_per_node = xmalloc(sizeof(uint16_t));
	arg->sockets_per_node[0] = 1;
	arg->sock_core_rep_count = xmalloc(sizeof(uint32_t));
	arg->sock_core_rep_count[0] = node_cnt;

	ctx->step_resp->cred = slurm_cred_faker(arg);

	/* Don't free, this memory will be free'd later */
	arg->job_hostlist = NULL;
	arg->step_hostlist = NULL;
	slurm_cred_free_args(arg);
}

/*
 * step_ctx_create - Create a job step and its context.
 * IN step_params - job step parameters
 * IN timeout - in milliseconds
 * IN srun_opt - srun options
 * OUT timed_out - indicate if poll timed-out
 * RET the step context or NULL on failure with slurm errno set
 * NOTE: Free allocated memory using step_ctx_destroy()
 */
extern slurm_step_ctx_t *step_ctx_create_timeout(
	job_step_create_request_msg_t *step_req, int timeout, bool *timed_out,
	srun_opt_t *srun_opt)
{
	struct slurm_step_ctx_struct *ctx = NULL;
	job_step_create_response_msg_t *step_resp = NULL;
	int i, rc, time_left;
	int sock = -1;
	uint16_t port = 0;
	int errnum = 0;
	struct pollfd fds[2];
	long elapsed_time;
	DEF_TIMERS;

	xassert(step_req);

	/* Don't need a port for async steps since they're fire and forget */
	if (!srun_opt->async) {
		/*
		 * We will handle the messages in the step_launch.c
		 * message handler, but we need to open the socket right
		 * now so we can tell the controller which port to use.
		 */
		if (slurm_init_msg_engine_srun_ports(&sock, &port) !=
		    SLURM_SUCCESS) {
			error("unable to initialize step request socket: %m");
			return NULL;
		}

		step_req->port = port;
	}

	rc = slurm_job_step_create(step_req, &step_resp);
	if ((rc < 0) && launch_step_retry_errno(errno)) {
		START_TIMER;
		errnum = errno;
		fds[0].fd = sock;
		fds[0].events = POLLIN;
		fds[1].fd = srun_sig_eventfd;
		fds[1].events = POLLIN;

		while (1) {
			END_TIMER;
			elapsed_time = (TIMER_DURATION_USEC() / 1000);
			if (elapsed_time >= timeout)
				break;
			time_left = timeout - elapsed_time;
			i = poll(fds, 2, time_left);

			/* Caught destroy signal during poll() */
			if (fds[1].revents & POLLIN)
				break;

			if (i == 0)
				*timed_out = true;
			if (i >= 0)
				break;
			if ((errno == EINTR) || (errno == EAGAIN))
				continue;
			break;
		}

		slurm_mutex_lock(&srun_destroy_sig_lock);
		if (srun_destroy_sig) {
			info("Cancelled pending job step with signal %d",
			     srun_destroy_sig);
			errnum = ESLURM_ALREADY_DONE;
		}
		slurm_mutex_unlock(&srun_destroy_sig_lock);

		if (sock != -1)
			close(sock);
		errno = errnum;
	} else if ((rc < 0) || (step_resp == NULL)) {
		if (sock != -1)
			close(sock);
	} else {
		if (step_resp->state == JOB_PENDING) {
			/*
			 * Async pending step: controller queued it
			 * with a real step_id but did not launch.
			 * Surface step_id to caller via step_req;
			 * no ctx (srun returns SUCCESS for --async
			 * without a live step).
			 */
			if (step_req->step_id.step_id == NO_VAL)
				step_req->step_id.step_id =
					step_resp->step_id.step_id;
			if (step_req->array_task_id != NO_VAL)
				step_req->step_id.job_id =
					step_resp->step_id.job_id;
			slurm_free_job_step_create_response_msg(step_resp);
			if (sock != -1)
				close(sock);
			errno = ESLURM_STEP_QUEUED;
			return NULL;
		}

		ctx = xmalloc(sizeof(struct slurm_step_ctx_struct));
		ctx->launch_state = NULL;
		ctx->magic	= STEP_CTX_MAGIC;
		ctx->job_id	= step_req->step_id.job_id;
		ctx->step_req   = step_req;
		/*
		 * Grab the step id here if we don't already have it, we will
		 * need to to send to the slurmd.
		 */
		if (step_req->step_id.step_id == NO_VAL)
			step_req->step_id.step_id = step_resp->step_id.step_id;

		if (step_req->array_task_id != NO_VAL) {
			step_req->step_id.job_id = step_resp->step_id.job_id;
			ctx->job_id = step_resp->step_id.job_id;
		}

		ctx->step_resp	= step_resp;
		ctx->launch_state = step_launch_state_create(ctx);
		ctx->launch_state->slurmctld_socket_fd = sock;
	}

	return (slurm_step_ctx_t *) ctx;
}

/*
 * step_ctx_create_no_alloc - Create a job step and its context without
 *                            getting an allocation.
 * IN step_params - job step parameters
 * IN step_id     - since we are faking it give me the id to use
 * RET the step context or NULL on failure with slurm errno set
 * NOTE: Free allocated memory using step_ctx_destroy()
 */
extern slurm_step_ctx_t *step_ctx_create_no_alloc(
	job_step_create_request_msg_t *step_req, uint32_t step_id)
{
	struct slurm_step_ctx_struct *ctx = NULL;
	job_step_create_response_msg_t *step_resp = NULL;
	int sock = -1;
	uint16_t port = 0;

	xassert(step_req);
	/* We will handle the messages in the step_launch.c message handler,
	 * but we need to open the socket right now so we can tell the
	 * controller which port to use.
	 */
	if (!(step_req->flags & SSF_ASYNC)) {
		if (slurm_init_msg_engine_srun_ports(&sock, &port) !=
		    SLURM_SUCCESS) {
			error("unable to initialize step context socket: %m");
			return NULL;
		}
		step_req->port = port;
	}

	/* Then make up a response with only certain things filled in */
	step_resp = (job_step_create_response_msg_t *)
		xmalloc(sizeof(job_step_create_response_msg_t));

	step_resp->step_layout = fake_slurm_step_layout_create(
		step_req->node_list,
		NULL, NULL,
		step_req->min_nodes,
		step_req->num_tasks,
		0);

	step_resp->step_id.step_id = step_id;

	ctx = xmalloc(sizeof(struct slurm_step_ctx_struct));
	ctx->launch_state = NULL;
	ctx->magic	= STEP_CTX_MAGIC;
	ctx->job_id	= step_req->step_id.job_id;
	ctx->step_req   = step_req;

	/*
	 * Grab the step id here if we don't already have it, we will
	 * need to to send to the slurmd.
	 */
	if (step_req->step_id.step_id == NO_VAL)
		step_req->step_id.step_id = step_resp->step_id.step_id;

	ctx->step_resp	= step_resp;
	ctx->launch_state = step_launch_state_create(ctx);
	ctx->launch_state->slurmctld_socket_fd = sock;

	_job_fake_cred(ctx);

	return (slurm_step_ctx_t *)ctx;
}
