Merge branch 'master' of https://github.com/worklist/hifi

assignment-client/src/main.cpp

@@ -7,7 +7,10 @@
 //
 
 #include <arpa/inet.h>
+#include <errno.h>
+#include <signal.h>
 #include <sys/time.h>
+#include <sys/wait.h>
 
 #include <Assignment.h>
 #include <AudioMixer.h>
@@ -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[]) {
+void childClient() {
-    
+    // this is one of the child forks or there is a single assignment client, continue assignment-client execution
-    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;
-            }
-        }
-    }
     
     // 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);
     
@@ -111,3 +85,107 @@ int main(int argc, const char* argv[]) {
         }
     }
 }
+
+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();
+    }
+}