X-Git-Url: http://git.liburcu.org/?a=blobdiff_plain;f=ustd%2Fustd.c;h=b6b362f4fcc3a44eb5c0cba3e84e6eeb49efaa60;hb=ab805ccdcd0269b28d7288be5e69e9042afbb61c;hp=470401649cdb8c55099a07e98a2946e80653a6ab;hpb=cd226f25dcef5ab527f79ec0ad536a7e6a99d839;p=lttng-ust.git

diff --git a/ustd/ustd.c b/ustd/ustd.c
index 47040164..b6b362f4 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;
@@ -59,12 +65,12 @@ int test_sigpipe(void)
 
 	result = sigemptyset(&sigset);
 	if(result == -1) {
-		perror("sigemptyset");
+		PERROR("sigemptyset");
 		return -1;
 	}
 	result = sigaddset(&sigset, SIGPIPE);
 	if(result == -1) {
-		perror("sigaddset");
+		PERROR("sigaddset");
 		return -1;
 	}
 
@@ -74,7 +80,7 @@ int test_sigpipe(void)
 		return 0;
 	}
 	else if(result == -1) {
-		perror("sigtimedwait");
+		PERROR("sigtimedwait");
 		return -1;
 	}
 	else if(result == SIGPIPE) {
@@ -88,32 +94,35 @@ int test_sigpipe(void)
 
 int get_subbuffer(struct buffer_info *buf)
 {
-	char *send_msg;
-	char *received_msg;
-	char *rep_code;
+	char *send_msg=NULL;
+	char *received_msg=NULL;
+	char *rep_code=NULL;
 	int retval;
 	int result;
 
 	asprintf(&send_msg, "get_subbuffer %s", buf->name);
 	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
-	free(send_msg);
 	if(test_sigpipe()) {
 		WARN("process %d destroyed before we could connect to it", buf->pid);
-		return GET_SUBBUF_DONE;
+		retval = GET_SUBBUF_DONE;
+		goto end;
 	}
 	else if(result < 0) {
 		ERR("get_subbuffer: ustcomm_send_request failed");
-		return -1;
+		retval = -1;
+		goto end;
 	}
 	else if(result == 0) {
 		DBG("app died while being traced");
-		return GET_SUBBUF_DIED;
+		retval = GET_SUBBUF_DIED;
+		goto end;
 	}
 
 	result = sscanf(received_msg, "%as %ld", &rep_code, &buf->consumed_old);
 	if(result != 2 && result != 1) {
 		ERR("unable to parse response to get_subbuffer");
-		return -1;
+		retval = -1;
+		goto end_rep;
 	}
 
 	DBG("received msg is %s", received_msg);
@@ -123,7 +132,8 @@ int get_subbuffer(struct buffer_info *buf)
 		retval = GET_SUBBUF_OK;
 	}
 	else if(nth_token_is(received_msg, "END", 0) == 1) {
-		return GET_SUBBUF_DONE;
+		retval = GET_SUBBUF_DONE;
+		goto end_rep;
 	}
 	else {
 		DBG("error getting subbuffer %s", buf->name);
@@ -131,33 +141,44 @@ int get_subbuffer(struct buffer_info *buf)
 	}
 
 	/* FIMXE: free correctly the stuff */
-	free(received_msg);
-	free(rep_code);
+end_rep:
+	if(rep_code)
+		free(rep_code);
+end:
+	if(send_msg)
+		free(send_msg);
+	if(received_msg)
+		free(received_msg);
+
 	return retval;
 }
 
 int put_subbuffer(struct buffer_info *buf)
 {
-	char *send_msg;
-	char *received_msg;
-	char *rep_code;
+	char *send_msg=NULL;
+	char *received_msg=NULL;
+	char *rep_code=NULL;
 	int retval;
 	int result;
 
 	asprintf(&send_msg, "put_subbuffer %s %ld", buf->name, buf->consumed_old);
 	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
+	if(result < 0 && errno == ECONNRESET) {
+		retval = PUT_SUBBUF_DIED;
+		goto end;
+	}
 	if(result < 0) {
 		ERR("put_subbuffer: send_message failed");
-		return -1;
+		retval = -1;
+		goto end;
 	}
-	free(send_msg);
 
 	result = sscanf(received_msg, "%as", &rep_code);
 	if(result != 1) {
 		ERR("unable to parse response to put_subbuffer");
-		return -1;
+		retval = -1;
+		goto end_rep;
 	}
-	free(received_msg);
 
 	if(!strcmp(rep_code, "OK")) {
 		DBG("subbuffer put %s", buf->name);
@@ -165,10 +186,20 @@ int put_subbuffer(struct buffer_info *buf)
 	}
 	else {
 		DBG("put_subbuffer: received error, we were pushed");
-		return PUT_SUBBUF_PUSHED;
+		retval = PUT_SUBBUF_PUSHED;
+		goto end_rep;
 	}
 
-	free(rep_code);
+end_rep:
+	if(rep_code)
+		free(rep_code);
+
+end:
+	if(send_msg)
+		free(send_msg);
+	if(received_msg)
+		free(received_msg);
+
 	return retval;
 }
 
