1 files changed, 182 insertions, 114 deletions

diff --git a/runtime/staprun/mainloop.c b/runtime/staprun/mainloop.c
index 2e0c3c5c..61963743 100644
--- a/runtime/staprun/mainloop.c
+++ b/runtime/staprun/mainloop.c
@@ -15,39 +15,98 @@
 
 /* globals */
 int ncpus;
-int use_old_transport = 0;
+static int use_old_transport = 0;
+//enum _stp_sig_type { sig_none, sig_done, sig_detach };
+//static enum _stp_sig_type got_signal = sig_none;
 
-static void sigproc(int signum)
+/**
+ *      send_request - send request to kernel over control channel
+ *      @type: the relay-app command id
+ *      @data: pointer to the data to be sent
+ *      @len: length of the data to be sent
+ *
+ *      Returns 0 on success, negative otherwise.
+ */
+int send_request(int type, void *data, int len)
 {
-	dbug(2, "sigproc %d (%s)\n", signum, strsignal(signum));
-
-	if (signum == SIGCHLD) {
-		pid_t pid = waitpid(-1, NULL, WNOHANG);
-		if (pid != target_pid)
-			return;
-		send_request(STP_EXIT, NULL, 0);
-	} else if (signum == SIGQUIT)
-		cleanup_and_exit(2);	
-	else if (signum == SIGINT || signum == SIGHUP || signum == SIGTERM)
-		send_request(STP_EXIT, NULL, 0);
+	char buf[1024];
+
+	/* Before doing memcpy, make sure 'buf' is big enough. */
+	if ((len + 4) > (int)sizeof(buf)) {
+		_err("exceeded maximum send_request size.\n");
+		return -1;
+	}
+	memcpy(buf, &type, 4);
+	memcpy(&buf[4], data, len);
+	return write(control_channel, buf, len + 4);
 }
 
-static void setup_main_signals(int cleanup)
+static void *signal_thread(void *arg)
 {
-	struct sigaction a;
-	memset(&a, 0, sizeof(a));
-	sigfillset(&a.sa_mask);
-	if (cleanup == 0) {
-		a.sa_handler = sigproc;
-		sigaction(SIGCHLD, &a, NULL);
-	} else 
-		a.sa_handler = SIG_IGN;
-	sigaction(SIGINT, &a, NULL);
-	sigaction(SIGTERM, &a, NULL);
-	sigaction(SIGHUP, &a, NULL);
-	sigaction(SIGQUIT, &a, NULL);
+	sigset_t *s = (sigset_t *) arg;
+	int signum, rc, btype = STP_EXIT;
+
+	while (1) {
+		if (sigwait(s, &signum) < 0) {
+			_perr("sigwait");
+			continue;
+		}
+		dbug(2, "sigproc %d (%s)\n", signum, strsignal(signum));
+		if (signum == SIGQUIT)
+			cleanup_and_exit(1);
+		else if (signum == SIGINT || signum == SIGHUP || signum == SIGTERM) {
+			// send STP_EXIT
+			rc = write(control_channel, &btype, sizeof(btype));			
+			break;
+		}
+	}
+	return NULL;
+}
+
+static void chld_proc(int signum)
+{
+	int32_t rc, btype = STP_EXIT;
+	dbug(2, "chld_proc %d (%s)\n", signum, strsignal(signum));
+	pid_t pid = waitpid(-1, NULL, WNOHANG);
+	if (pid != target_pid)
+		return;
+	// send STP_EXIT
+	rc = write(control_channel, &btype, sizeof(btype));
 }
 
+static void setup_main_signals(void)
+{
+	pthread_t tid;
+	struct sigaction sa;
+	sigset_t *s = malloc(sizeof(*s));
+	if (!s) {
+		_perr("malloc failed");
+		exit(1);
+	}
+	sigfillset(s);
+	pthread_sigmask(SIG_SETMASK, s, NULL);
+	memset(&sa, 0, sizeof(sa));
+	sigfillset(&sa.sa_mask);
+	sa.sa_handler = SIG_IGN;
+	sigaction(SIGINT, &sa, NULL);
+	sigaction(SIGTERM, &sa, NULL);
+	sigaction(SIGHUP, &sa, NULL);
+	sigaction(SIGQUIT, &sa, NULL);
+
+	sa.sa_handler = chld_proc;
+	sigaction(SIGCHLD, &sa, NULL);
+
+	sigemptyset(s);
+	sigaddset(s, SIGINT);
+	sigaddset(s, SIGTERM);
+	sigaddset(s, SIGHUP);
+	sigaddset(s, SIGQUIT);
+	pthread_sigmask(SIG_SETMASK, s, NULL);
+	if (pthread_create(&tid, NULL, signal_thread, s) < 0) {
+		_perr("failed to create thread");
+		exit(1);
+	}
+}
 
 /* 
  * start_cmd forks the command given on the command line
@@ -75,7 +134,7 @@ void start_cmd(void)
 	a.sa_handler = SIG_IGN;
 	sigaction(SIGINT, &a, NULL);
 
-	dbug (1, "execing target_cmd %s\n", target_cmd);
+	dbug(1, "execing target_cmd %s\n", target_cmd);
 	if ((pid = fork()) < 0) {
 		_perr("fork");
 		exit(1);
@@ -86,8 +145,8 @@ void start_cmd(void)
 		sigaction(SIGINT, &a, NULL);
 
 		/* commands we fork need to run at normal priority */
