/* ------------------------------------------------------------------------
 *
 * worker.c
 *
 * Short description: 
 *
 * ------------------------------------------------------------------------
 *
 * Copyright (c) 2017, Ericsson Canada Inc
 * All rights reserved.
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * Redistributions of source code must retain the above copyright notice, this
 * list of conditions and the following disclaimer.
 * Redistributions in binary form must reproduce the above copyright notice,
 * this list of conditions and the following disclaimer in the documentation
 * and/or other materials provided with the distribution.
 * Neither the name of the copyright holders nor the names of its
 * contributors may be used to endorse or promote products derived from this
 * software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
 * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
 * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
 * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
 * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
 * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
 * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
 * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
 * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
 * POSSIBILITY OF SUCH DAMAGE.
 *
 * ------------------------------------------------------------------------
 */

#define _GNU_SOURCE
#include <sys/types.h>
#include <sys/socket.h>
#include <sys/wait.h>
#include <sys/sysinfo.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/param.h>
#include <errno.h>
#include <unistd.h>
#include <netinet/in.h>
#include <tipcc.h>
#include <sys/epoll.h>
#include <sys/time.h>
#include <sys/prctl.h>
#include <assert.h>
#include <common.h>
#include <sched.h>
#include <time.h>

static char buf[TIPC_MAX_USER_MSG_SIZE];
int vlevel = 0;

struct context {
	int efd;
	int cfd;
	int tfd;
	int receivers;
	int msg_len;
	bool loopback;
	char mstr[64];
	uint32_t node;

	/* Statistics */
        struct timeval start_intv;
	uint32_t stats_intv;
	uint32_t prev_snt;
	uint32_t snt_uc;
	uint32_t snt_ac;
	uint32_t snt_mc;
	uint32_t snt_bc;
	uint32_t prev_rcvd;
	uint32_t rcv_uc;
	uint32_t rcv_ac;
	uint32_t rcv_mc;
	uint32_t rcv_bc;
	uint32_t rcv_first;
	bool stats;

	/* Transmitter parameters */
	int curr_test;
	uint32_t snd_nxt;
	struct tipc_addr dest;
	struct tipc_addr dest_sockid;

	/* Receiver parameters */
	int check_seqno;
	uint32_t rcv_nxt;
	int error;
};

static void bind_to_cpu(bool randomize)
{
	int num_cpus, upper_cpu;
        cpu_set_t cpumask;
	int cpu = 0;

	CPU_ZERO(&cpumask);
	srand(time(NULL) ^ (getpid() << 16));

	if (randomize) {
		num_cpus = get_nprocs();
		upper_cpu = num_cpus - 1;
		if (!upper_cpu)
			return;
		do {
			cpu = rand() / (RAND_MAX / num_cpus);
		} while (!cpu || cpu > upper_cpu);
	}
        CPU_SET(cpu, &cpumask);
        if (sched_setaffinity(0, sizeof(cpumask), &cpumask) == -1)
		die("sched_setaffinity");
}

static const char *msgstr(int mtyp)
{
	switch(mtyp) {
	case UNICAST:
		return "Unicast";
	case ANYCAST:
		return "Anycast";
	case MULTICAST:
		return "Multicast";
	case BROADCAST:
		return "Broadcast";
	default:
		return "Unknown Msg Type";
	}
}

static void report_result(struct context *ctx)
{
	struct tipc_addr commander = {COMMAND_GROUP, COMMANDER_ID, ctx->node};
	struct report_hdr hdr;
	int rc;

	hdr.error = htonl(ctx->error);
	hdr.rcv_first = htonl(ctx->rcv_first);
	hdr.rcv_last = htonl(ctx->rcv_nxt - 1);
	if (ctx->rcv_nxt == 0)
		hdr.rcv_last = 0;
	hdr.rcv_uc = htonl(ctx->rcv_uc);
	hdr.rcv_ac = htonl(ctx->rcv_ac);
	hdr.rcv_mc = htonl(ctx->rcv_mc);
	hdr.rcv_bc = htonl(ctx->rcv_bc);
	hdr.snt_uc = htonl(ctx->snt_uc);
	hdr.snt_ac = htonl(ctx->snt_ac);
	hdr.snt_mc = htonl(ctx->snt_mc);
	hdr.snt_bc = htonl(ctx->snt_bc);
	strcpy(hdr.mbrstr, ctx->mstr);
	rc = tipc_sendto(ctx->cfd, &hdr, sizeof(hdr), &commander);
	if (rc != sizeof(hdr))
		die("Failed to send report to commander\n");
}

