diff --git a/assignment-client/src/main.cpp b/assignment-client/src/main.cpp index 67d2a1e021..130cca9135 100644 --- a/assignment-client/src/main.cpp +++ b/assignment-client/src/main.cpp @@ -7,7 +7,10 @@ // #include +#include +#include #include +#include #include #include @@ -18,38 +21,13 @@ const long long ASSIGNMENT_REQUEST_INTERVAL_USECS = 1 * 1000 * 1000; +pid_t* childForks = NULL; +sockaddr_in customAssignmentSocket = {}; +const char* assignmentPool = NULL; +int numForks = 0; -int main(int argc, const char* argv[]) { - - setvbuf(stdout, NULL, _IOLBF, 0); - - sockaddr_in customAssignmentSocket = {}; - - // grab the overriden assignment-server hostname from argv, if it exists - const char* customAssignmentServer = getCmdOption(argc, argv, "-a"); - if (customAssignmentServer) { - customAssignmentSocket = socketForHostnameAndHostOrderPort(customAssignmentServer, ASSIGNMENT_SERVER_PORT); - } - - const char* NUM_FORKS_PARAMETER = "-n"; - const char* numForksIncludingParentString = getCmdOption(argc, argv, NUM_FORKS_PARAMETER); - - if (numForksIncludingParentString) { - int numForksIncludingParent = atoi(numForksIncludingParentString); - qDebug() << "Starting" << numForksIncludingParent << "assignment clients."; - - int processID = 0; - - // fire off as many children as we need (this is one less than the parent since the parent will run as well) - for (int i = 0; i < numForksIncludingParent - 1; i++) { - processID = fork(); - - if (processID == 0) { - // this is one of the children, break so we don't start a fork bomb - break; - } - } - } +void childClient() { + // this is one of the child forks or there is a single assignment client, continue assignment-client execution // create a NodeList as an unassigned client NodeList* nodeList = NodeList::createInstance(NODE_TYPE_UNASSIGNED); @@ -67,10 +45,6 @@ int main(int argc, const char* argv[]) { unsigned char packetData[MAX_PACKET_SIZE]; ssize_t receivedBytes = 0; - // grab the assignment pool from argv, if it was passed - const char* ASSIGNMENT_POOL_PARAMETER = "-p"; - const char* assignmentPool = getCmdOption(argc, argv, ASSIGNMENT_POOL_PARAMETER); - // create a request assignment, accept all assignments, pass the desired pool (if it exists) Assignment requestAssignment(Assignment::Request, Assignment::All, assignmentPool); @@ -110,4 +84,108 @@ int main(int argc, const char* argv[]) { nodeList->clear(); } } +} + +void sigchldHandler(int sig) { + pid_t processID; + int status; + + while ((processID = waitpid(-1, &status, WNOHANG)) != -1) { + if (processID == 0) { + // there are no more children to process, break out of here + break; + } + + qDebug() << "Handling death of" << processID; + + int newForkProcessID = 0; + + // find the dead process in the array of child forks + for (int i = 0; i < ::numForks; i++) { + if (::childForks[i] == processID) { + qDebug() << "Matched" << ::childForks[i] << "with" << processID; + + newForkProcessID = fork(); + if (newForkProcessID == 0) { + // this is the child, call childClient + childClient(); + + // break out so we don't fork bomb + break; + } else { + // this is the parent, replace the dead process with the new one + ::childForks[i] = newForkProcessID; + break; + } + } + } + } +} + +void parentMonitor() { + + struct sigaction sa; + + memset(&sa, 0, sizeof(sa)); + sa.sa_handler = sigchldHandler; + + sigaction(SIGCHLD, &sa, NULL); + + pid_t childID = 0; + + // don't bail until all children have finished + while ((childID = waitpid(-1, NULL, 0))) { + if (errno == ECHILD) { + break; + } + } + + // delete the array of pid_t holding the forked process IDs + delete[] ::childForks; +} + +int main(int argc, const char* argv[]) { + + setvbuf(stdout, NULL, _IOLBF, 0); + + // grab the overriden assignment-server hostname from argv, if it exists + const char* customAssignmentServer = getCmdOption(argc, argv, "-a"); + if (customAssignmentServer) { + ::customAssignmentSocket = socketForHostnameAndHostOrderPort(customAssignmentServer, ASSIGNMENT_SERVER_PORT); + } + + const char* NUM_FORKS_PARAMETER = "-n"; + const char* numForksString = getCmdOption(argc, argv, NUM_FORKS_PARAMETER); + + // grab the assignment pool from argv, if it was passed + const char* ASSIGNMENT_POOL_PARAMETER = "-p"; + ::assignmentPool = getCmdOption(argc, argv, ASSIGNMENT_POOL_PARAMETER); + + int processID = 0; + + if (numForksString) { + ::numForks = atoi(numForksString); + qDebug() << "Starting" << numForks << "assignment clients."; + + ::childForks = new pid_t[numForks]; + + // fire off as many children as we need (this is one less than the parent since the parent will run as well) + for (int i = 0; i < numForks; i++) { + processID = fork(); + + if (processID == 0) { + // this is in one of the children, break so we don't start a fork bomb + break; + } else { + // this is in the parent, save the ID of the forked process + childForks[i] = processID; + } + } + } + + if (processID == 0 || numForks == 0) { + childClient(); + } else { + parentMonitor(); + } } \ No newline at end of file