/*
 * idleconn -- a tool for opening any number of idle connections
 * Copyright (C) 2006  Henrik Nordstrom <henrik@henriknordstrom.net>
 * 
 * Some code may be based on idleconn from httpperf
 *
 * httperf -- a tool for measuring web server performance
 * Copyright (C) 2000  Hewlett-Packard Company
 * Contributed by David Mosberger-Tang <davidm@hpl.hp.com>
 * 
 * This program 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.
 * 
 * This program 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 this program; if not, write to the Free Software
 * Foundation, Inc., 59 Temple Place, Suite 330, Boston, MA
 * 02111-1307 USA
 */

#include <errno.h>
#include <netdb.h>
#include <signal.h>
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <fcntl.h>
#include <unistd.h>

#include <netinet/in.h>
#include <arpa/inet.h>

#include <sys/time.h>
#include <sys/types.h>
#include <sys/resource.h>
#include <sys/socket.h>

#include <event.h>

#define UNUSED __attribute__((unused))

const char *prog_name;

static struct fhandle {
    int open;
    struct event ev;
    struct fhandle *next;
}      *fhandles;
static struct fhandle *free_fhandle;
static int desired_connections;
static int current_connections, pending_connections;
struct sockaddr_in server_addr;
int error;

static struct sockaddr_in *my_addr = NULL;
static int nr_my_addr = 0;
static int last_addr = 0;
static unsigned short last_port = 1024;

static void
add_my_addr(char *ip)
{
    struct sockaddr_in *sin;
    nr_my_addr++;
    my_addr = realloc(my_addr, nr_my_addr * sizeof(*my_addr));
    sin = &my_addr[nr_my_addr - 1];
    memset(sin, 0, sizeof(*sin));
    sin->sin_family = AF_INET;
    inet_aton(ip, &sin->sin_addr);
}

static void
setlocalport(int sd)
{
    int i;
    if (!nr_my_addr)
	return;
    for (i = 0; i < 100; i++) {
	struct sockaddr_in *sin = &my_addr[last_addr++];
	if (last_addr >= nr_my_addr) {
	    last_addr = 0;
	    last_port++;
	    if (!last_port)
		last_port = 1024;
	}
	sin->sin_port = htons(last_port);
	if (bind(sd, (struct sockaddr *)sin, sizeof(*sin)) == 0)
	    break;
    }
}

static struct fhandle *
get_fhandle()
{
    struct fhandle *fh = free_fhandle;
    free_fhandle = fh->next;
    fh->next = NULL;
    fh->open = 0;
    return fh;
}

static void
put_fhandle(struct fhandle *fh)
{
    fh->next = free_fhandle;
    free_fhandle = fh;
}

static void
init(const char *server, const char *_port, const char *_desired)
{
    struct rlimit rlimit;
    struct hostent *he;
    int i;
    short port = atoi(_port);

    desired_connections = atoi(_desired);

    fhandles = calloc(sizeof(*fhandles), desired_connections);
    for (i = 0; i < desired_connections - 1; i++) {
	fhandles[i].next = &fhandles[i + 1];
    }
    free_fhandle = fhandles;

    /* boost open file limit to the max: */
    if (getrlimit(RLIMIT_NOFILE, &rlimit) < 0) {
	fprintf(stderr, "%s: failed to get number of open file limit: %s",
	    prog_name, strerror(errno));
	exit(1);
    }
    rlimit.rlim_cur = rlimit.rlim_max;
    if (setrlimit(RLIMIT_NOFILE, &rlimit) < 0) {
	fprintf(stderr, "%s: failed to increase number of open file limit: %s",
	    prog_name, strerror(errno));
	exit(1);
    }
    memset(&server_addr, 0, sizeof(server_addr));
    server_addr.sin_family = AF_INET;
    server_addr.sin_port = htons(port);

    he = gethostbyname(server);
    if (he) {
	if (he->h_addrtype != AF_INET || he->h_length != sizeof(server_addr.sin_addr)) {
	    perror(server);
	    exit(-1);
	}
	memcpy(&server_addr.sin_addr, he->h_addr_list[0],
	    sizeof(server_addr.sin_addr));
    } else if (!inet_aton(server, &server_addr.sin_addr)) {
	fprintf(stderr, "%s: invalid server address %s\n", prog_name, server);
	exit(-1);
    }
    event_init();
}