static void print_stats(struct context *ctx)
{
	struct timeval now, *start_intv = &ctx->start_intv;
        unsigned long long from, to, elapsed;
	unsigned long long msg_per_sec;
        unsigned long long thruput;
	unsigned long long snt, rcvd;
	char str[64];

	if (!ctx->stats && (vlevel < 2))
		return;

	gettimeofday(&now, 0);
        from = (start_intv->tv_sec * 1000000 + start_intv->tv_usec) / 1000;
        to = (now.tv_sec * 1000000 + now.tv_usec) / 1000;
	elapsed = to - from;
	if ((elapsed / 1000) < ctx->stats_intv)
		return;

	ctx->start_intv = now;
	sprintf(str, "%s: ", ctx->mstr);

	if (ctx->curr_test) {
		snt = ctx->snt_bc + ctx->snt_mc + ctx->snt_ac + ctx->snt_uc;
		msg_per_sec = ((snt - ctx->prev_snt) * 1000) / elapsed;
		thruput = (msg_per_sec * ctx->msg_len * 8) / 1000000;
		ctx->prev_snt = snt;
		printf("%sSent UC %u, AC %u, MC %u, BC %u, "
		       "throughput last intv %llu Mb/s\n",
		       str, ctx->snt_uc, ctx->snt_ac, ctx->snt_mc,
		       ctx->snt_bc, thruput);
		memset(str, 0x20, strlen(str));
	}

	if (!ctx->curr_test || ctx->loopback) {
		rcvd = ctx->rcv_bc + ctx->rcv_mc + ctx->rcv_ac + ctx->rcv_uc;
		msg_per_sec = ((rcvd - ctx->prev_rcvd) * 1000) / elapsed;
		thruput = (msg_per_sec * ctx->msg_len * 8) / 1000000;
		ctx->prev_rcvd = rcvd;
		printf("%sRecv UC %u, AC %u, MC %u, BC %u, "
		       "throughput last intv %llu Mb/s\n",
		       str, ctx->rcv_uc, ctx->rcv_ac, ctx->rcv_mc,
		       ctx->rcv_bc, thruput);
	}
}

void join_group(int *epollfd, int *sockfd, struct tipc_addr *memberid,
		struct tipc_addr *sockid, int events, bool evts, bool loopback)
{
	int efd = *epollfd;
	int sfd;
	struct epoll_event evt;

	if (!efd) {
		efd = epoll_create(1);
		if (efd == -1)
			die("Failed to create epoll object\n");
	}
	sfd = tipc_socket(SOCK_RDM);
	if (sfd < 0)
		die("Failed to create member socket\n");
	tipc_sock_non_block(sfd);

	evt.data.fd = sfd;
	evt.events = events;

	if (epoll_ctl(efd, EPOLL_CTL_ADD, sfd, &evt) == -1)
		die("epoll_mod failure");

	if (tipc_join(sfd, memberid, evts, loopback))
		die("Failed to join group\n");
	tipc_sockaddr(sfd, sockid);

	*epollfd = efd;
	*sockfd = sfd;
}

static void rcv_cmd_msg(struct context *ctx)
{
	struct tipc_addr memberid = {TRAFFIC_GROUP, RECEIVER_ID1, ctx->node};
	struct tipc_addr srcid;
	struct tipc_addr sockid;
	struct command_hdr hdr;
	char str[32], str2[32];
	int cmd;
	int rc, err;

	rc = tipc_recvfrom(ctx->cfd, &hdr, sizeof(hdr), 0, &srcid, &err);

	if (rc < 0)
		die("recvfrom() error on command bus\n");

	/* Membership events */
	if (rc == 0) {
		if (srcid.instance != COMMANDER_ID || !err)
			return;
		dbg1("Worker %s: Lost contact with commander, exiting...\n",
		     ctx->mstr);
		exit(EXIT_SUCCESS);
	}

	/* Receiver Commands */
	cmd = ntohl(hdr.cmd);

	if (cmd == JOIN) {
		/* Some multicast receivers will receive duplicate requests */
		if (ctx->tfd)
			return;
		memberid.instance = ntohl(hdr.memberid);
		ctx->msg_len = ntohl(hdr.msg_len);
		ctx->stats_intv = ntohl(hdr.stats_intv);
		ctx->loopback = hdr.loopback;
		join_group(&ctx->efd, &ctx->tfd, &memberid, &sockid,
			   EPOLLIN | EPOLLOUT, hdr.events, hdr.loopback);
		tipc_ntoa(&memberid, str, sizeof(str));
		tipc_ntoa(&sockid, str2, sizeof(str2));
		sprintf(ctx->mstr, "%s/%s", str, str2);
		dbg1("Worker: Joined Traffic Group as %s, evt/loop %u/%u\n",
		     ctx->mstr, hdr.events, hdr.loopback);
		ctx->error = NOT_RECEIVED;
		return;
	}
	if (cmd == LEAVE) {
		report_result(ctx);
		dbg1("Worker %s: Exiting\n", ctx->mstr);
		exit(EXIT_SUCCESS);
	}

	/* Transmitter Commands */ 
	dbg1("Transmitting Worker: Received Command %s\n", cmdstr(cmd));
	ctx->curr_test = cmd;
	ctx->msg_len = ntohl(hdr.msg_len);
	ctx->stats_intv = ntohl(hdr.stats_intv);
	ctx->dest.type = TRAFFIC_GROUP;
	ctx->dest.instance = ntohl(hdr.memberid);
	ctx->dest.node = ctx->node;
	ctx->error = NOT_SENT;
	gettimeofday(&ctx->start_intv, 0);
}

