/* -*- Mode: C; c-basic-offset:4 ; indent-tabs-mode:nil ; -*- */

/*

 *  (C) 2008 by Argonne National Laboratory.

 *      See COPYRIGHT in top-level directory.

 */

#include "mpi.h"

#include <pthread.h>

#include <signal.h>

#include <stdio.h>

#include <stdlib.h>

#include <string.h>

#include <time.h>

#include <unistd.h>

#include "mpitest.h"

#include "connectstuff.h"

int main(int argc, char **argv)

{

    MPI_Comm tmp, comm, startComm;

    char *fname;

    char *actualFname = NULL;

    char *globalFname = NULL;

    int totalSize, expectedRank, size, cachedRank;

    char portName[MPI_MAX_PORT_NAME];

    int rankToAccept = 1;

    /* Debug - print out where we picked up the MPICH build from */

#ifdef MPICHLIBSTR

    msg("MPICH library taken from: %s\n", MPICHLIBSTR);

#endif

    if (argc != 4) {

        printf("Usage: %s <fname> <totalSize> <idx-1-based>\n", argv[0]);

        exit(1);

    }

    /* This is the base name of the file into which we write the port */

    fname = argv[1];

    /* This is the total number of processes launched */

    totalSize = atoi(argv[2]);

    /* Each process knows its expected rank */

    expectedRank = atoi(argv[3]) - 1;

    /* Start a watchdog thread which will abort after 120 seconds, and will

     * print stack traces using GDB every 5 seconds if you don't call

     * strokeWatchdog() */

    startWatchdog(120);

    /* Print a debug header */

    msg("Waiting for: %d - my rank is %d\n", totalSize, expectedRank);

    /* Singleton init */

    MPI_Init(0, 0);

    /* Duplicate from MPI_COMM_SELF the starting point */

    MPI_Comm_dup(MPI_COMM_SELF, &startComm);

    if (expectedRank == 0) {

        /* This process opens the port, and writes the information to the file */

        MPI_Open_port(MPI_INFO_NULL, portName);

        /* Write the port to fname.<rank> so that the connecting processes can

         * wait their turn by checking for the correct file to show up */

        actualFname = writePortToFile(portName, "%s.%d", fname, rankToAccept++);

        /* The wrapper script I'm using checks for the existance of "fname", so

         * create that - even though it isn't used  */

        globalFname = writePortToFile(portName, fname);

        installExitHandler(globalFname);

        comm = startComm;

    } else {

        char *readPort;

        readPort = getPortFromFile("%s.%d", fname, expectedRank);

        strncpy(portName, readPort, MPI_MAX_PORT_NAME);

        free(readPort);

        msg("Read port <%s>\n", portName);

        MPI_Comm_connect(portName, MPI_INFO_NULL, 0, startComm, &comm);

        MPI_Intercomm_merge(comm, 1, &tmp);

        comm = tmp;

        MPI_Comm_size(comm, &size);

        msg("After my first merge, size is now: %d\n", size);

    }

    while (size < totalSize) {

        /* Make sure we don't print a stack until we stall */

        strokeWatchdog();

        /* Accept the connection */

        MPI_Comm_accept(portName, MPI_INFO_NULL, 0, comm, &tmp);

        /* Merge into intracomm */

        MPI_Intercomm_merge(tmp, 0, &comm);

        /* Free the intercomm */

        MPI_Comm_free(&tmp);

        /* See where we're up to */

        MPI_Comm_rank(comm, &cachedRank);

        MPI_Comm_size(comm, &size);

        if (expectedRank == 0) {

            msg("Up to size: %d\n", size);

            /* Delete the old file, create the new one */

            unlink(actualFname);

            free(actualFname);

            /* Allow the next rank to connect */

            actualFname = writePortToFile(portName, "%s.%d", fname, rankToAccept++);

        }

    }

    MPI_Comm_rank(comm, &cachedRank);

    msg("All done - I got rank: %d.\n", cachedRank);

    MPI_Barrier(comm);

    if (expectedRank == 0) {

        /* Cleanup on rank zero - delete some files */

        MTestSleep(4);

        unlink(actualFname);

        free(actualFname);

        unlink(globalFname);

        free(globalFname);

        /* This lets my wrapper script know that we did everything correctly */

        indicateConnectSucceeded();

    }

    MPI_Finalize();

    return 0;

}