/**
* Copyright (C) Mellanox Technologies Ltd. 2001-2011.  ALL RIGHTS RESERVED.
* This software product is a proprietary product of Mellanox Technologies Ltd.
* (the "Company") and all right, title, and interest and to the software product,
* including all associated intellectual property rights, are and shall
* remain exclusively with the Company.
*
* This software product is governed by the End User License Agreement
* provided with the software product.
* $COPYRIGHT$
* $HEADER$
*/

#if HAVE_CONFIG_H
#include "config.h"
#endif

#define _GNU_SOURCE
#include "libperf.h"

#include <mxm/api/mxm_api.h>
#include <mxm/api/mxm_config.h>

#include <sys/socket.h>
#include <sys/poll.h>
#include <sys/time.h>
#include <sys/fcntl.h>
#include <arpa/inet.h>
#include <stdlib.h>
#include <stdio.h>
#include <unistd.h>
#include <netdb.h>
#include <assert.h>
#include <getopt.h>
#include <libgen.h>
#include <errno.h>
#include <string.h>
#include <sys/types.h>
#include <signal.h>
#include <assert.h>
#include <time.h>
#include <locale.h>

#if HAVE_MPI
#include <mpi.h>
#endif


#define EP_ADDR_LEN       (256)
#define SIG_SOCK_CLOSED   SIGUSR1

enum test_command {
    TEST_CMD_SEND,
    TEST_CMD_PUT,
    TEST_CMD_GET,
    TEST_CMD_AM,
    TEST_CMD_UNDEFINED
};

enum test_method {
    TEST_BANDWIDTH,
    TEST_LATENCY,
};

enum {
    TEST_FLAG_RECV_ANY     = MXM_BIT(0),
    TEST_FLAG_NO_REQ_WAIT  = MXM_BIT(1),
    TEST_FLAG_ALL_SIZES    = MXM_BIT(2),
    TEST_FLAG_REG_MEM      = MXM_BIT(3),
    TEST_FLAG_REQ_SYNC     = MXM_BIT(4),
    TEST_FLAG_REQ_LAZY     = MXM_BIT(5),
    TEST_FLAG_BIDI         = MXM_BIT(6),
    TEST_FLAG_SUMMARY      = MXM_BIT(7),
    TEST_FLAG_SET_OPT      = MXM_BIT(8),
    TEST_FLAG_NUMERIC_FMT  = MXM_BIT(9)
};

struct perftest_user_params {
    enum test_command            test_cmd;
    enum test_method             test_method;
    mxm_perf_cycle_t             cycle_method;
    size_t                       size;
    long                         iters;
    unsigned                     window;
    uint32_t                     flags;
    mxm_req_data_type_t          data_type;
    unsigned                     iov_count;
    unsigned                     conn_cnt;
};

struct perftest_peer_info {
    unsigned long                vaddr;
    mxm_mem_key_t                mkey;
    char                         ep_addr[0];
};

struct perftest_context {
    mxm_h                        mxmh;
    mxm_ep_h                     ep;
    mxm_mq_h                     mq;
    mxm_conn_h                   *conns;

    const char                   *server_addr;
    int                          port;
    int                          server_loop;

    struct perftest_user_params  test_params;
    double                       latency_divider;

    void                         *mem_access_buf;
    size_t                       mem_access_buf_size;

    struct perftest_peer_info    *local, *remote;
};


mxm_error_t mxm_ep_create_internal(mxm_h context, mxm_ep_opts_t *opts,
                                   unsigned num_slots, mxm_ep_h *ep_p);

mxm_error_t mxm_ep_get_address_internal(mxm_ep_h ep, unsigned slot_index,
                                        void *address, size_t *addrlen_p);

mxm_error_t mxm_ep_connect_internal(mxm_ep_h ep, unsigned slot_index,
                                    void *address, mxm_conn_h *conn_p);


