/*
 * topsrv_tester.c - This program implements:
 * - logic to send subscriptions/cancellations both valid and
 *   invalid to topology server.
 * - Depends on flags specified it can wait/skip the response.
 * - If publisher flag is specified, it acts as a publisher in
 *   addition to being a subscriber.
 * It is purely used for stress testing the topology server.
 *
 * It implements:
 * - "N" subscriber threads sending requests.
 * - "1" receiver thread consumeing the topology server events.
 * - "1" main thread which spawns above threads and can publish
 *   services if specified.
 *
 */

#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <linux/tipc.h>
#include <sys/epoll.h>
#include <errno.h>
#include <string.h>
#include <unistd.h>
#include <pthread.h>
#include <fcntl.h>

#define OPT_SEND_INVALID_REQ 0x01
#define OPT_RESUBSCRIBE 0x02
#define OPT_PROCESS_EVENTS 0x04
#define OPT_BE_PUBLISHER_AND_CONTINUE 0x08
#define OPT_JUST_BE_PUBLISHER 0x10

#define MAX_EPOLL_EVENTS 10000

#define TOPSRV_TEST_TYPE 1000
#define SUBSCRIPTION_LOWER 0
#define SUBSCRIPTION_UPPER 0xffff

struct subscriber_stats {
	unsigned int subscription_withdraw;
	unsigned int subscription_publish;
	unsigned int subscription_timeout;
	unsigned int sent_req;
	unsigned int rcvd_rsp;
	unsigned int unknown;
};

struct connection_req {
	int fd;
	struct subscriber_stats stats;
};

unsigned int subscribers = 1;
unsigned int subscriptions = 1;
unsigned int timeout = TIPC_WAIT_FOREVER;
unsigned int flags = OPT_PROCESS_EVENTS;
unsigned int subscription_type = TOPSRV_TEST_TYPE;

pthread_mutex_t mutex;
pthread_cond_t subscribers_ready;
pthread_cond_t events_completed;
unsigned int subscribers_created;
int pending_events;

static void print_results(struct connection_req *tmp,
			  struct subscriber_stats *sum);

static void diep(const char *str)
{
	perror(str);
	exit(1);
}

static int socket_set_flag_nonblock(int fd)
{
	int flag;

	flag = fcntl (fd, F_GETFL, 0);
	if (flag == -1)
		diep("fcntl");

	flag |= O_NONBLOCK;
	if (fcntl (fd, F_SETFL, flag))
		diep("fcntl");

	return 0;
}

static int read_response(int fd, struct subscriber_stats *stats)
{
	unsigned int rsp_count = 0;
	int ret = 0;
	struct tipc_event event = { 0 };

	while (1) {
		ret = recv(fd, &event, sizeof(event),0);
		if ( ret < 0) {
			if (errno == EAGAIN)
				break;
			diep("recv");
		} else if (ret == 0) {
			fprintf(stderr, "Recv returned 0\n");
			continue;
		}

		rsp_count++;
		stats->rcvd_rsp++;
		switch (event.event) {
		case TIPC_PUBLISHED:
			stats->subscription_publish++;
			break;
		case TIPC_WITHDRAWN:
			stats->subscription_withdraw++;
			break;
		case TIPC_SUBSCR_TIMEOUT:
			stats->subscription_timeout++;
			break;
		default:
			stats->unknown++;
			break;
		}
	}
	return rsp_count;
}

