From 3158b808d3a8c45995ac7a40b022f9d569290b49 Mon Sep 17 00:00:00 2001
From: Pierre-Marc Fournier <pierre-marc.fournier@polymtl.ca>
Date: Wed, 9 Sep 2009 12:15:42 -0400
Subject: [PATCH] ustd: add system to elegantly end daemon on SIGTERM, after
 waiting that all traces are complete

---
 ustd/ustd.c | 51 +++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 51 insertions(+)

diff --git a/ustd/ustd.c b/ustd/ustd.c
index 4704016..fdf8d13 100644
--- a/ustd/ustd.c
+++ b/ustd/ustd.c
@@ -52,6 +52,12 @@
 char *sock_path=NULL;
 char *trace_path=NULL;
 
+/* Number of active buffers and the mutex to protect it. */
+int active_buffers = 0;
+pthread_mutex_t active_buffers_mutex = PTHREAD_MUTEX_INITIALIZER;
+/* Whether a request to end the program was received. */
+sig_atomic_t terminate_req = 0;
+
 int test_sigpipe(void)
 {
 	sigset_t sigset;
@@ -196,11 +202,20 @@ ssize_t patient_write(int fd, const void *buf, size_t count)
 	return bufc-(const char *)buf;
 }
 
+void decrement_active_buffers(void *arg)
+{
+	pthread_mutex_lock(&active_buffers_mutex);
+	active_buffers--;
+	pthread_mutex_unlock(&active_buffers_mutex);
+}
+
 void *consumer_thread(void *arg)
 {
 	struct buffer_info *buf = (struct buffer_info *) arg;
 	int result;
 
+	pthread_cleanup_push(decrement_active_buffers, NULL);
+
 	for(;;) {
 		/* get the subbuffer */
 		result = get_subbuffer(buf);
@@ -247,6 +262,8 @@ void *consumer_thread(void *arg)
 
 	/* FIXME: destroy, unalloc... */
 
+	pthread_cleanup_pop(1);
+
 	return NULL;
 }
 
@@ -413,6 +430,10 @@ int add_buffer(pid_t pid, char *bufname)
 	buf->file_fd = fd;
 	free(tmp);
 
+	pthread_mutex_lock(&active_buffers_mutex);
+	active_buffers++;
+	pthread_mutex_unlock(&active_buffers_mutex);
+
 	pthread_create(&thr, NULL, consumer_thread, buf);
 
 	return 0;
@@ -476,11 +497,31 @@ int parse_args(int argc, char **argv)
 	return 0;
 }
 
+void sigterm_handler(int sig)
+{
+	terminate_req = 1;
+}
+
 int main(int argc, char **argv)
 {
 	struct ustcomm_ustd ustd;
 	int result;
 	sigset_t sigset;
+	struct sigaction sa;
+
+	result = sigemptyset(&sigset);
+	if(result == -1) {
+		perror("sigemptyset");
+		return 1;
+	}
+	sa.sa_handler = sigterm_handler;
+	sa.sa_mask = sigset;
+	sa.sa_flags = SA_RESTART;
+	result = sigaction(SIGTERM, &sa, NULL);
+	if(result == -1) {
+		PERROR("sigaction");
+		return 1;
+	}
 
 	result = parse_args(argc, argv);
 	if(result == -1) {
@@ -493,6 +534,7 @@ int main(int argc, char **argv)
 		return 1;
 	}
 
+	/* setup handler for SIGPIPE */
 	result = sigemptyset(&sigset);
 	if(result == -1) {
 		perror("sigemptyset");
@@ -539,6 +581,15 @@ int main(int argc, char **argv)
 
 			free(recvbuf);
 		}
+
+		if(terminate_req) {
+			pthread_mutex_lock(&active_buffers_mutex);
+			if(active_buffers == 0) {
+				pthread_mutex_unlock(&active_buffers_mutex);
+				break;
+			}
+			pthread_mutex_unlock(&active_buffers_mutex);
+		}
 	}
 
 	return 0;
-- 
2.34.1