static void open_connections(void);

static void
handle_connection(int fd, UNUSED short event, void *data)
{
    struct fhandle *fh = data;
    char buf[64];
    int len;
    int do_close = 0;
    len = read(fd, buf, sizeof(buf));
    if (len < 0 && errno != EAGAIN) {
	error++;
	perror("read");
    }
    if (len >= 0 || errno != EAGAIN)
	do_close = 1;
    if (!fh->open) {
	fh->open = 1;
	pending_connections--;
	if (!do_close) {
	    fh->ev.ev_events = EV_READ | EV_PERSIST;
	    event_add(&fh->ev, NULL);
	}
    }
    if (do_close) {
	event_del(&fh->ev);
	close(fd);
	put_fhandle(fh);
	current_connections--;
    }
    if (!error)
	open_connections();
    if (!pending_connections && current_connections >= desired_connections)
	printf("All open\n");
}

static int
open_connection(void)
{
    int fl;
    int sd;
    struct fhandle *fh;
    int i;

    const int max_batch = 1000;

    if (current_connections >= desired_connections || pending_connections >= max_batch)
	return 0;

    fh = get_fhandle();

    /* create more idle connections */
    sd = socket(AF_INET, SOCK_STREAM, 0);
    if (sd < 0) {
	perror("socket");
	exit(-1);
    }
    fl = fcntl(sd, F_GETFL);
    fcntl(sd, F_SETFL, fl | O_NONBLOCK);
    i = 1;
    setsockopt(sd, SOL_SOCKET, SO_REUSEADDR, &i, sizeof(i));
    setlocalport(sd);
    if (connect(sd, (struct sockaddr *) &server_addr, sizeof(server_addr)) < 0) {
	switch (errno) {
	case EINPROGRESS:
	    /* Fine, don't care */
	    pending_connections++;
	    fh->open = 0;
	    break;
	case ECONNREFUSED:
	case ETIMEDOUT:
	    error++;
	    perror("connect");
	    close(sd);
	    put_fhandle(fh);
	    return 0;

	default:
	    perror("connect");
	    exit(-1);
	}
    } else {
	fh->open = 1;
    }

    ++current_connections;
    if (current_connections % 1000 == 0 || current_connections == desired_connections)
	printf("%d connections, %d pending\n", current_connections, pending_connections);

    if (fh->open)
	event_set(&fh->ev, sd, EV_READ | EV_PERSIST, handle_connection, fh);
    else
	event_set(&fh->ev, sd, EV_WRITE, handle_connection, fh);
    event_add(&fh->ev, NULL);
    return 1;
}

static void
open_connections(void)
{
    while (open_connection());
}


static void timer_event(UNUSED int fd, UNUSED short event, UNUSED void *data);

static void
setup_timer(void)
{
    static struct event timer;
    struct timeval tv;

    tv.tv_sec = 1;
    tv.tv_usec = 000000;

    evtimer_set(&timer, timer_event, NULL);
    evtimer_add(&timer, &tv);
}

static void
timer_event(UNUSED int fd, UNUSED short event, UNUSED void *data)
{
    if (pending_connections == 0) {
	error = 0;
	open_connection();
    }
    setup_timer();
}

int
main(int argc, char **argv)
{
    int i;
    prog_name = strrchr(argv[0], '/');
    if (prog_name)
	++prog_name;
    else
	prog_name = argv[0];

    if (argc < 4) {
	fprintf(stderr, "Usage: %s server port numidle [local_ip...]\n", prog_name);
	exit(-1);
    }
    init(argv[1], argv[2], argv[3]);

    for (i = 4; i < argc; i++)
	add_my_addr(argv[i]);

    printf("%s: creating and maintaining %d idle connections\n",
	prog_name, desired_connections);

    open_connection();

    setup_timer();

    event_dispatch();

    return 0;
}