static int init_ctx(struct perftest_context *ctx)
{
    mxm_context_opts_t *mxm_opts;
    mxm_ep_opts_t *ep_opts;
    mxm_error_t error;
    size_t addr_len;
    unsigned index;

    if (ctx->test_params.flags & TEST_FLAG_SET_OPT) {
        setenv("MXM_IB_RX_QUEUE_LEN", "1024",   0);
        setenv("MXM_ASYNC_MODE",      "signal", 0);
    }

    error = mxm_config_read_opts(&mxm_opts, &ep_opts, NULL, NULL, 0);
    if (error != MXM_OK) {
        fprintf(stderr, "Failed to parse options: %s\n", mxm_error_string(error));
        return -1;
    }

    error = mxm_init(mxm_opts, &ctx->mxmh);
    if (error != MXM_OK) {
        fprintf(stderr, "Failed to create MXM: %s\n", mxm_error_string(error));
        return -1;
    }

    error = mxm_ep_create_internal(ctx->mxmh, ep_opts, ctx->test_params.conn_cnt, &ctx->ep);
    if (error != MXM_OK) {
        fprintf(stderr, "Failed to create endpoint: %s\n", mxm_error_string(error));
        return -1;
    }

    ctx->local = malloc(sizeof(struct perftest_peer_info) + ctx->test_params.conn_cnt * EP_ADDR_LEN);
    if (ctx->local == NULL) {
        fprintf(stderr, "Failed to allocate memory for local peer info.\n");
        return -1;
    }

    ctx->remote = malloc(sizeof(struct perftest_peer_info) + ctx->test_params.conn_cnt * EP_ADDR_LEN);
    if (ctx->remote == NULL) {
        fprintf(stderr, "Failed to allocate memory for remote peer info.\n");
        return -1;
    }

    for (index = 0; index < ctx->test_params.conn_cnt; index++) {
        addr_len = EP_ADDR_LEN;
        error = mxm_ep_get_address_internal(ctx->ep, index,
                                            (char*)ctx->local->ep_addr + index * EP_ADDR_LEN,
                                            &addr_len);
        if (error != MXM_OK) {
            fprintf(stderr, "Failed to get endpoint address of slot #%i: %s\n",
                    index, mxm_error_string(error));
            return -1;
        }
    }

    /* allocate memory for rdma test */
    ctx->mem_access_buf = NULL;
    ctx->mem_access_buf_size = ctx->test_params.size;
    error = mxm_mem_map(ctx->mxmh, &ctx->mem_access_buf, &ctx->mem_access_buf_size,
                        0, NULL, 0);
    if (error != MXM_OK) {
        fprintf(stderr, "Failed to allocate memory: %s\n", mxm_error_string(error));
        return -1;
    }

    ctx->local->vaddr = (uintptr_t)ctx->mem_access_buf;
    error = mxm_mem_get_key(ctx->mxmh, ctx->mem_access_buf, &ctx->local->mkey);
    if (error != MXM_OK) {
        fprintf(stderr, "Failed to get memory key: %s.\n", mxm_error_string(error));
        return -1;
    }

    mxm_config_free_context_opts(mxm_opts);
    mxm_config_free_ep_opts(ep_opts);
    return 0;
}

static void cleanup_ctx(struct perftest_context *ctx)
{
    unsigned index;
    mxm_mem_unmap(ctx->mxmh, ctx->mem_access_buf, ctx->mem_access_buf_size, 0);

    mxm_mq_destroy(ctx->mq);
    for (index = 0; index < ctx->test_params.conn_cnt; index++) {
        mxm_ep_disconnect(ctx->conns[index]);
    }
    free(ctx->conns);
    mxm_ep_destroy(ctx->ep);
    mxm_cleanup(ctx->mxmh);
}

static int safe_send(int sock, void *data, size_t size)
{
    size_t total = 0;
    int ret;

    while (total < size) {
        ret = send(sock, (char*)data + total, size - total, 0);
        if (ret < 0) {
            fprintf(stderr, "send() failed: %m\n");
            return -1;
        }
        total += ret;
    }
    return 0;
}

static int safe_recv(int sock, void *data, size_t size)
{
    size_t total = 0;
    int ret;

    while (total < size) {
        ret = recv(sock, (char*)data + total, size - total, 0);
        if (ret < 0) {
            fprintf(stderr, "recv() failed: %m\n");
            return -1;
        }
        total += ret;
    }
    return 0;
}

