/****************************************************************************\
 *  dbd_conn.c - functions to manage the connection to the SlurmDBD
 *****************************************************************************
 *  Copyright (C) SchedMD LLC.
 *  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 "src/common/slurm_xlator.h"

#include <errno.h>
#include <sys/socket.h>
#include <unistd.h>

#include "src/common/fd.h"
#include "src/common/slurmdbd_pack.h"
#include "src/interfaces/conn.h"
#include "src/slurmctld/trigger_mgr.h"
#include "slurmdbd_agent.h"

#define SLURMDBD_TIMEOUT	900	/* Seconds SlurmDBD for response */

static void _acct_full(void)
{
	if (running_in_slurmctld())
		trigger_primary_ctld_acct_full();
}

static void _dbd_fail(void)
{
	if (running_in_slurmctld())
		trigger_primary_dbd_fail();
}

static void _dbd_res_op(void)
{
	if (running_in_slurmctld())
		trigger_primary_dbd_res_op();
}

static void _db_fail(void)
{
	if (running_in_slurmctld())
		trigger_primary_db_fail();
}

static void _db_res_op(void)
{
	if (running_in_slurmctld())
		trigger_primary_db_res_op();
}

static int _connect_dbd_conn(persist_conn_t *pc)
{

	int rc;
	char *backup_host = NULL;

	xassert(pc);

	/* Only setup a backup host on a non ext_dbd connection */
	if (!(pc->flags & PERSIST_FLAG_EXT_DBD))
		backup_host =
			xstrdup(slurm_conf.accounting_storage_backup_host);
again:
	// A connection failure is only an error if backup dne or also fails
	if (backup_host)
		pc->flags |= PERSIST_FLAG_SUPPRESS_ERR;
	else
		pc->flags &= (~PERSIST_FLAG_SUPPRESS_ERR);

	pc->r_uid = SLURM_AUTH_UID_ANY;

	if (((rc = slurm_persist_conn_open(pc)) != SLURM_SUCCESS) &&
	    backup_host) {
		xfree(pc->rem_host);
		// Force the next error to display
		pc->comm_fail_time = 0;
		pc->rem_host = backup_host;
		backup_host = NULL;
		goto again;
	}

	xfree(backup_host);

	if (rc == SLURM_SUCCESS) {
		/*
		 * Increase SLURMDBD_TIMEOUT to wait as long as we need for a
		 * query to complete.
		 */
		pc->timeout = MAX(pc->timeout, SLURMDBD_TIMEOUT * 1000);
		(pc->trigger_callbacks.dbd_resumed)();
		(pc->trigger_callbacks.db_resumed)();
	}

	if (rc == SLURM_SUCCESS) {
		debug("Sent PersistInit msg");
		/* clear errno (checked after this for errors) */
		errno = 0;
	} else {
		if (rc == ESLURM_DB_CONNECTION)
			(pc->trigger_callbacks.db_fail)();
		slurm_persist_conn_close(pc);

		/* This means errno was already set correctly */
		if (rc != SLURM_ERROR)
			errno = rc;
		error("Sending PersistInit msg: %m");
	}

	return rc;
}