-		setpriority (PRIO_PROCESS, 0, 0);
-		
+		setpriority(PRIO_PROCESS, 0, 0);
+
 		/* wait here until signaled */
 		sigwait(&usrset, &signum);
 
@@ -107,11 +166,11 @@ void system_cmd(char *cmd)
 {
 	pid_t pid;
 
-	dbug (2, "system %s\n", cmd);
+	dbug(2, "system %s\n", cmd);
 	if ((pid = fork()) < 0) {
 		_perr("fork");
 	} else if (pid == 0) {
-		setpriority (PRIO_PROCESS, 0, 0);
+		setpriority(PRIO_PROCESS, 0, 0);
 		if (execl("/bin/sh", "sh", "-c", cmd, NULL) < 0)
 			perr("%s", cmd);
 		_exit(1);
@@ -128,7 +187,7 @@ static void read_buffer_info(void)
 	if (!use_old_transport)
 		return;
 
- 	if (statfs("/sys/kernel/debug", &st) == 0 && (int) st.f_type == (int) DEBUGFS_MAGIC)
+	if (statfs("/sys/kernel/debug", &st) == 0 && (int)st.f_type == (int)DEBUGFS_MAGIC)
 		return;
 
 	if (sprintf_chk(buf, "/proc/systemtap/%s/bufsize", modname))
@@ -152,7 +211,6 @@ static void read_buffer_info(void)
 	return;
 }
 
-
 /**
  *	init_stapio - initialize the app
  *	@print_summary: boolean, print summary or not at end of run
@@ -164,7 +222,7 @@ int init_stapio(void)
 	dbug(2, "init_stapio\n");
 
 	/* create control channel */
-	use_old_transport = init_ctl_channel(0);
+	use_old_transport = init_ctl_channel(modname, 1);
 	if (use_old_transport < 0) {
 		err("Failed to initialize control channel.\n");
 		return -1;
@@ -177,7 +235,7 @@ int init_stapio(void)
 			if (init_oldrelayfs() < 0) {
 				close_ctl_channel();
 				return -1;
-			} 
+			}
 		} else {
 			if (init_relayfs() < 0) {
 				close_ctl_channel();
@@ -192,17 +250,12 @@ int init_stapio(void)
 	if (target_cmd)
 		start_cmd();
 
-
 	return 0;
 }
 
-/* cleanup_and_exit() closed channels and frees memory
- * then exits with the following status codes:
- * 1 - failed to initialize.
- * 2 - disconnected
- * 3 - initialized
- */
-void cleanup_and_exit (int closed)
+/* cleanup_and_exit() closed channels, frees memory,
+ * removes the module (if necessary) and exits. */
+void cleanup_and_exit(int detach)
 {
 	pid_t err;
 	static int exiting = 0;
@@ -211,32 +264,34 @@ void cleanup_and_exit (int closed)
 		return;
 	exiting = 1;
 
-	setup_main_signals(1);
+	setup_main_signals();
 
-	dbug(1, "CLEANUP AND EXIT  closed=%d\n", closed);
+	dbug(1, "detach=%d\n", detach);
 
 	/* what about child processes? we will wait for them here. */
 	err = waitpid(-1, NULL, WNOHANG);
 	if (err >= 0)
 		err("\nWaiting for processes to exit\n");
-	while(wait(NULL) > 0) ;
+	while (wait(NULL) > 0) ;
 
 	if (use_old_transport)
-		close_oldrelayfs(closed == 2);
+		close_oldrelayfs(detach);
 	else
 		close_relayfs();
 
 	dbug(1, "closing control channel\n");
 	close_ctl_channel();
 
-	if (initialized == 2 && closed == 2) {
-		err("\nDisconnecting from systemtap module.\n"		\
-		    "To reconnect, type \"staprun -A %s\"\n", modname);
-	} else if (initialized)
-		closed = 3;
-	else
-		closed = 1;
-	exit(closed);
+	if (detach) {
+		err("\nDisconnecting from systemtap module.\n" "To reconnect, type \"staprun -A %s\"\n", modname);
+	} else {
+		dbug(2, "removing %s\n", modname);
+		if (execl(BINDIR "/staprun", "staprun", "-d", modname, NULL) < 0) {
+			perror(modname);
+			_exit(1);
+		}
+	}
+	_exit(0);
 }
 
 /**
@@ -247,90 +302,103 @@ int stp_main_loop(void)
 {
 	ssize_t nb;
 	void *data;
-	int type;
+	uint32_t type;
 	FILE *ofp = stdout;
 	char recvbuf[8196];
 
 	setvbuf(ofp, (char *)NULL, _IOLBF, 0);
-	setup_main_signals(0);
+	setup_main_signals();
 
 	dbug(2, "in main loop\n");
 	send_request(STP_READY, NULL, 0);
 
-	while (1) { /* handle messages from control channel */
+	/* handle messages from control channel */
+	while (1) {
 		nb = read(control_channel, recvbuf, sizeof(recvbuf));
+		dbug(2, "nb=%d\n", (int)nb);
 		if (nb <= 0) {
 			if (errno != EINTR)
 				_perr("Unexpected EOF in read (nb=%ld)", (long)nb);
 			continue;
 		}
-		
-		type = *(int *)recvbuf;
-		data = (void *)(recvbuf + sizeof(int));
 
-		switch (type) { 
+		type = *(uint32_t *) recvbuf;
+		data = (void *)(recvbuf + sizeof(uint32_t));
+		nb -= sizeof(uint32_t);
+
+		switch (type) {
 #ifdef STP_OLD_TRANSPORT
 		case STP_REALTIME_DATA:
-		{
-			ssize_t bw = write(out_fd[0], data, nb - sizeof(int));
-			if (bw >= 0 && bw != (nb - (ssize_t)sizeof(int))) {
-				nb = nb - bw; 
-				bw = write(out_fd[0], data, nb - sizeof(int));
+			{
+				ssize_t bw = write(out_fd[0], data, nb);
+				if (bw >= 0 && bw != nb) {
+					nb = nb - bw;
+					bw = write(out_fd[0], data, nb);
+				}
+				if (bw != nb) {
+					_perr("write error (nb=%ld)", (long)nb);
+					cleanup_and_exit(0);
+				}
+				break;
 			}
-			if (bw != (nb - (ssize_t)sizeof(int))) {
-				_perr("write error (nb=%ld)", (long)nb);
-				cleanup_and_exit(1);
-			}
-                        break;
-		}
 #endif
 		case STP_OOB_DATA:
-			fputs ((char *)data, stderr);
-			break;
-		case STP_EXIT: 
-		{
-			/* module asks us to unload it and exit */
-			int *closed = (int *)data;
-			dbug(2, "got STP_EXIT, closed=%d\n", *closed);
-			cleanup_and_exit(*closed);
+			fputs((char *)data, stderr);
 			break;
-		}
-		case STP_START: 
-		{
-			struct _stp_msg_start *t = (struct _stp_msg_start *)data;
-			dbug(2, "probe_start() returned %d\n", t->res);
-			if (t->res < 0) {
-				if (target_cmd)
-					kill (target_pid, SIGKILL);
-				cleanup_and_exit(1);
-			} else if (target_cmd)
-				kill (target_pid, SIGUSR1);
-			break;
-		}
+		case STP_EXIT:
+			{
+				/* module asks us to unload it and exit */
+				dbug(2, "got STP_EXIT\n");
+				cleanup_and_exit(0);
+				break;
+			}
+		case STP_START:
+			{
+				struct _stp_msg_start *t = (struct _stp_msg_start *)data;
+				dbug(2, "probe_start() returned %d\n", t->res);
+				if (t->res < 0) {
+					if (target_cmd)
+						kill(target_pid, SIGKILL);
+					cleanup_and_exit(0);
+				} else if (target_cmd)
+					kill(target_pid, SIGUSR1);
+				break;
+			}
 		case STP_SYSTEM:
-		{
-			struct _stp_msg_cmd *c = (struct _stp_msg_cmd *)data;
-			dbug(2, "STP_SYSTEM: %s\n", c->cmd);
-			system_cmd(c->cmd);
-			break;
-		}
+			{
+				struct _stp_msg_cmd *c = (struct _stp_msg_cmd *)data;
+				dbug(2, "STP_SYSTEM: %s\n", c->cmd);
+				system_cmd(c->cmd);
+				break;
+			}
 		case STP_TRANSPORT:
-		{
-			struct _stp_msg_start ts;
-			if (use_old_transport) {
-				if (init_oldrelayfs() < 0)
-					cleanup_and_exit(1);
-			} else {
-				if (init_relayfs() < 0)
+			{
+				struct _stp_msg_start ts;
+				if (use_old_transport) {
+					if (init_oldrelayfs() < 0)
+						cleanup_and_exit(0);
+				} else {
+					if (init_relayfs() < 0)
+						cleanup_and_exit(0);
+				}
+				ts.target = target_pid;
+				send_request(STP_START, &ts, sizeof(ts));
+				if (load_only)
 					cleanup_and_exit(1);
+				break;
+			}
+		case STP_UNWIND:
+			{
+				int len;
+				char *ptr = (char *)data;
+				while (nb > 0) {
+					send_unwind_data(ptr);
+					len = strlen(ptr) + 1;
+					ptr += len;
+					nb -= len;
+				}
+				break;
 			}
-			ts.target = target_pid;
-			initialized = 2;
-			send_request(STP_START, &ts, sizeof(ts));
-			if (load_only)
-				cleanup_and_exit(2);
-			break;
-		}
 		default:
 			err("WARNING: ignored message of type %d\n", (type));
 		}