[ust.git] / ustd / ustd.c

/* Copyright (C) 2009  Pierre-Marc Fournier
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2.1 of the License, or (at your option) any later version.
 *
 * This library is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public
 * License along with this library; if not, write to the Free Software
 * Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA  02110-1301 USA
 */

#define _GNU_SOURCE

#include <sys/types.h>
#include <sys/stat.h>
#include <sys/shm.h>
#include <fcntl.h>
#include <unistd.h>
#include <pthread.h>
#include <signal.h>

#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <errno.h>
#include <assert.h>
#include <getopt.h>

#include "ustd.h"
#include "localerr.h"
#include "ustcomm.h"

/* return value: 0 = subbuffer is finished, it won't produce data anymore
 *               1 = got subbuffer successfully
 *               <0 = error
 */

#define GET_SUBBUF_OK 1
#define GET_SUBBUF_DONE 0
#define GET_SUBBUF_DIED 2

#define PUT_SUBBUF_OK 1
#define PUT_SUBBUF_DIED 0
#define PUT_SUBBUF_PUSHED 2

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;
	int result;

	result = sigemptyset(&sigset);
	if(result == -1) {
		perror("sigemptyset");
		return -1;
	}
	result = sigaddset(&sigset, SIGPIPE);
	if(result == -1) {
		perror("sigaddset");
		return -1;
	}

	result = sigtimedwait(&sigset, NULL, &(struct timespec){0,0});
	if(result == -1 && errno == EAGAIN) {
		/* no signal received */
		return 0;
	}
	else if(result == -1) {
		perror("sigtimedwait");
		return -1;
	}
	else if(result == SIGPIPE) {
		/* received sigpipe */
		return 1;
	}
	else {
		assert(0);
	}
}

int get_subbuffer(struct buffer_info *buf)
{
	char *send_msg;
	char *received_msg;
	char *rep_code;
	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;
	}
	else if(result < 0) {
		ERR("get_subbuffer: ustcomm_send_request failed");
		return -1;
	}
	else if(result == 0) {
		DBG("app died while being traced");
		return GET_SUBBUF_DIED;
	}

	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;
	}

	DBG("received msg is %s", received_msg);

	if(!strcmp(rep_code, "OK")) {
		DBG("got subbuffer %s", buf->name);
		retval = GET_SUBBUF_OK;
	}
	else if(nth_token_is(received_msg, "END", 0) == 1) {
		return GET_SUBBUF_DONE;
	}
	else {
		DBG("error getting subbuffer %s", buf->name);
		retval = -1;
	}

	/* FIMXE: free correctly the stuff */
	free(received_msg);
	free(rep_code);
	return retval;
}

int put_subbuffer(struct buffer_info *buf)
{
	char *send_msg;
	char *received_msg;
	char *rep_code;
	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) {
		ERR("put_subbuffer: send_message failed");
		return -1;
	}
	free(send_msg);

	result = sscanf(received_msg, "%as", &rep_code);
	if(result != 1) {
		ERR("unable to parse response to put_subbuffer");
		return -1;
	}
	free(received_msg);

	if(!strcmp(rep_code, "OK")) {
		DBG("subbuffer put %s", buf->name);
		retval = PUT_SUBBUF_OK;
	}
	else {
		DBG("put_subbuffer: received error, we were pushed");
		return PUT_SUBBUF_PUSHED;
	}

	free(rep_code);
	return retval;
}

/* This write is patient because it restarts if it was incomplete.
 */