/* rcv_traffic_msgs(): non-blocking receiption of messages until EAGAIN
 */
static void rcv_traffic_msgs(struct context *ctx)
{
	struct traffic_hdr *hdr = (void*) buf;
	struct tipc_addr sockid, memberid;
	int seqno, mtyp;
	int rc, err;
	int i = 0;

	do {
		rc = tipc_recvfrom(ctx->tfd, buf, sizeof(buf),
				   &sockid, &memberid, &err);

		if (rc < 0) {
			if (errno == EAGAIN)
				return;
			die("recvfrom() failure");
		}

		/* Membership event */
		if (!rc) {
			if (!err)
				ctx->receivers++;
			else
				ctx->receivers--;
			if (!err && ctx->receivers == 1)
				memcpy(&ctx->dest_sockid, &sockid, sizeof(sockid));
			if (err && ctx->receivers == 0)
				ctx->curr_test = NONE;
			dbg2("Transmitting Worker: Traffic group %s evt, now %u receivers\n",
			     err ? "DOWN" : "UP", ctx->receivers);
			continue;
		}

		/* Traffic Message */
	        mtyp = ntohl(hdr->type);
		seqno = ntohl(hdr->seqno);
		ctx->check_seqno |= hdr->check_seqno;
		ctx->error = err;
		if (memberid.type != TRAFFIC_GROUP)
			ctx->error = BAD_SRC_ADDR;
		else if ((memberid.instance != TRANSMITTER_ID) &&
			 (memberid.instance != RECEIVER_ID1) &&
			 (memberid.instance != RECEIVER_ID2))
			ctx->error = BAD_SRC_ADDR;
		else if (hdr->check_length && rc != ctx->msg_len)
			ctx->error = BAD_LENGTH;
                
		if (mtyp == UNICAST)
			ctx->rcv_uc++;
                else  if (mtyp == ANYCAST)
			ctx->rcv_ac++;
                else  if (mtyp == MULTICAST)
			ctx->rcv_mc++;
                else if (mtyp == BROADCAST)
			ctx->rcv_bc++;
                else
			ctx->error = BAD_TYPE;

		if (!ctx->rcv_nxt) {
			ctx->rcv_first = seqno;
			ctx->rcv_nxt = seqno;
		}
		if (ctx->check_seqno && seqno != ctx->rcv_nxt) {
			dbg1("Worker: Rcv msg with bad seqno %u, expected %u\n",
			     seqno, ctx->rcv_nxt);
			ctx->error = BAD_SEQNO;
		}
		ctx->rcv_nxt = seqno + 1;

		dbg3("Worker %s: Received %s msg #%u of len %u\n",
		     ctx->mstr, msgstr(mtyp), seqno, rc);
		if (!ctx->error)
			continue;

		report_result(ctx);
		die("%s: Reporting Error %s, Exiting...\n", ctx->mstr,
		    errstr(ctx->error));
        } while (i++ < 100);
}

/* snd_traffic_msgs(): send non-blocking messages until EAGAIN
 */
