/* 
 *  rocks/hb.c
 *
 *  The heartbeat connection failure detection mechanism.
 *
 *  Copyright (C) 2001 Victor Zandy
 *  See COPYING for distribution terms.
 */

/*
  Heartbeats run over a separate connection.  There would be numerous
  problems running heartbeats over the data connection:
      - Some systems, such as Linux, copy old out-of-band urgent data
        in the receive queue into the data stream if it is not consumed
        soon enough.
      - Neither urgent data nor urgent notification are sent when the
        send queue is full on Linux.
      - Race-free implementations are complicated.  */
#include <sys/time.h>
#include <sys/types.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <errno.h>
#include <assert.h>

#include "rs.h"
#include "log.h"

struct hb {
	int period;		/* Beat period (multiple of alarm period) */
	int count;		/* Beats that have occurred within period */
	int limit;		/* Missed beats limit */
	int missed;		/* Current number of missed beats */
	int s;   		/* Connection to peer hb socket (plain) */
	rs_t rs;		/* Pointer to owning rock */
	int lshm;		/* Last count seen in shm */
	int ushm;		/* Beats in which shm has not changed */
};

/* These are NULL unless there is a HB for the descriptor. */
static hb_t hbs[RS_MAXFD];

hb_t
rs_new_heartbeat(rs_t rs)
{
	hb_t hb;
	hb = (hb_t) malloc(sizeof(struct hb));
	if (!hb)
		return NULL;
	hb->limit = rs_opt_max_alarm_misses;
	hb->period = 1;
	hb->count = 0;
	hb->missed = 0;
	hb->s = -1;
	hb->rs = rs;
	hb->lshm = 0;
	hb->ushm = 0;
	return hb;
}

int
rs_hb_save(hb_t hb, int fd)
{
	if (0 > rs_xwrite(fd, hb, sizeof(*hb)))
		return -1;
	return 0;
}

hb_t
rs_hb_restore(rs_t rs, int fd)
{
	struct hb h;
	hb_t hb;

	if (0 > rs_xread(fd, &h, sizeof(h), 0))
		return NULL;
	hb = rs_new_heartbeat(NULL);
	if (!hb)
		return NULL;
	*hb = h;
	rs->hb = hb;
	hb->rs = rs;
	hbs[hb->s] = hb;
	return hb;
}

void
rs_free_heartbeat(hb_t hb)
{
	rs_mode_native();
	close(hb->s);
	rs_mode_pop();
	free(hb);
}

void
rs_become_hb_owner(hb_t hb)
{
	if (0 > fcntl(hb->s, F_SETOWN, getpid()))
		assert(0);
}

static hb_t
hb_lookup(int s)
{
	if (s >= RS_MAXFD || s < 0)
		return NULL;
	return hbs[s];
}

/* ROCK is established data connection descriptor.
   HB is allocated, initialized HB structure.
   Establish HB connection and update HB accordingly.
   ROLE is used to decide ties. */
int
rs_hb_establish(int rock, hb_t hb, rs_role role)
{
	int rv, len, ls;
	struct sockaddr_in laddr, paddr;

	rs_mode_native();

	/* Bind a new listening socket with IP address of ROCK */
	len = sizeof(laddr);
	if (0 > getsockname(rock, (struct sockaddr *)&laddr, &len))
		goto out_noclose;
	ls = socket(AF_INET, SOCK_STREAM, 0);
	if (0 > ls)
		goto out;
	len = sizeof(laddr);
	laddr.sin_port = htons(0);
	rv = bind(ls, (struct sockaddr *) &laddr, sizeof(laddr));
	if (0 > rv)
		goto out;

	/* Non blocking avoids timing problem described in
	   Stevens, Network Programming 1, 15.6 */
	if (0 > rs_nonblock(ls, 1))
		goto out;

	/* Exchange socket addresses with peer - network order */
	len = sizeof(laddr);
	if (0 > getsockname(ls, (struct sockaddr *)&laddr, &len))
		goto out;
	if (0 >= rs_xwrite(rock, &laddr, sizeof(laddr)))
		goto out;
	if (0 >= rs_xread(rock, &paddr, sizeof(paddr), 0))
		goto out;

	hb->s = rs_1of2(&laddr, &paddr, ls, NULL, role);
	assert(hb->s >= 0);
	close(ls);
	hbs[hb->s] = hb;
	rs_become_hb_owner(hb);
	rs_mode_pop();
	rs_log("Heartbeat connection established for <%p>.", hb->rs);
	return 0;
out:
	close(ls);
out_noclose:
	rs_mode_pop();
	rs_log("Cannot establish heartbeat connection for <%p>.", hb->rs);
	return -1;
}