static void *process_events(void *data)
{
	struct connection_req *req = data;
	int epfd, event_cnt, event_received;
	struct epoll_event ev_req, *ev_rsp, *ev_tmp;
	struct connection_req *tmp = req;
	int cnt;

	pthread_mutex_lock(&mutex);
	while (subscribers_created < subscribers) {
		pthread_cond_wait(&subscribers_ready, &mutex);
	}
#ifdef PTHREAD_DEBUG
	printf("pthread_cond_ok : subscribers_ready\n");
#endif
	pthread_mutex_unlock(&mutex);

	epfd = epoll_create(1);
	if (epfd == -1)
		diep("epoll_create");

	tmp = req;
	for (cnt = 0; cnt < subscribers; cnt++, tmp++) {
		socket_set_flag_nonblock(tmp->fd);
		ev_req.data.ptr = (struct connection_req *) tmp;
		ev_req.events = EPOLLIN | EPOLLET;
		if (epoll_ctl(epfd, EPOLL_CTL_ADD, tmp->fd, &ev_req) == -1)
			diep("epoll_ctl");
	}

	ev_rsp = malloc(MAX_EPOLL_EVENTS * sizeof(*ev_rsp));

	while (pending_events > 0) {
		event_cnt = epoll_wait(epfd, ev_rsp, MAX_EPOLL_EVENTS, -1);
		if (event_cnt <= 0)
			diep("epoll_wait");

		event_received = 0;
		ev_tmp = ev_rsp;
		for (cnt = 0; cnt < event_cnt; cnt++, ev_tmp++) {
			struct connection_req *rsp;
			struct subscriber_stats *stats;

			rsp = (struct connection_req *) ev_tmp->data.ptr;
			stats = &rsp->stats;

			if (ev_tmp->events & EPOLLIN) {
				event_received += read_response(rsp->fd, stats);
			} else if (ev_tmp->events & (EPOLLHUP | EPOLLERR)) {
				stats->unknown++;
				diep("Hangup received from TIPC_TOP_SRV\n");
			}
		}

#ifdef STAT_DEBUG
		struct subscriber_stats stat_sum;
		memset(&stat_sum, 0, sizeof(stat_sum));
		tmp = req;
		print_results(tmp, &stat_sum);
#endif
		pthread_mutex_unlock(&mutex);
		pending_events -= event_received;
		pthread_mutex_unlock(&mutex);
	}

	tmp = req;
	for (cnt = 0; cnt < subscribers; cnt++, tmp++) {
		ev_req.data.ptr = (struct connection_req *) tmp;
		ev_req.events = EPOLLIN;
		if (epoll_ctl(epfd, EPOLL_CTL_DEL, tmp->fd, NULL) == -1)
			diep("epoll_ctl");
	}

	pthread_mutex_lock(&mutex);
	pthread_cond_broadcast(&events_completed);
	pthread_mutex_unlock(&mutex);
#ifdef PTHREAD_DEBUG
	printf("pthread_cond_broadcast : events_completed\n");
#endif
	free(ev_rsp);
	close(epfd);
	pthread_exit(NULL);
}

static void create_subscriber(int *fd)
{
	*fd = socket(AF_TIPC, SOCK_SEQPACKET, 0);
	if (*fd < 0)
		diep("failed to create socket\n");
}

static void connect_socket(int fd)
{
	struct sockaddr_tipc topsrv;

	memset(&topsrv, 0, sizeof(topsrv));
	topsrv.family = AF_TIPC;
	topsrv.addrtype = TIPC_ADDR_NAME;
	topsrv.addr.name.name.type = TIPC_TOP_SRV;
	topsrv.addr.name.name.instance = TIPC_TOP_SRV;

	/* Connect to topology server: */
	if (0 > connect(fd, (struct sockaddr*)&topsrv, sizeof(topsrv)))
		diep("cannot connect to topology server\n");
}

static void create_subscription(int fd)
{
	struct tipc_subscr subscr;
	int *handle = (int *)&subscr.usr_handle;
	unsigned int cnt;

	memset(&subscr, 0, sizeof(subscr));
	subscr.seq.lower = SUBSCRIPTION_LOWER;
	subscr.seq.upper = SUBSCRIPTION_UPPER;
	subscr.timeout = timeout;
	subscr.filter = TIPC_SUB_PORTS;
	*handle = fd;
	for (cnt = 0; cnt < subscriptions; cnt++) {
		subscr.seq.type = subscription_type + cnt;
		if( send(fd, &subscr, sizeof(subscr), 0) != sizeof(subscr))
			diep("failed to send subscription\n");
	}
}

static void cancel_subscription(int fd, unsigned int filter)
{
	struct tipc_subscr cancel;
	int *handle = (int *)&cancel.usr_handle;
	unsigned int cnt;

	memset(&cancel, 0, sizeof(cancel));
	cancel.seq.lower = SUBSCRIPTION_LOWER;
	cancel.seq.upper = SUBSCRIPTION_UPPER;
	cancel.timeout = timeout;
	cancel.filter = filter;
	*handle = fd;
	for (cnt = 0; cnt < subscriptions; cnt++) {
		cancel.seq.type = subscription_type + cnt;
		if( send(fd, &cancel, sizeof(cancel), 0) != sizeof(cancel))
			diep("failed to cancel subscription\n");
	}
}