ssize_t patient_write(int fd, const void *buf, size_t count)
{
	const char *bufc = (const char *) buf;
	int result;

	for(;;) {
		result = write(fd, bufc, count);
		if(result <= 0) {
			return result;
		}
		count -= result;
		bufc += result;

		if(count == 0) {
			break;
		}
	}

	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);
		if(result == -1) {
			ERR("error getting subbuffer");
			continue;
		}
		else if(result == GET_SUBBUF_DONE) {
			/* this is done */
			break;
		}
		else if(result == GET_SUBBUF_DIED) {
			finish_consuming_dead_subbuffer(buf);
			break;
		}

		/* write data to file */
		result = patient_write(buf->file_fd, buf->mem + (buf->consumed_old & (buf->n_subbufs * buf->subbuf_size-1)), buf->subbuf_size);
		if(result == -1) {
			PERROR("write");
			/* FIXME: maybe drop this trace */
		}

		/* put the subbuffer */
		result = put_subbuffer(buf);
		if(result == -1) {
			ERR("unknown error putting subbuffer (channel=%s)", buf->name);
			break;
		}
		else if(result == PUT_SUBBUF_PUSHED) {
			ERR("Buffer overflow (channel=%s), reader pushed. This channel will not be usable passed this point.", buf->name);
			break;
		}
		else if(result == PUT_SUBBUF_DIED) {
			WARN("application died while putting subbuffer");
			/* FIXME: probably need to skip the first subbuffer in finish_consuming_dead_subbuffer */
			finish_consuming_dead_subbuffer(buf);
		}
		else if(result == PUT_SUBBUF_OK) {
		}
	}

	DBG("thread for buffer %s is stopping", buf->name);

	/* FIXME: destroy, unalloc... */

	pthread_cleanup_pop(1);

	return NULL;
}

int create_dir_if_needed(char *dir)
{
	int result;
	result = mkdir(dir, 0777);
	if(result == -1) {
		if(errno != EEXIST) {
			perror("mkdir");
			return -1;
		}
	}

	return 0;
}

int is_directory(const char *dir)
{
	int result;
	struct stat st;

	result = stat(dir, &st);
	if(result == -1) {
		PERROR("stat");
		return 0;
	}

	if(!S_ISDIR(st.st_mode)) {
		return 0;
	}

	return 1;
}

int add_buffer(pid_t pid, char *bufname)
{
	struct buffer_info *buf;
	char *send_msg;
	char *received_msg;
	int result;
	char *tmp;
	int fd;
	pthread_t thr;
	struct shmid_ds shmds;

	buf = (struct buffer_info *) malloc(sizeof(struct buffer_info));
	if(buf == NULL) {
		ERR("add_buffer: insufficient memory");
		return -1;
	}

	buf->name = bufname;
	buf->pid = pid;

	/* connect to app */
	result = ustcomm_connect_app(buf->pid, &buf->conn);
	if(result) {
		WARN("unable to connect to process, it probably died before we were able to connect");
		return -1;
	}

	/* get shmid */
	asprintf(&send_msg, "get_shmid %s", buf->name);
	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
	free(send_msg);
	if(result == -1) {
		ERR("problem in ustcomm_send_request(get_shmid)");
		return -1;
	}

	result = sscanf(received_msg, "%d %d", &buf->shmid, &buf->bufstruct_shmid);
	if(result != 2) {
		ERR("unable to parse response to get_shmid");
		return -1;
	}
	free(received_msg);
	DBG("got shmids %d %d", buf->shmid, buf->bufstruct_shmid);

	/* get n_subbufs */
	asprintf(&send_msg, "get_n_subbufs %s", buf->name);
	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
	free(send_msg);
	if(result == -1) {
		ERR("problem in ustcomm_send_request(g_n_subbufs)");
		return -1;
	}

	result = sscanf(received_msg, "%d", &buf->n_subbufs);
	if(result != 1) {
		ERR("unable to parse response to get_n_subbufs");
		return -1;
	}
	free(received_msg);
	DBG("got n_subbufs %d", buf->n_subbufs);

	/* get subbuf size */
	asprintf(&send_msg, "get_subbuf_size %s", buf->name);
	ustcomm_send_request(&buf->conn, send_msg, &received_msg);
	free(send_msg);

	result = sscanf(received_msg, "%d", &buf->subbuf_size);
	if(result != 1) {
		ERR("unable to parse response to get_subbuf_size");
		return -1;
	}
	free(received_msg);
	DBG("got subbuf_size %d", buf->subbuf_size);

	/* attach memory */
	buf->mem = shmat(buf->shmid, NULL, 0);
	if(buf->mem == (void *) 0) {
		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");
		return -1;
	}
	DBG("successfully attached buffer bufstruct memory");

	/* obtain info on the memory segment */
	result = shmctl(buf->shmid, IPC_STAT, &shmds);
	if(result == -1) {
		perror("shmctl");
		return -1;
	}
	buf->memlen = shmds.shm_segsz;

	/* open file for output */
	if(!trace_path) {
		/* Only create the directory if using the default path, because
		 * of the risk of typo when using trace path override. We don't
		 * want to risk creating plenty of useless directories in that case.
		 */
		result = create_dir_if_needed(USTD_DEFAULT_TRACE_PATH);
		if(result == -1) {
			ERR("could not create directory %s", USTD_DEFAULT_TRACE_PATH);
			return -1;
		}

		trace_path = USTD_DEFAULT_TRACE_PATH;
	}

	asprintf(&tmp, "%s/%u", trace_path, buf->pid);
	result = create_dir_if_needed(tmp);
	if(result == -1) {
		ERR("could not create directory %s", tmp);
		free(tmp);
		return -1;
	}
	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);
	if(result == -1) {
		PERROR("open");
		ERR("failed opening trace file %s", tmp);
		return -1;
	}
	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;
}