int
rs_hb_cancel(hb_t hb)
{
	if (hb->s >= 0) {
		hbs[hb->s] = NULL;
		close(hb->s);
	}
	hb->s = -1;
	hb->count = 0;
	hb->missed = 0;
	return 0;
}

void
rs_hb_init_shm(rs_t rs)
{
	rs->shm->hb_owner = getpid();
	rs->shm->hb_count = 0;
	rs->hb->lshm = 0;
	rs->hb->ushm = 0;
}

static pid_t
hb_owner(hb_t hb)
{
	if (!rs_rock_is_shared(hb->rs))
		return getpid();
	else
		return hb->rs->shm->hb_owner;
}

static int
hb_ping_owner(hb_t hb)
{
	if (hb->lshm == hb->rs->shm->hb_count) {
		hb->ushm++;
		if (hb->ushm > 2)
			return 0;
		else
			return 1;
	}
	hb->lshm = hb->rs->shm->hb_count;
	hb->ushm = 0;
	return 1;
}

static int
hb_takeover(hb_t hb)
{
	rs_shm_lock(hb->rs->shm);
	hb->rs->shm->hb_owner = getpid();
	hb->rs->shm->hb_count++; /* notifies others that hb is okay */
	rs_become_hb_owner(hb);
	rs_log("new owner of HB <%p> (%d)", hb->rs, hb->rs->sd);
	rs_shm_unlock(hb->rs->shm);
	return 0;
}

/* Return 0 if sent hb okay.  Return -1 if the connection went away. */
static int
send_hb(int sd)
{
	int rv, e;
retry:
	rv = send(sd, "x", 1, MSG_OOB);
	if (rv == 1)
		return 0;
	assert(rv != 0);
	e = errno;
	switch(e) {
	case EINTR:
		goto retry;
		break;
	case EAGAIN:
	case ENOSPC:
	case EIO:
	case EFAULT:
		/* These shouldn't happen */
		assert(0);
		break;
	case EBADF:
	case EINVAL:  /* Linux can do this if you remove the nic */
	default:
 	        /* The connection went away */
		break;
	}
	return -1;
}

static void
handle_sigalrm(int sig)
{
	fd_set fdset;
	rs_t rs;
	int max;
	int sd;
	pid_t mypid;

	rs_mode_native();
	mypid = getpid();
	FD_ZERO(&fdset);
	max = rs_fdset(&fdset);
	for (sd = 0; sd < max; sd++) {
		if (!FD_ISSET(sd, &fdset))
			continue;
		rs = rs_lookup(sd);
		assert(rs);

		/* Update shared heartbeats; even for suspended rocks */
		if (rs_rock_is_shared(rs)) {
			if (mypid != hb_owner(rs->hb)) {
				if (!hb_ping_owner(rs->hb))
					hb_takeover(rs->hb);
				continue; /* non-owner does not touch hb */
			} else if (rs_shm_has_one_owner(rs))
				/* garbage collect */
				rs_shm_detach(rs);
			else
				rs->shm->hb_count++;
		}
		if (RS_SUSPENDED == rs->state)
			continue;

		/* Count beats toward this heartbeat's period */
		rs->hb->count++;
		if (rs->hb->count < rs->hb->period)
			continue;
		
		/* Period expired */
		rs->hb->count = 0;
		rs->hb->missed++;
		
		/* Send hb */
		if (0 > send_hb(rs->hb->s)) {
			rs_log("reconnect for hb send failure");
			rs_reconnect(rs, RS_NOBLOCK);
			continue;
		}
		
		/* Check missed hb */
		if (rs->hb->missed >= rs->hb->limit) {
			rs_log("reconnect for too many missed hbs");
			rs_reconnect(rs, RS_NOBLOCK);
		} else if (rs->hb->missed > 1)
			rs_log("missed heartbeat <%p> (suspending after %d more)",
			       rs, rs->hb->limit - rs->hb->missed);
	}
	rs_mode_pop();
}