static void publish_subscriber_ready(void)
{
	pthread_mutex_lock(&mutex);
	subscribers_created++;
	if (subscribers_created == subscribers) {
#ifdef PTHREAD_DEBUG
		printf("subscribers_created : %d\n", subscribers_created);
#endif
		pthread_cond_signal(&subscribers_ready);
	}
	pthread_mutex_unlock(&mutex);
}

static void *subscribe(void *data)
{
	struct connection_req *req = data;

	create_subscriber(&req->fd);
	connect_socket(req->fd);
	create_subscription(req->fd);

	publish_subscriber_ready();
	req->stats.sent_req += subscriptions;

	if (flags & OPT_PROCESS_EVENTS) {
		pthread_mutex_lock(&mutex);
		while (pending_events > 0) {
			pthread_cond_wait(&events_completed, &mutex);
		}
		pthread_mutex_unlock(&mutex);
	}

	/* Send malformed request, this request should be ignored */
	if (flags & OPT_SEND_INVALID_REQ) {
		cancel_subscription(req->fd, TIPC_SUB_CANCEL);
	}

	cancel_subscription(req->fd, TIPC_SUB_PORTS|TIPC_SUB_CANCEL);

	if (flags & OPT_RESUBSCRIBE) {
		create_subscription(req->fd);
	}

	if(close (req->fd))
		diep("failed to close\n");

	pthread_exit(NULL);
}

static int start_subscribers(pthread_t *subscriber_thread,
			     struct connection_req *req)
{
	struct connection_req *tmp = req;
	int cnt;

	for (cnt=0; cnt < subscribers; tmp++, subscriber_thread++, cnt++) {
		if (pthread_create(subscriber_thread, NULL, subscribe, tmp))
			return -1;
	}
	return 0;
}

static int shutdown_subscribers(pthread_t *subscriber_thread)
{
	int cnt;

	for (cnt=0; cnt < subscribers; subscriber_thread++, cnt++) {
		if (pthread_join(*subscriber_thread, NULL))
			diep("pthread_join");
	}
	return 0;
}

static void publish_service(void)
{
	int fd;
	struct sockaddr_tipc sa;
	unsigned int cnt;

	fd = socket(AF_TIPC, SOCK_RDM, 0);

	sa.family = AF_TIPC;
	sa.addrtype = TIPC_ADDR_NAMESEQ;
	sa.addr.nameseq.type = subscription_type;
	sa.addr.nameseq.lower = 1;
	sa.addr.nameseq.upper = 20;
	sa.scope = TIPC_CLUSTER_SCOPE;

	for (cnt = 0; cnt < subscriptions; cnt++) {
		sa.addr.nameseq.type = subscription_type + cnt;
		if (0 != bind(fd, (struct sockaddr*)&sa, sizeof(sa)))
			diep("bind()");
	}
}

static void print_stat(unsigned int cnt, char *description)
{
	if (cnt)
		printf("\n%s:%u", description, cnt);
}

static void print_results(struct connection_req *tmp,
			  struct subscriber_stats *sum)
{
	int cnt;

	for (cnt=0; cnt < subscribers; tmp++, cnt++) {
		sum->rcvd_rsp += tmp->stats.rcvd_rsp;
		sum->sent_req += tmp->stats.sent_req;
		sum->subscription_publish += tmp->stats.subscription_publish;
		sum->subscription_withdraw += tmp->stats.subscription_withdraw;
		sum->subscription_timeout += tmp->stats.subscription_timeout;
		sum->unknown += tmp->stats.unknown;
	}

	print_stat(sum->sent_req, "Subscription Sent");
	print_stat(sum->rcvd_rsp, "Event Received");
	print_stat(sum->subscription_publish, "published");
	print_stat(sum->subscription_withdraw, "Withdrawn");
	print_stat(sum->subscription_timeout, "Timeout");
	print_stat(sum->unknown, "Unknown");
	fflush(stdout);
}

static int check_results(struct connection_req *tmp)
{
	struct subscriber_stats stat_sum;

	memset(&stat_sum, 0, sizeof(stat_sum));
	print_results(tmp, &stat_sum);

	if (stat_sum.rcvd_rsp < stat_sum.sent_req) {
		printf("\nTotal Req:%u Total Rsp:%u \n",
		       stat_sum.sent_req, stat_sum.rcvd_rsp);
		return -1;
	}
	if (flags & OPT_BE_PUBLISHER_AND_CONTINUE) {
		if (stat_sum.sent_req == stat_sum.subscription_publish)
			return 0;
		printf("\nTotal Req:%u Total Published:%u \n",
		       stat_sum.sent_req, stat_sum.subscription_publish);
		return -2;
	}
	if (stat_sum.sent_req != stat_sum.subscription_timeout) {
		printf("\nTotal Req:%u Total Timeout:%u \n",
		       stat_sum.sent_req, stat_sum.subscription_timeout);
		return -2;
	}
	return 0;
}