extern persist_conn_t *dbd_conn_open(uint16_t *persist_conn_flags,
				     char *cluster_name,
				     char *rem_host,
				     uint16_t rem_port)
{
	persist_conn_t *pc = xmalloc(sizeof(*pc));

	if (persist_conn_flags)
		pc->flags = *persist_conn_flags;
	pc->flags |= (PERSIST_FLAG_DBD | PERSIST_FLAG_RECONNECT);
	pc->persist_type = PERSIST_TYPE_DBD;
	if (cluster_name)
		pc->cluster_name = xstrdup(cluster_name);
	else
		pc->cluster_name = xstrdup(slurm_conf.cluster_name);
	pc->timeout = (slurm_conf.msg_timeout + 35) * 1000;
	if (rem_host)
		pc->rem_host = xstrdup(rem_host);
	else
		pc->rem_host = xstrdup(slurm_conf.accounting_storage_host);
	if (rem_port)
		pc->rem_port = rem_port;
	else
		pc->rem_port = slurm_conf.accounting_storage_port;
	pc->version = SLURM_PROTOCOL_VERSION;

	/* Initialize the callback pointers */
	pc->trigger_callbacks.acct_full = _acct_full;
	pc->trigger_callbacks.dbd_fail = _dbd_fail;
	pc->trigger_callbacks.dbd_resumed = _dbd_res_op;
	pc->trigger_callbacks.db_fail = _db_fail;
	pc->trigger_callbacks.db_resumed = _db_res_op;

	(void)_connect_dbd_conn(pc);

	if (persist_conn_flags)
		*persist_conn_flags = pc->flags;

	return pc;
}

extern int dbd_conn_check_and_reopen(persist_conn_t *pc)
{
	xassert(pc);

	if (pc && pc->conn) {
		debug("Attempt to re-open slurmdbd socket");
		/* clear errno (checked after this for errors) */
		errno = 0;
		return SLURM_SUCCESS;
	}

	/*
	 * Bail out if slurm_conf hasn't published accounting_storage_host
	 * yet. This races the slurmctld startup path where the agent
	 * thread can call us while main is mid-slurm_conf_reinit() and
	 * the host string is transiently NULL. _connect_dbd_conn()'s
	 * _open_persist_conn() would xassert on the NULL rem_host. The
	 * agent's outer loop retries after a backoff, so returning an
	 * error here is fine -- we'll try again once the conf is
	 * loaded.
	 */
	if (!slurm_conf.accounting_storage_host) {
		debug("%s: accounting_storage_host not set; skipping reconnect",
		      __func__);
		return SLURM_ERROR;
	}

	/*
	 * Reset the rem_host just in case we were connected to the backup
	 * before.
	 */
	xfree(pc->rem_host);
	pc->rem_host = xstrdup(slurm_conf.accounting_storage_host);

	return _connect_dbd_conn(pc);
}