static int connect_eps(struct perftest_context *ctx)
{
    unsigned index;
    mxm_error_t error;

    ctx->conns = malloc(ctx->test_params.conn_cnt * sizeof(mxm_conn_h));
    if (ctx->conns == NULL) {
        fprintf(stderr, "Failed to allocate memory for connections.\n");
        return -1;
    }

    for (index = 0; index < ctx->test_params.conn_cnt; index++) {
        error = mxm_ep_connect_internal(ctx->ep, index,
                                        (char*)ctx->remote->ep_addr + index * EP_ADDR_LEN,
                                        &ctx->conns[index]);
        if (error != MXM_OK) {
            free(ctx->conns);
            fprintf(stderr, "mxm_ep_connect_internal() failed on slot #%i: %s.\n",
                    index, mxm_error_string(error));
            return -1;
        }
    }

    error = mxm_mq_create(ctx->mxmh, 0x1337, &ctx->mq);
    if (error != MXM_OK) {
        free(ctx->conns);
        fprintf(stderr, "Failed to create MQ: %s.\n", mxm_error_string(error));
        return -1;
    }

    error = mxm_ep_wireup(ctx->ep);
    if (error != MXM_OK) {
        free(ctx->conns);
        fprintf(stderr, "Failed to wire-up all connections: %s.\n", mxm_error_string(error));
        return -1;
    }
    return 0;
}

static int exchange_info(struct perftest_context *ctx, int sockfd, int send_first)
{
    size_t size;
    int i, ret;

    size = sizeof(struct perftest_peer_info) + ctx->test_params.conn_cnt * EP_ADDR_LEN;
    for (i = 0; i < 2; ++i) {
        if (i == !send_first) {
            ret = safe_send(sockfd, ctx->local, size);
        } else {
            ret = safe_recv(sockfd, ctx->remote, size);
        }
        if (ret < 0) {
            return ret;
        }
    }
    return ret;
}

static int server_init_ctx(struct perftest_context *ctx, int sockfd)
{
    int dummy = 0x8e110;

    if (safe_send(sockfd, &dummy, sizeof(dummy)) < 0) {
        return -1;
    }

    /* Server receive test params */
    if (safe_recv(sockfd, &ctx->test_params, sizeof(ctx->test_params)) < 0) {
        return -1;
    }

    /* Initialize context */
    if (init_ctx(ctx) < 0) {
        return -1;
    }

    if (exchange_info(ctx, sockfd, 1) < 0) {
        return -1;
    }

    if (connect_eps(ctx) < 0) {
        return -1;
    }

    return 0;
}

static int client_init_ctx(struct perftest_context *ctx, int sockfd)
{
    int dummy;

    if (safe_recv(sockfd, &dummy, sizeof(dummy)) < 0) {
        return -1;
    }

    /* Client send test params */
    if (safe_send(sockfd, &ctx->test_params, sizeof(ctx->test_params)) < 0) {
        return -1;
    }

    /* Initialize context */
    if (init_ctx(ctx) < 0) {
        return -1;
    }

    if (exchange_info(ctx, sockfd, 0) < 0) {
        return -1;
    }

    if (connect_eps(ctx) < 0) {
        return -1;
    }

    return 0;
}

static void print_header(void)
{
    printf("+--------------+-----------------------------+---------------------+-----------------------+\n");
    printf("|              |       latency (usec)        |   bandwidth (MB/s)  |  message rate (msg/s) |\n");
    printf("+--------------+---------+---------+---------+----------+----------+-----------+-----------+\n");
    printf("| # iterations | typical | average | overall |  average |  overall |   average |   overall |\n");
    printf("+--------------+---------+---------+---------+----------+----------+-----------+-----------+\n");
}

static void print_progress(mxm_perf_result_t *result, void *user_data)
{
    static const char *fmt_numeric =  "%'14.0f %9.3f %9.3f %9.3f %10.2f %10.2f %'11.0f %'11.0f\n";
    static const char *fmt_plain   =  "%14.0f %9.3f %9.3f %9.3f %10.2f %10.2f %11.0f %11.0f\n";

    struct perftest_context *ctx = user_data;

    printf((ctx->test_params.flags & TEST_FLAG_NUMERIC_FMT) ? fmt_numeric : fmt_plain,
           (double)result->iters,
           result->latency.typical * 1000000.0 / ctx->latency_divider,
           result->latency.moment_average * 1000000.0 / ctx->latency_divider,
           result->latency.total_average * 1000000.0 / ctx->latency_divider,
           result->bandwidth.moment_average / (1024.0 * 1024.0),
           result->bandwidth.total_average / (1024.0 * 1024.0),
           result->msgrate.moment_average,
           result->msgrate.total_average);
    fflush(stdout);
}