static void manage_workers(pthread_t *subscriber_thread,
			   struct connection_req *req)
{
	pthread_t event_thread;

	subscribers_created = 0;
	pthread_cond_init(&subscribers_ready, NULL);
	pthread_cond_init(&events_completed, NULL);
	pending_events = subscribers * subscriptions;

	start_subscribers(subscriber_thread, req);

	if (flags & OPT_PROCESS_EVENTS) {
		if (pthread_create(&event_thread, NULL, process_events, req))
			diep("pthread_create_event_thread");
		if (pthread_join(event_thread, NULL))
			diep("pthread_join");
	} else {
		pthread_mutex_lock(&mutex);
		while (subscribers_created < subscribers) {
			pthread_cond_wait(&subscribers_ready, &mutex);
		}
		pthread_cond_broadcast(&events_completed);
		pending_events = 0;
		pthread_mutex_unlock(&mutex);
	}

	shutdown_subscribers(subscriber_thread);
}

int32_t main(int32_t argc, char *argv[], char *envp[])
{
	unsigned int num_of_runs=1;
	struct connection_req *req;
	pthread_t *subscriber_thread;
	int result = 0;
	char c;

	while ((c = getopt(argc, argv, ":n:o:t:c:y:irehpw")) != -1)
	{
		switch(c) {
		case 'n':
			num_of_runs = atoi(optarg);
		break;
		case 'o':
			subscribers = atoi(optarg);
		break;
		case 't':
			timeout = atoi(optarg);
		break;
		case 'c':
			subscriptions = atoi(optarg);
		break;
		case 'i':
			flags |= OPT_SEND_INVALID_REQ;
		break;
		case 'r':
			flags |= OPT_RESUBSCRIBE;
		break;
		case 'e':
			flags &= ~OPT_PROCESS_EVENTS;
		break;
		case 'p':
			flags |= OPT_BE_PUBLISHER_AND_CONTINUE;
		break;
		case 'y':
			subscription_type = atoi(optarg);
		break;
		case 'w':
			if (flags & OPT_BE_PUBLISHER_AND_CONTINUE) {
				printf("\nWarning: ignoring option -p as -w is specified\n");
			}
			flags |= OPT_JUST_BE_PUBLISHER;
		break;
		case 'h':
		default:
			printf("%s [options] [flags]\n", argv[0]);
			printf("  options:\n"
			       "	-n <num_of_attempts>\n"
			       "	-y <subscription_type>\n"
			       "	-o <open_subscribers>\n"
			       "	-c <subscriptions_per_subscriber>\n"
			       "	-t <subscription_timeout>\n");
			printf("  flags:\n"
			       "	-h \"print this help string\"\n"
			       "	-i \"send invalid request\"\n"
			       "	-r \"resend subscription after cancel\"\n"
			       "	-e \"ignore events from topology server\"\n"
			       "	-p \"Publish the required service - bind to port and continue\"\n"
			       "	-w \"Just Publish the required service - bind to port and wait\"\n");
			exit(1);
		break;
		}
	}


	if (flags & OPT_JUST_BE_PUBLISHER) {
		publish_service();
		while(1);
	}

	if (flags & OPT_BE_PUBLISHER_AND_CONTINUE)
		publish_service();

	req = malloc(subscribers * sizeof(*req));
	subscriber_thread = malloc(subscribers * sizeof(pthread_t));
	memset(req, 0, subscribers * sizeof(*req));

	while (num_of_runs) {
		manage_workers(subscriber_thread, req);
		num_of_runs--;
		if ((num_of_runs % 25000) == 0) {
			printf(".");
			fflush(stdout);
		}
	}

	if (flags & OPT_PROCESS_EVENTS)
		result = check_results(req);

	printf("\nTIPC Topology server:%s\n", result ? "NOK" : "OK");

	pthread_mutex_destroy(&mutex);
	pthread_cond_destroy(&subscribers_ready);
	pthread_cond_destroy(&events_completed);
	free(subscriber_thread);
	free(req);

	return result;
}