@@ -196,11 +227,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 +287,8 @@ void *consumer_thread(void *arg)
 
 	/* FIXME: destroy, unalloc... */
 
+	pthread_cleanup_pop(1);
+
 	return NULL;
 }
 
@@ -256,7 +298,7 @@ int create_dir_if_needed(char *dir)
 	result = mkdir(dir, 0777);
 	if(result == -1) {
 		if(errno != EEXIST) {
-			perror("mkdir");
+			PERROR("mkdir");
 			return -1;
 		}
 	}
@@ -309,6 +351,23 @@ int add_buffer(pid_t pid, char *bufname)
 		return -1;
 	}
 
+	/* get pidunique */
+	asprintf(&send_msg, "get_pidunique");
+	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
+	free(send_msg);
+	if(result == -1) {
+		ERR("problem in ustcomm_send_request(get_pidunique)");
+		return -1;
+	}
+
+	result = sscanf(received_msg, "%lld", &buf->pidunique);
+	if(result != 1) {
+		ERR("unable to parse response to get_pidunique");
+		return -1;
+	}
+	free(received_msg);
+	DBG("got pidunique %lld", buf->pidunique);
+
 	/* get shmid */
 	asprintf(&send_msg, "get_shmid %s", buf->name);
 	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
@@ -359,14 +418,14 @@ int add_buffer(pid_t pid, char *bufname)
 	/* attach memory */
 	buf->mem = shmat(buf->shmid, NULL, 0);
 	if(buf->mem == (void *) 0) {
-		perror("shmat");
+		PERROR("shmat");
 		return -1;
 	}
 	DBG("successfully attached buffer memory");
 
 	buf->bufstruct_mem = shmat(buf->bufstruct_shmid, NULL, 0);
 	if(buf->bufstruct_mem == (void *) 0) {
-		perror("shmat");
+		PERROR("shmat");
 		return -1;
 	}
 	DBG("successfully attached buffer bufstruct memory");
@@ -374,7 +433,7 @@ int add_buffer(pid_t pid, char *bufname)
 	/* obtain info on the memory segment */
 	result = shmctl(buf->shmid, IPC_STAT, &shmds);
 	if(result == -1) {
-		perror("shmctl");
+		PERROR("shmctl");
 		return -1;
 	}
 	buf->memlen = shmds.shm_segsz;
@@ -394,7 +453,7 @@ int add_buffer(pid_t pid, char *bufname)
 		trace_path = USTD_DEFAULT_TRACE_PATH;
 	}
 
-	asprintf(&tmp, "%s/%u", trace_path, buf->pid);
+	asprintf(&tmp, "%s/%u_%lld", trace_path, buf->pid, buf->pidunique);
 	result = create_dir_if_needed(tmp);
 	if(result == -1) {
 		ERR("could not create directory %s", tmp);
@@ -403,8 +462,8 @@ int add_buffer(pid_t pid, char *bufname)
 	}
 	free(tmp);
 
-	asprintf(&tmp, "%s/%u/%s_0", trace_path, buf->pid, buf->name);
-	result = fd = open(tmp, O_WRONLY | O_CREAT | O_TRUNC, 00600);
+	asprintf(&tmp, "%s/%u_%lld/%s_0", trace_path, buf->pid, buf->pidunique, buf->name);
+	result = fd = open(tmp, O_WRONLY | O_CREAT | O_TRUNC | O_EXCL, 00600);
 	if(result == -1) {
 		PERROR("open");
 		ERR("failed opening trace file %s", tmp);
@@ -413,6 +472,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 +539,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,19 +576,20 @@ int main(int argc, char **argv)
 		return 1;
 	}
 
+	/* setup handler for SIGPIPE */
 	result = sigemptyset(&sigset);
 	if(result == -1) {
-		perror("sigemptyset");
+		PERROR("sigemptyset");
 		return 1;
 	}
 	result = sigaddset(&sigset, SIGPIPE);
 	if(result == -1) {
-		perror("sigaddset");
+		PERROR("sigaddset");
 		return 1;
 	}
 	result = sigprocmask(SIG_BLOCK, &sigset, NULL);
 	if(result == -1) {
-		perror("sigprocmask");
+		PERROR("sigprocmask");
 		return 1;
 	}
 
@@ -539,6 +623,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;