static void print_summary(mxm_perf_result_t *result, void *user_data)
{
    struct perftest_context *ctx = user_data;

    static const char *fmt_numeric =  "latency(usec):   %.3f\nbandwidth(MB/s): %.2f\nmsgrate(msg/s):  %'.0f\n";
    static const char *fmt_plain   =  "latency(usec):   %.3f\nbandwidth(MB/s): %.2f\nmsgrate(msg/s):  %.0f\n";

    printf((ctx->test_params.flags & TEST_FLAG_NUMERIC_FMT) ? fmt_numeric : fmt_plain,
           result->latency.total_average * 1000000.0 / ctx->latency_divider,
           result->bandwidth.total_average / (1024.0 * 1024.0),
           result->msgrate.total_average);
}

static int run_test(struct perftest_context *ctx, int index)
{
    mxm_perf_test_t test;
    mxm_error_t error;
    mxm_perf_result_t result;

    switch (ctx->test_params.test_cmd) {
    case TEST_CMD_SEND:
        if (ctx->test_params.test_method == TEST_BANDWIDTH) {
            if (ctx->test_params.flags & TEST_FLAG_BIDI) {
                test.command = MXM_PERF_CMD_SEND_RECV;
            } else {
                test.command = (index == 0) ? MXM_PERF_CMD_RECV : MXM_PERF_CMD_SEND;
            }
            ctx->latency_divider = 1.0;
        } else {
            test.command = MXM_PERF_CMD_SEND_PINGPONG;
            ctx->latency_divider = 2.0;
        }
        break;
    case TEST_CMD_PUT:
        if (ctx->test_params.test_method == TEST_BANDWIDTH) {
            if ((index > 0) || (ctx->test_params.flags & TEST_FLAG_BIDI)) {
                test.command = MXM_PERF_CMD_PUT;
            } else {
                test.command = MXM_PERF_CMD_PROGRESS;
            }
            ctx->latency_divider = 1.0;
        } else {
            test.command = MXM_PERF_CMD_PUT_PINGPONG;
            ctx->latency_divider = 2.0;
        }
        break;
    case TEST_CMD_GET:
        assert(ctx->test_params.test_method == TEST_BANDWIDTH);
        if ((index > 0) || (ctx->test_params.flags & TEST_FLAG_BIDI)) {
            test.command = MXM_PERF_CMD_GET;
        } else {
            test.command = MXM_PERF_CMD_PROGRESS;
        }
        ctx->latency_divider = 1.0;
        break;
    case TEST_CMD_AM:
        if (ctx->test_params.test_method == TEST_BANDWIDTH) {
            if ((index > 0) || (ctx->test_params.flags & TEST_FLAG_BIDI)) {
                test.command = MXM_PERF_CMD_AM;
            } else {
                test.command = MXM_PERF_CMD_PROGRESS;
            }
            ctx->latency_divider = 1.0;
        } else {
            test.command = MXM_PERF_CMD_AM_PINGPONG;
            ctx->latency_divider = 2.0;
        }
        break;
    default:
        return -1;
    }

    test.index = index;
    test.flags = MXM_PERF_TEST_FLAG_WARMUP;
    if (!(ctx->test_params.flags & TEST_FLAG_RECV_ANY)) {
        test.flags |= MXM_PERF_TEST_FLAG_RECV_CONN;
    }
    if (!(ctx->test_params.flags & TEST_FLAG_NO_REQ_WAIT)) {
        test.flags |= MXM_PERF_TEST_FLAG_REQ_WAIT;
    }
    if (ctx->test_params.flags & TEST_FLAG_REG_MEM) {
        test.flags |= MXM_PERF_TEST_FLAG_REG_MEM;
    }
    if (ctx->test_params.flags & TEST_FLAG_REQ_SYNC) {
        test.flags |= MXM_PERF_TEST_FLAG_REQ_SYNC;
    }
    if (ctx->test_params.flags & TEST_FLAG_REQ_LAZY) {
        test.flags |= MXM_PERF_TEST_FLAG_REQ_LAZY;
    }
    test.message_size = ctx->test_params.size;
    test.max_iters = ctx->test_params.iters;
    test.window = ctx->test_params.window;
    test.max_time = 0.0;
    test.data_type = ctx->test_params.data_type;
    test.iov_count = ctx->test_params.iov_count;
    test.cycle = ctx->test_params.cycle_method;

    test.local_addr  = ctx->mem_access_buf;
    test.remote_addr = ctx->remote->vaddr;
    test.remote_mkey = ctx->remote->mkey;

    if (index == 0) {
        /* Server */
        error = mxm_perf_test_run(ctx->conns, ctx->test_params.conn_cnt, ctx->mq,
                                  ctx->ep, ctx->mxmh, &test,
                                  NULL, NULL, 0.5, &result);
        if (error != MXM_OK) {
            fprintf(stderr, "Failed to run test: %s\n", mxm_error_string(error));
            return -1;
        }
    } else {
        /* Client */
        if (!(ctx->test_params.flags & TEST_FLAG_SUMMARY)) {
            print_header();
        }
        setlocale(LC_ALL, "en_US");
        error = mxm_perf_test_run(ctx->conns, ctx->test_params.conn_cnt, ctx->mq,
                                  ctx->ep, ctx->mxmh, &test,
                                  (ctx->test_params.flags & TEST_FLAG_SUMMARY) ? NULL : print_progress,
                                  ctx, 0.5, &result);
        if (error != MXM_OK) {
            fprintf(stderr, "Failed to run test: %s\n", mxm_error_string(error));
            return -1;
        }

        if (ctx->test_params.flags & TEST_FLAG_SUMMARY) {
            print_summary(&result, ctx);
        } else {
            printf("Overall:\n");
            print_progress(&result, ctx);
        }
    }
    return 0;
}