extern void dbd_conn_close(persist_conn_t **pc)
{
	persist_msg_t req = { 0 };
	dbd_fini_msg_t get_msg;
	int resp_code = SLURM_SUCCESS;
	int rc;

	if (!pc)
		return;

	/*
	 * Only send the FINI message if we haven't shutdown
	 * (i.e. not slurmctld)
	 */
	if (*(*pc)->shutdown) {
		log_flag(NET, "We are shutdown, not sending DB_FINI to %s:%u",
			 (*pc)->rem_host,
			 (*pc)->rem_port);
		goto destroy_conn;
	}

	/* If the connection is already gone, we don't need to send a fini. */
	if (slurm_persist_conn_writeable(*pc) == -1) {
		log_flag(NET, "unable to send DB_FINI msg to %s:%u",
			 (*pc)->rem_host,
			 (*pc)->rem_port);
		goto destroy_conn;
	}

	memset(&get_msg, 0, sizeof(dbd_fini_msg_t));

	get_msg.close_conn = 0;
	get_msg.commit = false;

	req.msg_type = DBD_FINI;
	req.pcon = *pc;
	req.data = &get_msg;

	if ((rc = dbd_conn_send_recv_rc_msg(SLURM_PROTOCOL_VERSION, &req,
					    &resp_code))) {
		log_flag(NET, "unable to send/recv DB_FINI msg to %s:%u: %s",
			 (*pc)->rem_host, (*pc)->rem_port,
			 slurm_strerror(rc));
		goto destroy_conn;
	}

	if (resp_code != SLURM_SUCCESS) {
		log_flag(NET, "got error in response to DB_FINI msg from %s:%u: %s",
			 (*pc)->rem_host, (*pc)->rem_port,
			 slurm_strerror(resp_code));
		rc = resp_code;
		goto destroy_conn;
	}

	log_flag(NET, "sent DB_FINI msg to %s:%u rc(%d):%s",
		 (*pc)->rem_host, (*pc)->rem_port,
		 rc, slurm_strerror(rc));

	/*
	 * In slurmctld only: synchronize with slurmdbd's connection
	 * teardown before we let the caller (and eventually the next
	 * slurmctld) proceed. Other callers (sacctmgr, slurmstepd, etc.)
	 * just close and exit and don't need to wait.
	 *
	 * Slurmdbd runs _connection_fini_callback() on the recv thread
	 * after that thread sees EOF on our socket. The callback is what
	 * removes us from registered_clusters and NULLs pcon_send. Only
	 * after the callback returns does slurmdbd's _service_connection
	 * call slurm_persist_conn_free_thread_loc(), which destroys the
	 * conn and closes its TCP side.
	 *
	 * So: half-close our write side (FIN to slurmdbd), then wait
	 * for activity on our read side (POLLIN or POLLHUP). By the
	 * time wait_fd() returns, slurmdbd has run the callback and
	 * we won't race a subsequent reconnect against a stale entry
	 * in registered_clusters / pcon_send.
	 *
	 * Bounded by a short poll timeout so we don't hang if slurmdbd
	 * is itself wedged.
	 */
	if (running_in_slurmctld() && (*pc)->conn) {
		int fd = conn_g_get_fd((*pc)->conn);
		if (fd >= 0) {
			char drain[256];
			ssize_t n;

			(void) conn_blocking_g_shutdown((*pc)->conn);
			(void) shutdown(fd, SHUT_WR);

			if (wait_fd(fd, 5, POLLIN)) {
				error("%s: [fd:%d] problem waiting for slurmdbd to close connection on %s:%u",
				      __func__, fd, (*pc)->rem_host,
				      (*pc)->rem_port);
			} else {
				do {
					n = read(fd, drain, sizeof(drain));
				} while ((n < 0) && (errno == EINTR));
			}

			(*pc)->skip_conn_shutdown = true;
		}
	}

destroy_conn:
	slurm_persist_conn_destroy(*pc);
	*pc = NULL;
}

/*
 * Send an RPC to the SlurmDBD and wait for an arbitrary reply message.
 * The RPC will not be queued if an error occurs.
 * The "resp" message must be freed by the caller.
 * Returns SLURM_SUCCESS or an error code
 */
extern int dbd_conn_send_recv_direct(uint16_t rpc_version,
				     persist_msg_t *req,
				     persist_msg_t *resp)
{
	int rc = SLURM_SUCCESS;
	buf_t *buffer;
	persist_conn_t *use_conn = req->pcon;

	xassert(req);
	xassert(resp);
	xassert(use_conn);

	if (!use_conn->conn) {
		/* The connection has been closed, reopen */
		rc = dbd_conn_check_and_reopen(use_conn);

		if (rc != SLURM_SUCCESS || !use_conn->conn) {
			rc = SLURM_ERROR;
			goto end_it;
		}
	}

	if (!(buffer = pack_slurmdbd_msg(req, rpc_version))) {
		rc = SLURM_ERROR;
		goto end_it;
	}

	rc = slurm_persist_send_msg(use_conn, buffer);
	FREE_NULL_BUFFER(buffer);
	if (rc != SLURM_SUCCESS) {
		error("Sending message type %s: %d: %s",
		      slurmdbd_msg_type_2_str(req->msg_type, 1), rc,
		      slurm_strerror(rc));
		goto end_it;
	}

	buffer = slurm_persist_recv_msg(use_conn);
	if (buffer == NULL) {
		error("Getting response to message type: %s",
		      slurmdbd_msg_type_2_str(req->msg_type, 1));
		rc = SLURM_ERROR;
		goto end_it;
	}

	rc = unpack_slurmdbd_msg(resp, rpc_version, buffer);
	/* check for the rc of the start job message */
	if (rc == SLURM_SUCCESS && resp->msg_type == DBD_ID_RC)
		rc = ((dbd_id_rc_msg_t *)resp->data)->return_code;

	FREE_NULL_BUFFER(buffer);
end_it:

	log_flag(PROTOCOL, "msg_type:%s protocol_version:%hu return_code:%d response_msg_type:%s",
		 slurmdbd_msg_type_2_str(req->msg_type, 1),
		 rpc_version, rc, slurmdbd_msg_type_2_str(resp->msg_type, 1));

	return rc;
}