void usage(void)
{
	fprintf(stderr, "Usage:\nustd OPTIONS\n\nOptions:\n"
			"\t-h\t\tDisplay this usage.\n"
			"\t-o DIR\t\tSpecify the directory where to output the traces.\n"
			"\t-s PATH\t\tSpecify the path to use for the daemon socket.\n");
}

int parse_args(int argc, char **argv)
{
	int c;

	while (1) {
		int option_index = 0;
		static struct option long_options[] = {
			{"help", 0, 0, 'h'},
			{"version", 0, 0, 'V'},
			{0, 0, 0, 0}
		};

		c = getopt_long(argc, argv, "hs:o:", long_options, &option_index);
		if (c == -1)
			break;

		switch (c) {
		case 0:
			printf("option %s", long_options[option_index].name);
			if (optarg)
				printf(" with arg %s", optarg);
			printf("\n");
			break;
		case 's':
			sock_path = optarg;
			break;
		case 'o':
			trace_path = optarg;
			if(!is_directory(trace_path)) {
				ERR("Not a valid directory. (%s)", trace_path);
				return -1;
			}
			break;
		case 'h':
			usage();
			exit(0);
		case 'V':
			printf("Version 0.0\n");
			break;

		default:
			/* unknown option or other error; error is
			printed by getopt, just return */
			return -1;
		}
	}

	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) {
		exit(1);
	}

	result = ustcomm_init_ustd(&ustd, sock_path);
	if(result == -1) {
		ERR("failed to initialize socket");
		return 1;
	}

	/* setup handler for SIGPIPE */
	result = sigemptyset(&sigset);
	if(result == -1) {
		perror("sigemptyset");
		return 1;
	}
	result = sigaddset(&sigset, SIGPIPE);
	if(result == -1) {
		perror("sigaddset");
		return 1;
	}
	result = sigprocmask(SIG_BLOCK, &sigset, NULL);
	if(result == -1) {
		perror("sigprocmask");
		return 1;
	}

	/* app loop */
	for(;;) {
		char *recvbuf;

		/* check for requests on our public socket */
		result = ustcomm_ustd_recv_message(&ustd, &recvbuf, NULL, 100);
		if(result == -1) {
			ERR("error in ustcomm_ustd_recv_message");
			continue;
		}
		if(result > 0) {
			if(!strncmp(recvbuf, "collect", 7)) {
				pid_t pid;
				char *bufname;
				int result;

				result = sscanf(recvbuf, "%*s %d %50as", &pid, &bufname);
				if(result != 2) {
					fprintf(stderr, "parsing error: %s\n", recvbuf);
				}

				result = add_buffer(pid, bufname);
				if(result < 0) {
					ERR("error in add_buffer");
					continue;
				}
			}

			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;
}
Commit	Line	Data
c39c72ee	1	/* Copyright (C) 2009 Pierre-Marc Fournier
1f8b0dff	2	*
c39c72ee PMF	3	* This library is free software; you can redistribute it and/or
	4	* modify it under the terms of the GNU Lesser General Public
	5	* License as published by the Free Software Foundation; either
	6	* version 2.1 of the License, or (at your option) any later version.
1f8b0dff	7	*
c39c72ee	8	* This library is distributed in the hope that it will be useful,
1f8b0dff	9	* but WITHOUT ANY WARRANTY; without even the implied warranty of
c39c72ee PMF	10	* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
c39c72ee PMF	11	* Lesser General Public License for more details.
1f8b0dff	12	*
c39c72ee PMF	13	* You should have received a copy of the GNU Lesser General Public
	14	* License along with this library; if not, write to the Free Software
	15	* Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
1f8b0dff PMF	16	*/
1f8b0dff PMF	17
3796af9b PMF	18	#define _GNU_SOURCE
	19
	20	#include <sys/types.h>
cd226f25	21	#include <sys/stat.h>
3796af9b	22	#include <sys/shm.h>
688760ef PMF	23	#include <fcntl.h>
688760ef PMF	24	#include <unistd.h>
3a7b90de	25	#include <pthread.h>
a3cdd4a7	26	#include <signal.h>
3796af9b PMF	27
	28	#include <stdlib.h>
	29	#include <stdio.h>
	30	#include <string.h>
a3cdd4a7 PMF	31	#include <errno.h>
a3cdd4a7 PMF	32	#include <assert.h>
cd226f25	33	#include <getopt.h>
3796af9b	34
0b0cd937	35	#include "ustd.h"
3796af9b PMF	36	#include "localerr.h"
	37	#include "ustcomm.h"
	38
3a7b90de PMF	39	/* return value: 0 = subbuffer is finished, it won't produce data anymore
	40	* 1 = got subbuffer successfully
	41	* <0 = error
	42	*/
3796af9b	43
8cefc145 PMF	44	#define GET_SUBBUF_OK 1
	45	#define GET_SUBBUF_DONE 0
	46	#define GET_SUBBUF_DIED 2
	47
a3cdd4a7 PMF	48	#define PUT_SUBBUF_OK 1
	49	#define PUT_SUBBUF_DIED 0
	50	#define PUT_SUBBUF_PUSHED 2
	51
c97d4437	52	char *sock_path=NULL;
cd226f25 PMF	53	char *trace_path=NULL;
cd226f25 PMF	54
3158b808 PMF	55	/* Number of active buffers and the mutex to protect it. */
	56	int active_buffers = 0;
	57	pthread_mutex_t active_buffers_mutex = PTHREAD_MUTEX_INITIALIZER;
	58	/* Whether a request to end the program was received. */
	59	sig_atomic_t terminate_req = 0;
	60
a3cdd4a7 PMF	61	int test_sigpipe(void)
	62	{
	63	sigset_t sigset;
	64	int result;
	65
	66	result = sigemptyset(&sigset);
	67	if(result == -1) {
	68	perror("sigemptyset");
	69	return -1;
	70	}
	71	result = sigaddset(&sigset, SIGPIPE);
	72	if(result == -1) {
	73	perror("sigaddset");
	74	return -1;
	75	}
	76
	77	result = sigtimedwait(&sigset, NULL, &(struct timespec){0,0});
	78	if(result == -1 && errno == EAGAIN) {
	79	/* no signal received */
	80	return 0;
	81	}
	82	else if(result == -1) {
	83	perror("sigtimedwait");
	84	return -1;
	85	}
	86	else if(result == SIGPIPE) {
	87	/* received sigpipe */
	88	return 1;
	89	}
	90	else {
	91	assert(0);
	92	}
	93	}
	94
688760ef PMF	95	int get_subbuffer(struct buffer_info *buf)
	96	{
	97	char *send_msg;
	98	char *received_msg;
	99	char *rep_code;
	100	int retval;
	101	int result;
	102
	103	asprintf(&send_msg, "get_subbuffer %s", buf->name);
3bb56863	104	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
8cefc145	105	free(send_msg);
a3cdd4a7 PMF	106	if(test_sigpipe()) {
	107	WARN("process %d destroyed before we could connect to it", buf->pid);
	108	return GET_SUBBUF_DONE;
	109	}
	110	else if(result < 0) {
3bb56863	111	ERR("get_subbuffer: ustcomm_send_request failed");
688760ef PMF	112	return -1;
688760ef PMF	113	}
8cefc145 PMF	114	else if(result == 0) {
	115	DBG("app died while being traced");
	116	return GET_SUBBUF_DIED;
	117	}
688760ef PMF	118
688760ef PMF	119	result = sscanf(received_msg, "%as %ld", &rep_code, &buf->consumed_old);
3a7b90de	120	if(result != 2 && result != 1) {
688760ef PMF	121	ERR("unable to parse response to get_subbuffer");
	122	return -1;
	123	}
3a7b90de PMF	124
3a7b90de PMF	125	DBG("received msg is %s", received_msg);
688760ef PMF	126
	127	if(!strcmp(rep_code, "OK")) {
	128	DBG("got subbuffer %s", buf->name);
8cefc145	129	retval = GET_SUBBUF_OK;
688760ef	130	}
3a7b90de	131	else if(nth_token_is(received_msg, "END", 0) == 1) {
8cefc145	132	return GET_SUBBUF_DONE;
3a7b90de	133	}
688760ef	134	else {
3a7b90de PMF	135	DBG("error getting subbuffer %s", buf->name);
3a7b90de PMF	136	retval = -1;
688760ef PMF	137	}
688760ef PMF	138
3a7b90de PMF	139	/* FIMXE: free correctly the stuff */
3a7b90de PMF	140	free(received_msg);
688760ef PMF	141	free(rep_code);
	142	return retval;
	143	}
	144
	145	int put_subbuffer(struct buffer_info *buf)
	146	{
	147	char *send_msg;
	148	char *received_msg;
	149	char *rep_code;
	150	int retval;
	151	int result;
	152
	153	asprintf(&send_msg, "put_subbuffer %s %ld", buf->name, buf->consumed_old);
3bb56863	154	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
688760ef PMF	155	if(result < 0) {
	156	ERR("put_subbuffer: send_message failed");
	157	return -1;
	158	}
	159	free(send_msg);
	160
	161	result = sscanf(received_msg, "%as", &rep_code);
	162	if(result != 1) {
	163	ERR("unable to parse response to put_subbuffer");
	164	return -1;
	165	}
	166	free(received_msg);
	167
	168	if(!strcmp(rep_code, "OK")) {
	169	DBG("subbuffer put %s", buf->name);
a3cdd4a7	170	retval = PUT_SUBBUF_OK;
688760ef PMF	171	}
688760ef PMF	172	else {
a3cdd4a7 PMF	173	DBG("put_subbuffer: received error, we were pushed");
a3cdd4a7 PMF	174	return PUT_SUBBUF_PUSHED;
688760ef PMF	175	}
	176
	177	free(rep_code);
	178	return retval;
	179	}
	180
a3cdd4a7 PMF	181	/* This write is patient because it restarts if it was incomplete.
	182	*/
	183
688760ef PMF	184	ssize_t patient_write(int fd, const void *buf, size_t count)
	185	{
	186	const char bufc = (const char ) buf;
	187	int result;
	188
	189	for(;;) {
	190	result = write(fd, bufc, count);
	191	if(result <= 0) {
	192	return result;
	193	}
	194	count -= result;
	195	bufc += result;
	196
	197	if(count == 0) {
	198	break;
	199	}
	200	}
	201
	202	return bufc-(const char *)buf;
	203	}
	204
3158b808 PMF	205	void decrement_active_buffers(void *arg)
	206	{
	207	pthread_mutex_lock(&active_buffers_mutex);
	208	active_buffers--;
	209	pthread_mutex_unlock(&active_buffers_mutex);
	210	}
	211
3a7b90de PMF	212	void consumer_thread(void arg)
	213	{
	214	struct buffer_info buf = (struct buffer_info ) arg;
	215	int result;
	216
3158b808 PMF	217	pthread_cleanup_push(decrement_active_buffers, NULL);
3158b808 PMF	218
3a7b90de	219	for(;;) {
8cefc145	220	/* get the subbuffer */
0b0cd937 PMF	221	result = get_subbuffer(buf);
	222	if(result == -1) {
	223	ERR("error getting subbuffer");
	224	continue;
3a7b90de	225	}
0b0cd937 PMF	226	else if(result == GET_SUBBUF_DONE) {
	227	/* this is done */
	228	break;
	229	}
	230	else if(result == GET_SUBBUF_DIED) {
	231	finish_consuming_dead_subbuffer(buf);
	232	break;
3a7b90de PMF	233	}
	234
	235	/* write data to file */
	236	result = patient_write(buf->file_fd, buf->mem + (buf->consumed_old & (buf->n_subbufs * buf->subbuf_size-1)), buf->subbuf_size);
	237	if(result == -1) {
	238	PERROR("write");
	239	/* FIXME: maybe drop this trace */
	240	}
	241
8cefc145	242	/* put the subbuffer */
0b0cd937 PMF	243	result = put_subbuffer(buf);
0b0cd937 PMF	244	if(result == -1) {
a3cdd4a7 PMF	245	ERR("unknown error putting subbuffer (channel=%s)", buf->name);
	246	break;
	247	}
	248	else if(result == PUT_SUBBUF_PUSHED) {
	249	ERR("Buffer overflow (channel=%s), reader pushed. This channel will not be usable passed this point.", buf->name);
0b0cd937	250	break;
3a7b90de	251	}
a3cdd4a7 PMF	252	else if(result == PUT_SUBBUF_DIED) {
	253	WARN("application died while putting subbuffer");
	254	/* FIXME: probably need to skip the first subbuffer in finish_consuming_dead_subbuffer */
	255	finish_consuming_dead_subbuffer(buf);
	256	}
	257	else if(result == PUT_SUBBUF_OK) {
	258	}
3a7b90de PMF	259	}
	260
	261	DBG("thread for buffer %s is stopping", buf->name);
	262
8cefc145 PMF	263	/* FIXME: destroy, unalloc... */
8cefc145 PMF	264
3158b808 PMF	265	pthread_cleanup_pop(1);
3158b808 PMF	266
3a7b90de PMF	267	return NULL;
	268	}
	269
72ebd39a PMF	270	int create_dir_if_needed(char *dir)
	271	{
	272	int result;
	273	result = mkdir(dir, 0777);
	274	if(result == -1) {
	275	if(errno != EEXIST) {
	276	perror("mkdir");
	277	return -1;
	278	}
	279	}
	280
	281	return 0;
	282	}
	283
cd226f25 PMF	284	int is_directory(const char *dir)
	285	{
	286	int result;
	287	struct stat st;
	288
	289	result = stat(dir, &st);
	290	if(result == -1) {
	291	PERROR("stat");
	292	return 0;
	293	}
	294
	295	if(!S_ISDIR(st.st_mode)) {
	296	return 0;
	297	}
	298
	299	return 1;
	300	}
	301
3a7b90de PMF	302	int add_buffer(pid_t pid, char *bufname)
	303	{
	304	struct buffer_info *buf;
	305	char *send_msg;
	306	char *received_msg;
	307	int result;
	308	char *tmp;
	309	int fd;
	310	pthread_t thr;
a3cdd4a7	311	struct shmid_ds shmds;
3a7b90de PMF	312
	313	buf = (struct buffer_info *) malloc(sizeof(struct buffer_info));
	314	if(buf == NULL) {
	315	ERR("add_buffer: insufficient memory");
	316	return -1;
	317	}
	318
	319	buf->name = bufname;
	320	buf->pid = pid;
	321
4e2a8808 PMF	322	/* connect to app */
	323	result = ustcomm_connect_app(buf->pid, &buf->conn);
	324	if(result) {
a3cdd4a7	325	WARN("unable to connect to process, it probably died before we were able to connect");
4e2a8808 PMF	326	return -1;
	327	}
	328
3a7b90de PMF	329	/* get shmid */
3a7b90de PMF	330	asprintf(&send_msg, "get_shmid %s", buf->name);
a3cdd4a7	331	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
3a7b90de	332	free(send_msg);
a3cdd4a7 PMF	333	if(result == -1) {
	334	ERR("problem in ustcomm_send_request(get_shmid)");
	335	return -1;
	336	}
3a7b90de	337
8cefc145 PMF	338	result = sscanf(received_msg, "%d %d", &buf->shmid, &buf->bufstruct_shmid);
8cefc145 PMF	339	if(result != 2) {
3a7b90de PMF	340	ERR("unable to parse response to get_shmid");
	341	return -1;
	342	}
	343	free(received_msg);
8cefc145	344	DBG("got shmids %d %d", buf->shmid, buf->bufstruct_shmid);
3a7b90de PMF	345
	346	/* get n_subbufs */
	347	asprintf(&send_msg, "get_n_subbufs %s", buf->name);
a3cdd4a7	348	result = ustcomm_send_request(&buf->conn, send_msg, &received_msg);
3a7b90de	349	free(send_msg);
a3cdd4a7 PMF	350	if(result == -1) {
	351	ERR("problem in ustcomm_send_request(g_n_subbufs)");
	352	return -1;
	353	}
3a7b90de PMF	354
	355	result = sscanf(received_msg, "%d", &buf->n_subbufs);
	356	if(result != 1) {
	357	ERR("unable to parse response to get_n_subbufs");
	358	return -1;
	359	}
	360	free(received_msg);
	361	DBG("got n_subbufs %d", buf->n_subbufs);
	362
	363	/* get subbuf size */
	364	asprintf(&send_msg, "get_subbuf_size %s", buf->name);
4e2a8808	365	ustcomm_send_request(&buf->conn, send_msg, &received_msg);
3a7b90de PMF	366	free(send_msg);
	367
	368	result = sscanf(received_msg, "%d", &buf->subbuf_size);
	369	if(result != 1) {
	370	ERR("unable to parse response to get_subbuf_size");
	371	return -1;
	372	}
	373	free(received_msg);
	374	DBG("got subbuf_size %d", buf->subbuf_size);
	375
	376	/* attach memory */
	377	buf->mem = shmat(buf->shmid, NULL, 0);
	378	if(buf->mem == (void *) 0) {
	379	perror("shmat");
	380	return -1;
	381	}
8cefc145 PMF	382	DBG("successfully attached buffer memory");
	383
	384	buf->bufstruct_mem = shmat(buf->bufstruct_shmid, NULL, 0);
	385	if(buf->bufstruct_mem == (void *) 0) {
	386	perror("shmat");
	387	return -1;
	388	}
	389	DBG("successfully attached buffer bufstruct memory");
3a7b90de	390
a3cdd4a7 PMF	391	/* obtain info on the memory segment */
	392	result = shmctl(buf->shmid, IPC_STAT, &shmds);
	393	if(result == -1) {
	394	perror("shmctl");
	395	return -1;
	396	}
	397	buf->memlen = shmds.shm_segsz;
	398
3a7b90de	399	/* open file for output */
cd226f25 PMF	400	if(!trace_path) {
	401	/* Only create the directory if using the default path, because
	402	* of the risk of typo when using trace path override. We don't
	403	* want to risk creating plenty of useless directories in that case.
	404	*/
	405	result = create_dir_if_needed(USTD_DEFAULT_TRACE_PATH);
	406	if(result == -1) {
	407	ERR("could not create directory %s", USTD_DEFAULT_TRACE_PATH);
	408	return -1;
	409	}
	410
	411	trace_path = USTD_DEFAULT_TRACE_PATH;
72ebd39a PMF	412	}
72ebd39a PMF	413
cd226f25	414	asprintf(&tmp, "%s/%u", trace_path, buf->pid);
72ebd39a PMF	415	result = create_dir_if_needed(tmp);
	416	if(result == -1) {
	417	ERR("could not create directory %s", tmp);
	418	free(tmp);
	419	return -1;
	420	}
	421	free(tmp);
	422
cd226f25	423	asprintf(&tmp, "%s/%u/%s_0", trace_path, buf->pid, buf->name);
3a7b90de PMF	424	result = fd = open(tmp, O_WRONLY \| O_CREAT \| O_TRUNC, 00600);
	425	if(result == -1) {
	426	PERROR("open");
6cb88bc0	427	ERR("failed opening trace file %s", tmp);
3a7b90de PMF	428	return -1;
	429	}
	430	buf->file_fd = fd;
	431	free(tmp);
	432
3158b808 PMF	433	pthread_mutex_lock(&active_buffers_mutex);
	434	active_buffers++;
	435	pthread_mutex_unlock(&active_buffers_mutex);
	436
3a7b90de PMF	437	pthread_create(&thr, NULL, consumer_thread, buf);
	438
	439	return 0;
	440	}
	441
cd226f25 PMF	442	void usage(void)
	443	{
	444	fprintf(stderr, "Usage:\nustd OPTIONS\n\nOptions:\n"
	445	"\t-h\t\tDisplay this usage.\n"
	446	"\t-o DIR\t\tSpecify the directory where to output the traces.\n"
	447	"\t-s PATH\t\tSpecify the path to use for the daemon socket.\n");
	448	}
	449
	450	int parse_args(int argc, char **argv)
	451	{
	452	int c;
	453
	454	while (1) {
	455	int option_index = 0;
	456	static struct option long_options[] = {
	457	{"help", 0, 0, 'h'},
	458	{"version", 0, 0, 'V'},
	459	{0, 0, 0, 0}
	460	};
	461
	462	c = getopt_long(argc, argv, "hs:o:", long_options, &option_index);
	463	if (c == -1)
	464	break;
	465
	466	switch (c) {
	467	case 0:
	468	printf("option %s", long_options[option_index].name);
	469	if (optarg)
	470	printf(" with arg %s", optarg);
	471	printf("\n");
	472	break;
	473	case 's':
	474	sock_path = optarg;
	475	break;
	476	case 'o':
	477	trace_path = optarg;
	478	if(!is_directory(trace_path)) {
	479	ERR("Not a valid directory. (%s)", trace_path);
	480	return -1;
	481	}
	482	break;
	483	case 'h':
	484	usage();
	485	exit(0);
	486	case 'V':
	487	printf("Version 0.0\n");
	488	break;
	489
	490	default:
	491	/* unknown option or other error; error is
	492	printed by getopt, just return */
	493	return -1;
	494	}
	495	}
	496
	497	return 0;
	498	}
	499
3158b808 PMF	500	void sigterm_handler(int sig)
	501	{
	502	terminate_req = 1;
	503	}
	504
3796af9b PMF	505	int main(int argc, char **argv)
	506	{
	507	struct ustcomm_ustd ustd;
	508	int result;
a3cdd4a7	509	sigset_t sigset;
3158b808 PMF	510	struct sigaction sa;
	511
	512	result = sigemptyset(&sigset);
	513	if(result == -1) {
	514	perror("sigemptyset");
	515	return 1;
	516	}
	517	sa.sa_handler = sigterm_handler;
	518	sa.sa_mask = sigset;
	519	sa.sa_flags = SA_RESTART;
	520	result = sigaction(SIGTERM, &sa, NULL);
	521	if(result == -1) {
	522	PERROR("sigaction");
	523	return 1;
	524	}
3796af9b	525
cd226f25 PMF	526	result = parse_args(argc, argv);
	527	if(result == -1) {
	528	exit(1);
	529	}
	530
	531	result = ustcomm_init_ustd(&ustd, sock_path);
3796af9b PMF	532	if(result == -1) {
	533	ERR("failed to initialize socket");
	534	return 1;
	535	}
	536
3158b808	537	/* setup handler for SIGPIPE */
a3cdd4a7 PMF	538	result = sigemptyset(&sigset);
	539	if(result == -1) {
	540	perror("sigemptyset");
	541	return 1;
	542	}
	543	result = sigaddset(&sigset, SIGPIPE);
	544	if(result == -1) {
	545	perror("sigaddset");
	546	return 1;
	547	}
	548	result = sigprocmask(SIG_BLOCK, &sigset, NULL);
	549	if(result == -1) {
	550	perror("sigprocmask");
	551	return 1;
	552	}
	553
688760ef	554	/* app loop */
3796af9b PMF	555	for(;;) {
	556	char *recvbuf;
	557
3a7b90de	558	/* check for requests on our public socket */
688760ef PMF	559	result = ustcomm_ustd_recv_message(&ustd, &recvbuf, NULL, 100);
	560	if(result == -1) {
	561	ERR("error in ustcomm_ustd_recv_message");
	562	continue;
	563	}
	564	if(result > 0) {
	565	if(!strncmp(recvbuf, "collect", 7)) {
	566	pid_t pid;
	567	char *bufname;
	568	int result;
3796af9b	569
688760ef PMF	570	result = sscanf(recvbuf, "%*s %d %50as", &pid, &bufname);
	571	if(result != 2) {
	572	fprintf(stderr, "parsing error: %s\n", recvbuf);
	573	}
3796af9b	574
688760ef PMF	575	result = add_buffer(pid, bufname);
	576	if(result < 0) {
	577	ERR("error in add_buffer");
	578	continue;
	579	}
3796af9b PMF	580	}
3796af9b PMF	581
688760ef	582	free(recvbuf);
3796af9b	583	}
3158b808 PMF	584
	585	if(terminate_req) {
	586	pthread_mutex_lock(&active_buffers_mutex);
	587	if(active_buffers == 0) {
	588	pthread_mutex_unlock(&active_buffers_mutex);
	589	break;
	590	}
	591	pthread_mutex_unlock(&active_buffers_mutex);
	592	}
3796af9b PMF	593	}
	594
	595	return 0;
	596	}