Infinitely wait for subsequent connection until timeout

Makefile

@@ -136,7 +136,7 @@ build/$(JATTACH): src/jattach/*.c src/jattach/*.h
 	$(CC) $(CFLAGS) -DJATTACH_VERSION=\"$(PROFILER_VERSION)-ap\" -o $@ src/jattach/*.c
 
 build/fdtransfer: src/fdtransfer/*.cpp src/fdtransfer/*.h src/jattach/psutil.c src/jattach/psutil.h
-	$(CXX) $(CFLAGS) -o $@ src/fdtransfer/*.cpp src/jattach/psutil.c
+	$(CXX) $(CFLAGS) -o $@ src/fdtransfer/*.cpp src/jattach/psutil.c src/arguments.cpp
 
 build/$(API_JAR): $(API_SOURCES)
 	mkdir -p build/api

profiler.sh

@@ -98,7 +98,11 @@ fdtransfer() {
     if [ "$USE_FDTRANSFER" = "true" ]; then
         FDTRANSFER_PATH="@async-profiler-$(od -An -N3 -i /dev/random | xargs)"
         PARAMS="$PARAMS,fdtransfer=$FDTRANSFER_PATH"
-        "$FDTRANSFER" "$FDTRANSFER_PATH" "$PID"
+        if [ -n "$ACCEPT_TIMEOUT" ]; then
+          "$FDTRANSFER" "$FDTRANSFER_PATH" "$PID" "$ACCEPT_TIMEOUT"
+        else
+          "$FDTRANSFER" "$FDTRANSFER_PATH" "$PID"
+        fi 
     fi
 }
 
@@ -140,6 +144,7 @@ FDTRANSFER_PATH=""
 PROFILER=$SCRIPT_DIR/build/libasyncProfiler.so
 ACTION="collect"
 DURATION="60"
+ACCEPT_TIMEOUT=""
 FILE=""
 USE_TMP="true"
 OUTPUT=""
@@ -235,6 +240,7 @@ while [ $# -gt 0 ]; do
             if [ "$ACTION" = "collect" ]; then
                 ACTION="start"
             fi
+            ACCEPT_TIMEOUT="$2"
             PARAMS="$PARAMS,${1#--}=$2"
             shift
             ;;

src/arguments.h

@@ -130,7 +130,6 @@ class Arguments {
     static long long hash(const char* arg);
     static Output detectOutputFormat(const char* file);
     static long parseUnits(const char* str, const Multiplier* multipliers);
-    static int parseTimeout(const char* str);
 
   public:
     Action _action;
@@ -231,6 +230,8 @@ class Arguments {
 
     const char* file();
 
+    static int parseTimeout(const char* str);
+
     bool hasOutputFile() const {
         return _file != NULL &&
             (_action == ACTION_STOP || _action == ACTION_DUMP ? _output != OUTPUT_JFR : _action >= ACTION_STATUS);

src/fdtransfer/fdtransfer.h

@@ -24,6 +24,8 @@
 #include <sys/socket.h>
 #include <sys/un.h>
 
+#define DEFAULT_ACCEPT_TIMEOUT 15
+
 #define ARRAY_SIZE(arr)  (sizeof(arr) / sizeof(arr[0]))
 
 #define RESTARTABLE(call)  ({ ssize_t ret; while ((ret = call) < 0 && errno == EINTR); ret; })

src/fdtransfer/fdtransferServer.cpp

@@ -31,6 +31,7 @@
 
 #include "fdtransfer.h"
 #include "../jattach/psutil.h"
+#include "../arguments.h"
 
 
 class FdTransferServer {
@@ -50,6 +51,7 @@ class FdTransferServer {
 
 int FdTransferServer::_server;
 int FdTransferServer::_peer;
+static char *_accept_timeout = NULL;
 
 bool FdTransferServer::bindServer(struct sockaddr_un *sun, socklen_t addrlen, int accept_timeout) {
     _server = socket(AF_UNIX, SOCK_SEQPACKET, 0);
@@ -86,7 +88,11 @@ bool FdTransferServer::bindServer(struct sockaddr_un *sun, socklen_t addrlen, in
 bool FdTransferServer::acceptPeer(int *peer_pid) {
     _peer = accept(_server, NULL, NULL);
     if (_peer == -1) {
-        perror("FdTransfer accept()");
+        if (errno == EAGAIN || errno == EWOULDBLOCK) {
+            fprintf(stderr, "[INFO] accept timeout, FdTransferServer exit now.");
+        } else {
+            perror("FdTransfer accept()");
+        }
         return false;
     }
 
@@ -301,21 +307,48 @@ static int single_pid_server(int pid, const char *path) {
         }
     }
 
-    if (!FdTransferServer::bindServer(&sun, addrlen, 10)) {
+    int accept_timeout_seconds = DEFAULT_ACCEPT_TIMEOUT;
+    if (_accept_timeout) {
+        int timeout = Arguments::parseTimeout(_accept_timeout);
+        if (timeout > 0) {
+            accept_timeout_seconds = timeout + DEFAULT_ACCEPT_TIMEOUT;
+        }
+    }
+    fprintf(stderr, "[INFO] set accept timeout to %d seconds\n", accept_timeout_seconds);
+
+    if (!FdTransferServer::bindServer(&sun, addrlen, accept_timeout_seconds)) {
         return 1;
     }
 
+    if (0 != fork()) {
+        // Exit now, let our caller continue.
+        return 0;
+    }
+
     if (!enter_ns(pid, "pid") == -1) {
         fprintf(stderr, "Failed to enter the PID NS of target process %d\n", pid);
         return 1;
     }
 
-    // CLONE_NEWPID affects children only - so we fork here.
-    if (0 == fork()) {
-        return FdTransferServer::acceptPeer(&nspid) && FdTransferServer::serveRequests(nspid) ? 0 : 1;
-    } else {
-        // Exit now, let our caller continue.
-        return 0;
+    // Infinitely wait for subsequent connection until timeout, this is necessary when use "--loop" option.
+    while (1) {
+        if (!FdTransferServer::acceptPeer(&nspid)) {
+            return 1;
+        }
+        // CLONE_NEWPID affects children only - so we fork here.
+        int child = fork();
+        if (-1 == child) {
+            perror("unable to fork");
+            return 1;
+        } else if (0 == child) {
+            return FdTransferServer::serveRequests(nspid) ? 0 : 1;
+        } else {
+            // Exit now, let our caller continue.
+            // return 0;
+            FdTransferServer::closePeer();
+            // Wait subprocess to exit.
+            RESTARTABLE(waitpid(child, NULL, 0));
+        }
     }
 }
 
@@ -370,7 +403,10 @@ static int path_server(const char *path) {
 
 int main(int argc, const char** argv) {
     int pid = 0;
-    if (argc == 3) {
+    if (argc >= 3) {
+        if (argc > 3) {
+            _accept_timeout = (char *)argv[3];
+        }
         pid = atoi(argv[2]);
     } else if (argc != 2) {
         fprintf(stderr, "Usage: %s <path> [<pid>]\n", argv[0]);