extern int dbd_conn_send_recv_rc_comment_msg(uint16_t rpc_version,
					     persist_msg_t *req,
					     int *resp_code,
					     char **comment)
{
	int rc;
	persist_msg_t resp;

	xassert(req);
	xassert(resp_code);

	memset(&resp, 0, sizeof(persist_msg_t));
	rc = dbd_conn_send_recv(rpc_version, req, &resp);
	if (rc != SLURM_SUCCESS) {
		;	/* error message already sent */
	} else if (resp.msg_type != PERSIST_RC) {
		error("response is not type PERSIST_RC: %s(%u)",
		      slurmdbd_msg_type_2_str(resp.msg_type, 1),
		      resp.msg_type);
		rc = SLURM_ERROR;
	} else {	/* resp.msg_type == PERSIST_RC */
		persist_rc_msg_t *msg = resp.data;
		*resp_code = msg->rc;

		if ((msg->ret_info == DBD_REGISTER_CTLD) &&
		    (msg->rc & RC_AS_CLUSTER_ID))
			msg->rc = SLURM_SUCCESS;

		if (msg->rc != SLURM_SUCCESS &&
		    msg->rc != ACCOUNTING_FIRST_REG &&
		    msg->rc != ACCOUNTING_TRES_CHANGE_DB &&
		    msg->rc != ACCOUNTING_NODES_CHANGE_DB) {
			char *comment = msg->comment;
			if (!comment)
				comment = slurm_strerror(msg->rc);
			if (!req->pcon &&
			    (msg->ret_info == DBD_REGISTER_CTLD) &&
			    slurm_conf.accounting_storage_enforce) {
				error("Issue with call "
				      "%s(%u): %u(%s)",
				      slurmdbd_msg_type_2_str(
					      msg->ret_info, 1),
				      msg->ret_info, msg->rc,
				      comment);
				fatal("You need to add this cluster "
				      "to accounting if you want to "
				      "enforce associations, or no "
				      "jobs will ever run.");
			} else
				debug("Issue with call "
				      "%s(%u): %u(%s)",
				      slurmdbd_msg_type_2_str(
					      msg->ret_info, 1),
				      msg->ret_info, msg->rc,
				      comment);
		}

		if (comment) {
			*comment = msg->comment;
			msg->comment = NULL;
		}

		slurm_persist_free_rc_msg(msg);
	}

	log_flag(PROTOCOL, "msg_type:%s protocol_version:%hu return_code:%d",
		 slurmdbd_msg_type_2_str(req->msg_type, 1),
		 rpc_version, rc);

	return rc;
}

extern int dbd_conn_send_recv_rc_msg(uint16_t rpc_version,
				     persist_msg_t *req,
				     int *resp_code)
{
	return dbd_conn_send_recv_rc_comment_msg(
		rpc_version, req, resp_code, NULL);
}

extern int dbd_conn_send_recv(uint16_t rpc_version,
			      persist_msg_t *req,
			      persist_msg_t *resp)
{
	if (running_in_slurmctld() &&
	    (!req->pcon || (req->pcon == slurmdbd_conn)))
		return slurmdbd_agent_send_recv(rpc_version, req, resp);
	else
		return dbd_conn_send_recv_direct(rpc_version, req, resp);
}