int safe_fcntl(int fd, int cmd, int arg)
{
    int ret = fcntl(fd, cmd, arg);
    if (ret < 0) {
        fprintf(stderr, "fcntl(fd=%d cmd=%d) failed: %m\n", fd, cmd);
    }
    return ret;
}

static void sock_closed_handler(int signo)
{
    fprintf(stderr, "Remove side disconnected, exiting.\n");
    exit(1);
}

static void exit_on_close(struct perftest_context *test_ctx, int fd)
{
    signal(SIG_SOCK_CLOSED, sock_closed_handler);
    safe_fcntl(fd, F_SETSIG, SIG_SOCK_CLOSED);
    safe_fcntl(fd, F_SETOWN, getpid());
    safe_fcntl(fd, F_SETFL, safe_fcntl(fd, F_GETFL, 0) | O_ASYNC);
}

static void cancel_exit_on_close(struct perftest_context *test_ctx, int fd)
{
    safe_fcntl(fd, F_SETFL, safe_fcntl(fd, F_GETFL, 0) & ~O_ASYNC);
}

static int run_server(struct perftest_context *test_ctx)
{
    struct perftest_context ctx = *test_ctx;
    struct sockaddr_in local_addr;
    struct sockaddr_in client_addr;
    socklen_t client_addrlen;
    char client_name[128];
    int servsock, connfd;
    int optval = 1;
    int ret;

    servsock = socket(AF_INET, SOCK_STREAM, 0);
    if (servsock < 0) {
        fprintf(stderr, "socket() failed: %m\n");
        return -1;
    }

    optval = 1;
    ret = setsockopt(servsock, SOL_SOCKET, SO_REUSEADDR, &optval, sizeof(optval));
    if (ret < 0) {
        fprintf(stderr, "setsockopt(SO_REUSEADDR) failed: %m\n");
        goto out_close_servsock;
    }

    local_addr.sin_family = AF_INET;
    local_addr.sin_port = htons(ctx.port);
    local_addr.sin_addr.s_addr = INADDR_ANY;
    memset(local_addr.sin_zero, 0, sizeof(local_addr.sin_zero));
    ret = bind(servsock, (struct sockaddr*)&local_addr, sizeof(local_addr));
    if (ret < 0) {
        fprintf(stderr, "bind() failed: %m\n");
        goto out_close_servsock;
    }

    ret = listen(servsock, 10);
    if (ret < 0) {
        fprintf(stderr, "listen() failed: %m\n");
        goto out_close_servsock;
    }

    do {
        printf("Waiting for connection...\n");

        /* Accept next connection */
        ret = accept(servsock, NULL, NULL);
        if (ret < 0) {
            fprintf(stderr, "accept() failed: %m\n");
            goto out_close_servsock;
        }
        connfd = ret;

        /* Print client name */
        client_addrlen = sizeof(client_addr);
        ret = getsockname(connfd, (struct sockaddr *)&client_addr, &client_addrlen);
        if (ret < 0) {
            fprintf(stderr, "getsockname() failed: %m\n");
            close(connfd);
            goto out_close_servsock;
        }

        inet_ntop(client_addr.sin_family, &client_addr.sin_addr, client_name,
                  sizeof(client_name));
        printf("Accepted connection from %s\n", client_name);

        ret = server_init_ctx(&ctx, connfd);
        if (ret < 0) {
            close(connfd);
            goto out_close_servsock;
        }

        exit_on_close(&ctx, connfd);
        ret = run_test(&ctx, 0);
        cancel_exit_on_close(&ctx, connfd);

        cleanup_ctx(&ctx);
        close(connfd);

    } while (ctx.server_loop);

    ret = 0;

out_close_servsock:
    close(servsock);
    return ret;
}