static void snd_traffic_msgs(struct context *ctx)
{
	struct traffic_hdr *hdr = (void*) buf;
	struct tipc_addr *dst = &ctx->dest;
	struct tipc_addr *dstsock = &ctx->dest_sockid;
	char str[60];
	int sd = ctx->tfd;
        int test = ctx->curr_test;
	int len, rc = 0;
	int i = 0;

	if (test == NONE)
		return;

        do {
		len = ctx->msg_len;
		hdr->seqno = htonl(ctx->snd_nxt);
		hdr->check_seqno = ~0;
		hdr->check_length = ~0;

		switch(test) {
                case UNICAST_TEST:
			tipc_ntoa(dst, str, sizeof(str));
			hdr->type = htonl(UNICAST);
                        rc = tipc_sendto(sd, buf, len, dstsock);
                        if (rc == len)
                                ctx->snt_uc++;
			break;
                case ANYCAST_TEST:
                case MULTIPOINT_TEST:
			hdr->type = htonl(ANYCAST);
			hdr->check_seqno = 0;
                        rc = tipc_sendto(sd, buf, len, dst);
                        if (rc == len)
                                ctx->snt_ac++;
			break;
                case MULTICAST_TEST:
			hdr->type = htonl(MULTICAST);
                        rc = tipc_mcast(sd, buf, len, dst);
                        if (rc == len)
                                ctx->snt_mc++;
			break;
                case BROADCAST_TEST:
			hdr->type = htonl(BROADCAST);
                        rc = tipc_send(sd, buf, len);
                        if (rc == len)
                                ctx->snt_bc++;
			break;
                case UC_AFT_BC_SEQ_TEST:
			/* Cycle: 7 large broadcasts, 1 small unicast */
			if (ctx->snd_nxt % 8) {
				hdr->type = htonl(BROADCAST);
				hdr->check_seqno = 0;
				rc = tipc_send(sd, buf, len);
				if (rc == len)
					ctx->snt_bc++;
			} else {
				hdr->type = htonl(UNICAST);
				hdr->check_length = 0;
				len = 100;
				rc = tipc_sendto(sd, buf, len, dstsock);
				if (rc == len)
					ctx->snt_uc++;
			}
			break;
                case BC_AFT_UC_SEQ_TEST:
			/* Cycle: 6 large unicasts, 2 small broadcasts */
			if (((ctx->snd_nxt + 1) % 8) && (ctx->snd_nxt % 8)) {
				hdr->type = htonl(UNICAST);
				rc = tipc_sendto(sd, buf, len, dstsock);
				if (rc == len)
					ctx->snt_uc++;
			} else {
				hdr->type = htonl(BROADCAST);
				len = 100;
				hdr->check_seqno = 0;
				hdr->check_length = 0;
				rc = tipc_send(sd, buf, len);
				if (rc == len)
					ctx->snt_bc++;
			}
			break;
                default:
			break;
                }
			
		if (rc != len)
			break;

		ctx->snd_nxt++;
		ctx->error = OK;
        } while (i++ < 100);
}

void worker_main(uint32_t instance, bool stats, uint32_t node)
{
	struct tipc_addr memberid = {COMMAND_GROUP, instance, node};
	struct tipc_addr sockid;
	struct epoll_event events[4];
	struct context ctx = {0,};
	int revents, rv, i, fd;
	char str[60], str2[60];
	pid_t pid;

	pid = fork();
	if (pid < 0)
		die("Worker: failed to create process\n");
	if (pid)
		return;

	prctl(PR_SET_PDEATHSIG, SIGTERM);
	bind_to_cpu(!stats);
	ctx.node = node;
	ctx.stats = stats;
	ctx.stats_intv = 100000000;
	join_group(&ctx.efd, &ctx.cfd, &memberid,
		   &sockid, EPOLLIN, true, false);
	tipc_ntoa(&memberid, str, sizeof(str));
	tipc_ntoa(&sockid, str2, sizeof(str2));
	gettimeofday(&ctx.start_intv, 0);
	dbg1("Worker: Joined Command Group as %s/%s\n", str, str2);

	while (1) {
		rv = epoll_wait(ctx.efd, events, 4, -1);
		if (rv <= 0)
			die("epoll() failed\n");

		for (i = 0; i < rv; i++) {
			fd = events[i].data.fd;
			revents = events[i].events;
			if (fd == ctx.tfd) {
				if (revents & EPOLLIN)
					rcv_traffic_msgs(&ctx);
				if (revents & EPOLLOUT)
					snd_traffic_msgs(&ctx);
			} else if (fd == ctx.cfd) {
				if (revents & EPOLLIN)
					rcv_cmd_msg(&ctx);
				if (revents & EPOLLOUT)
					die("EPOLLOUT on command socket\n");
			}
		}
		print_stats(&ctx);
	}
	exit(EXIT_SUCCESS);
}