static void
handle_sigurg(int sig)
{
	fd_set fdr, fde, fds;
	fd_set fdrx, fdex;
	int max, maxh, n, sd, rv;
	struct timeval tv;
	char x;
	
	rs_mode_native();

retry_all:
	FD_ZERO(&fdrx);
	FD_ZERO(&fdex);
	FD_ZERO(&fds);
	max = rs_fdset(&fds);

	maxh = 0;
	for (sd = 0; sd < max; sd++)
		if (FD_ISSET(sd, &fds)) {
			int x;
			rs_t rs = rs_lookup(sd);
			assert(rs);
			if (RS_SUSPENDED == rs->state) {
				FD_SET(rs->sd, &fdrx);
				x = rs->sd;
				rs_log("chking %d in read to select call",
				       rs->sd);
			} else {
				FD_SET(rs->hb->s, &fdex);
				x = rs->hb->s;
			}
			if (x > maxh)
				maxh = x;
		}
	maxh++;
retry_select:
	memcpy(&fdr, &fdrx, sizeof(fdrx));
	memcpy(&fde, &fdex, sizeof(fdex));
	tv.tv_sec = 0;
	tv.tv_usec = 0;
	n = select(maxh, &fdr, NULL, &fde, &tv);
	if (0 > n) {
		if (errno == EBADF
		    && 0 > rs_recover_bad_rocks(maxh, &fds))
			assert(0);
		goto retry_select;
	}

	for (sd = 0; n > 0 && sd < maxh; sd++) {
		if (FD_ISSET(sd, &fdr)) {
			rs_t rs;
			rs = rs_lookup(sd);
			assert(rs);
			assert(rs->state == RS_SUSPENDED);
			
			/* finished reconnecting */
			rs_rec_complete(rs, RS_NOBLOCK);
			continue;
		}
		if (FD_ISSET(sd, &fde)) {
			hb_t hb;
			hb = hb_lookup(sd);
			assert(hb);
			/* heartbeat received */
			rv = recv(hb->s, &x, 1, MSG_OOB);
			if (0 > rv) {
				rs_log("reconnect for hb recv failure");
				rs_reconnect(hb->rs, RS_NOBLOCK);
				/* Descriptors in fdseth may be invalid
				   after reconnection. */
				goto retry_all;
			} else
				hb->missed = 0;
			n--;
		}
	}
	rs_mode_pop();
}

void 
rs_stop_heartbeat(sigset_t *oset)
{
	sigset_t set;
	sigemptyset(&set);
	sigaddset(&set, SIGALRM);
	sigprocmask(SIG_BLOCK, &set, oset);
}

void
rs_resume_heartbeat(sigset_t *oset)
{
	sigprocmask(SIG_SETMASK, oset, NULL);
}

int
rs_init_heartbeat()
{
	struct sigaction sa;
	struct itimerval it;
	int ret = 0;

	memset(&sa, 0, sizeof(sa));
	sigfillset(&sa.sa_mask);
	sigdelset(&sa.sa_mask, SIGTERM);
	sigdelset(&sa.sa_mask, SIGINT);
	sa.sa_flags = SA_RESTART;
	sa.sa_handler = handle_sigurg;
	if (0 > sigaction(SIGURG, &sa, NULL))
		return -1;
	sa.sa_handler = handle_sigalrm;
	if (0 > sigaction(SIGALRM, &sa, NULL))
		return -1;

	it.it_value.tv_sec = it.it_interval.tv_sec = rs_opt_alarm_period;
	it.it_value.tv_usec = it.it_interval.tv_usec = 0;
	rs_mode_native();
	ret = setitimer(ITIMER_REAL, &it, NULL);
	rs_mode_pop();
	return ret;
}

int
rs_setitimer(int which, const struct itimerval *value, struct itimerval *ov)
{
	return 0;
}