static int run_client(struct perftest_context *ctx)
{
    struct sockaddr_in conn_addr;
    struct hostent *he;
    int connfd;
    int ret;

    connfd = socket(AF_INET, SOCK_STREAM, 0);
    if (connfd < 0) {
        fprintf(stderr, "socket() failed: %m\n");
        return -1;
    }

    he = gethostbyname(ctx->server_addr);
    if (he == NULL || he->h_addr_list == NULL) {
        fprintf(stderr, "host %s not found: %s\n", ctx->server_addr, hstrerror(h_errno));
        goto err_close_conn;
    }

    conn_addr.sin_family = he->h_addrtype;
    conn_addr.sin_port = htons(ctx->port);
    assert(he->h_length == sizeof(conn_addr.sin_addr));
    memcpy(&conn_addr.sin_addr, he->h_addr_list[0], he->h_length);
    memset(conn_addr.sin_zero, 0, sizeof(conn_addr.sin_zero));

    ret = connect(connfd, (struct sockaddr*)&conn_addr, sizeof(conn_addr));
    if (ret < 0) {
        fprintf(stderr, "connect() failed: %m\n");
        goto err_close_conn;
    }

    if (client_init_ctx(ctx, connfd) < 0) {
        goto err_close_conn;
    }

    ret = run_test(ctx, 1);
    cleanup_ctx(ctx);
    close(connfd);

    return ret;

err_close_conn:
    close(connfd);
    return -1;
}

static void usage(struct perftest_context *ctx, const char *program)
{
    printf(
           "Usage: %s [ server ] [ options ]\n"
           "\n"
           "  Common options:\n"
           "     -h           Show help.\n"
           "     -p <port>    TCP port to use. Default: %d.\n"
           "     -o           Don't configure MXM for maximal performance, be more portable.\n"
           "\n"
           "  Client options:\n"
           "     -t <test>    Test to run:\n"
           "                     send_lat : send/recv latency.\n"
           "                     send_bw  : send/recv bandwidth.\n"
           "                     put_lat  : put latency.\n"
           "                     put_bw   : put bandwidth.\n"
           "                     get_bw   : get bandwidth.\n"
           "                     am_lat   : active message latency.\n"
           "                     am_bw    : active message bandwidth.\n"
           "     -n <iters>   Number of iterations to run. Default: %ld.\n"
           "     -W <length>  Send this many requests at once. Default: %u.\n"
           "     -s <size>    Message size. Default: %lu.\n"
           "     -d <type>    Data type to use:\n"
           "                     buffer   : Plain data buffer (default).\n"
           "                     stream   : Streaming callback.\n"
           "                     iov      : IOV list.\n"
           "     -b           Run bi-directional test. Relevant only for bandwidth tests.\n"
           "     -e <count>   Number of connections on each endpoint.\n"
           "     -m <method>  Method to cycle through connections:\n"
           "                     rr       : Round-Robin (default).\n"
           "                     uniform  : Random connection, uniform distribution.\n"
           "                     nonuiform: Random connection, non-uniform distribution.\n"
           "     -c <count>   Number of entries in IOV list, if used. Default: %u.\n"
           "     -a           Receive from ANY.\n"
           "     -i           Don't set REQ_WAIT flag.\n"
           "     -r           Register data buffers and pass memory key to MXM.\n"
           "     -A           Test with all message sizes.\n"
           "     -S           Use SYNC operations - for SEND/PUT.\n"
           "     -L           Set the LAZY flag on requests.\n"
           "     -j           Print only summary of performance.\n"
           "     -N           Print numbers without numberic formatting.\n"
           "\n"
           "  Server options:\n"
           "     -l           Accept clients in an infinite loop\n"
           "\n"
           "\n NOTE: if server name is 'fork', the a server is forked in the background."
           "\n",

           program,
           ctx->port,
           ctx->test_params.iters,
           ctx->test_params.window,
           ctx->test_params.size,
           ctx->test_params.iov_count
        );
}

static int parse_opts(struct perftest_context *ctx, int argc, char **argv)
{
    char c;

    ctx->server_addr = NULL;
    ctx->port = 13337;
    ctx->server_loop = 0;
    ctx->test_params.cycle_method = MXM_PERF_TEST_CYCLE_RR;
    ctx->test_params.test_cmd = TEST_CMD_UNDEFINED;
    ctx->test_params.iters = 1000000l;
    ctx->test_params.window = 1;
    ctx->test_params.size = 0;
    ctx->test_params.flags = TEST_FLAG_SET_OPT | TEST_FLAG_NUMERIC_FMT;
    ctx->test_params.data_type = MXM_REQ_DATA_BUFFER;
    ctx->test_params.iov_count = 2;
    ctx->test_params.conn_cnt = 1;

    while ((c = getopt (argc, argv, "p:n:t:s:d:c:m:e:airAlSW:LobhjN")) != -1) {
        switch (c) {
        case 'p':
            ctx->port = atoi(optarg);
            break;
        case 'n':
            ctx->test_params.iters = atol(optarg);
            break;
        case 't':
            if (0 == strcmp(optarg, "send_lat")) {
                ctx->test_params.test_cmd = TEST_CMD_SEND;
                ctx->test_params.test_method = TEST_LATENCY;
            } else if (0 == strcmp(optarg, "send_bw")) {
                ctx->test_params.test_cmd = TEST_CMD_SEND;
                ctx->test_params.test_method = TEST_BANDWIDTH;
            } else if (0 == strcmp(optarg, "get_bw")) {
                ctx->test_params.test_cmd = TEST_CMD_GET;
                ctx->test_params.test_method = TEST_BANDWIDTH;
            } else if (0 == strcmp(optarg, "put_lat")) {
                ctx->test_params.test_cmd = TEST_CMD_PUT;
                ctx->test_params.test_method = TEST_LATENCY;
            } else if (0 == strcmp(optarg, "put_bw")) {
                ctx->test_params.test_cmd = TEST_CMD_PUT;
                ctx->test_params.test_method = TEST_BANDWIDTH;
            } else if (0 == strcmp(optarg, "am_lat")) {
                ctx->test_params.test_cmd = TEST_CMD_AM;
                ctx->test_params.test_method = TEST_LATENCY;
            } else if (0 == strcmp(optarg, "am_bw")) {
                ctx->test_params.test_cmd = TEST_CMD_AM;
                ctx->test_params.test_method = TEST_BANDWIDTH;
            } else {
                fprintf(stderr, "Invalid option argument for -t\n");
                return -1;
            }
            break;
        case 's':
            ctx->test_params.size = atol(optarg);
            break;
        case 'e':
            ctx->test_params.conn_cnt = atol(optarg);
            break;
        case 'd':
            if (0 == strcmp(optarg, "buffer")) {
                ctx->test_params.data_type = MXM_REQ_DATA_BUFFER;
            } else if (0 == strcmp(optarg, "stream")) {
                ctx->test_params.data_type = MXM_REQ_DATA_STREAM;
            } else if (0 == strcmp(optarg, "iov")) {
                ctx->test_params.data_type = MXM_REQ_DATA_IOV;
            } else {
                fprintf(stderr, "Invalid option argument for -d\n");
                return -1;
            }
            break;
        case 'm':
            if (0 == strcmp(optarg, "rr")) {
                ctx->test_params.cycle_method = MXM_PERF_TEST_CYCLE_RR;
            } else if (0 == strcmp(optarg, "uniform")) {
                ctx->test_params.cycle_method = MXM_PERF_TEST_CYCLE_UNIFORM;
            } else if (0 == strcmp(optarg, "nonuniform")) {
                ctx->test_params.cycle_method = MXM_PERF_TEST_CYCLE_NONUNIFORM;
            } else {
                fprintf(stderr, "Invalid option argument for -m\n");
                return -1;
            }
            break;
        case 'c':
            ctx->test_params.iov_count = atoi(optarg);
            break;
        case 'a':
            ctx->test_params.flags |= TEST_FLAG_RECV_ANY;
            break;
        case 'i':
            ctx->test_params.flags |= TEST_FLAG_NO_REQ_WAIT;
            break;
        case 'r':
            ctx->test_params.flags |= TEST_FLAG_REG_MEM;
            break;
        case 'A':
            ctx->test_params.flags |= TEST_FLAG_ALL_SIZES;
            break;
        case 'S':
            ctx->test_params.flags |= TEST_FLAG_REQ_SYNC;
            break;
        case 'L':
            ctx->test_params.flags |= TEST_FLAG_REQ_LAZY;
            break;
        case 'W':
            ctx->test_params.window = atoi(optarg);
            break;
        case 'l':
            ctx->server_loop = 1;
            break;
        case 'o':
            ctx->test_params.flags &= ~TEST_FLAG_SET_OPT;
            break;
        case 'b':
            ctx->test_params.flags |= TEST_FLAG_BIDI;
            break;
        case 'j':
            ctx->test_params.flags |= TEST_FLAG_SUMMARY;
            break;
        case 'N':
            ctx->test_params.flags &= ~TEST_FLAG_NUMERIC_FMT;
            break;
        case 'h':
        default:
           usage(ctx, basename(argv[0]));
           return -1;
        }
    }

    if (optind < argc) {
        ctx->server_addr = argv[optind];
    }

    return 0;
}

int run_standalone(int argc, char **argv)
{
    struct perftest_context ctx;
    pid_t pid;
    int ret;

    memset(&ctx, 0, sizeof(ctx));
    if (parse_opts(&ctx, argc, argv) < 0) {
        return -1;
    }

    if (ctx.server_addr == NULL) {
        return run_server(&ctx);
    } else {
        if (ctx.test_params.test_cmd == TEST_CMD_UNDEFINED) {
            fprintf(stderr, "Must specify test type with -t\n");
            return -1;
        }

        if (!strcmp(ctx.server_addr, "fork")) {
            srand(time(NULL));
            ctx.port = rand() % 10000 + 10000;
            pid = fork();
            if (pid == 0) {
                ctx.server_addr = NULL;
                return run_server(&ctx);
            }

            ctx.server_addr = "localhost";
        } else {
            pid = 0;
        }

        ret = run_client(&ctx);

        if (pid > 0) {
            kill(pid, SIGKILL);
        }

        return ret;
    }
}

int run_mpi(int argc, char **argv)
{
#if HAVE_MPI
    MPI_Status status;
    struct perftest_context ctx;
    int num_ranks, my_rank, dest_rank;
    size_t size;
    int ret;

    /* Don't try MPI when running interactively */
    if (isatty(0)) {
        return -ENOTSUP;
    }

    ret = MPI_Init(&argc, &argv);
    if (ret != 0) {
        return -ENOTSUP;
    }

    /* Use MPI only if we have at least 2 ranks */
    MPI_Comm_size(MPI_COMM_WORLD, &num_ranks);
    if (num_ranks == 1) {
        ret = -ENOTSUP;
        goto out;
    }

    MPI_Comm_rank(MPI_COMM_WORLD, &my_rank);

    if (my_rank < 2) {
        dest_rank = 1 - my_rank;

        if (my_rank == 0) {
            memset(&ctx, 0, sizeof(ctx));
            ret = parse_opts(&ctx, argc, argv);
            if (ret == 0) {
                if (ctx.test_params.test_cmd == TEST_CMD_UNDEFINED) {
                    fprintf(stderr, "Must specify test type with -t\n");
                    goto out;
                }
            }
        }

        MPI_Bcast(&ret,  1, MPI_INT, 0, MPI_COMM_WORLD);
        if (ret < 0) {
            goto out;
        }

        MPI_Bcast(&ctx.test_params,  sizeof(ctx.test_params), MPI_BYTE, 0, MPI_COMM_WORLD);

        ret = init_ctx(&ctx);
        if (ret < 0) {
            goto out;
        }

        size = sizeof(struct perftest_peer_info) + ctx.test_params.conn_cnt * EP_ADDR_LEN;
        MPI_Sendrecv(ctx.local,  size, MPI_BYTE, dest_rank, 1,
                     ctx.remote, size, MPI_BYTE, dest_rank, 1,
                     MPI_COMM_WORLD, &status);

        ret = connect_eps(&ctx);
        if (ret < 0) {
            goto out;
        }

        ret = run_test(&ctx, my_rank);
        cleanup_ctx(&ctx);
    }

out:
    MPI_Finalize();
    return ret;
#else
    return -ENOTSUP;
#endif
}

int main(int argc, char **argv)
{
    int ret;

    ret = run_mpi(argc, argv);
    if (ret == -ENOTSUP) {
        return run_standalone(argc, argv);
    } else {
        return ret;
    }
}
