path: root/mcastseed/src/dgrambuf.c

/*
 * dgrambuf.c - C datagrams buffer.
 * 
 * Copyright 2016 by Ludovic Pouzenc <ludovic@pouzenc.fr>
 */
#define _GNU_SOURCE /* See feature_test_macros(7) */
#include "dgrambuf.h"

#include "config.h"

#include <sys/socket.h> /* recvmmsg() _GNU_SOURCE */
#include <stdlib.h> /* calloc(), free(), qsort() */
#include <stdio.h> /* perror() */
#include <errno.h> /* errno */
#include <string.h> /* memset() */
#include <sys/uio.h> /* writev() */
#include <stdint.h> /* uint8_t, uint64_t */
#include <signal.h> /* sigaction() */
#include <unistd.h> /* alarm() */
#include <limits.h> /* SSIZE_MAX */
#include "gl_rbtree_oset.h" /* Red-Black Tree Ordered Set, gnulib-tool --import rbtree-oset */

struct uint_pair {
	unsigned int index;
	unsigned int value;
};

struct dgrambuf_stats_t {
	uint64_t dgrambuf_read_on_full;
	uint64_t recvmmsg_calls, recv_dgrams, recv_byte;
	uint64_t dgram_invalid, dgram_past, dgram_future, dgram_dup, dgram_end_marker;
	uint64_t qsort_calls;
	uint64_t writev_calls, write_partial, write_byte;
};

struct dgrambuf_t {
	/* dgram validation after receive, takes dgram len and a pointer to the start of dgram data
	   Must returns dgram seq number or 0 if invalid dgram */
	int (*validate_func)(unsigned int, void *, unsigned int*);

	struct dgrambuf_stats_t stats;
	struct sigaction sa_sigalrm;

	size_t dgram_slots;
	size_t dgram_free_count;
	size_t dgram_max_size;
	size_t dgram_header_size;

	size_t iovec_slots;
	struct mmsghdr *msgs;
	struct iovec *iov_recv;
	struct iovec *iov_write; /* malloc'ed array */

	struct iovec *partial_write_iov; /* Pointer to an item of iov_write[] */
	size_t partial_write_remaining_iovcnt;
	size_t partial_write_remaining_bytes;

	unsigned int dgram_seq_last;
	unsigned int dgram_seq_base;
	unsigned int *dgram_len;
	unsigned int *dgram_seq_numbers; /* Decoded datagram sequence number for each dgram_slot of buf */
	int dgram_ordered_seq_numbers_is_dirty;
	struct uint_pair *dgram_ordered_seq_numbers; /* Pairs to track original items ranks after qsort() */

	uint8_t *buf; /* malloc-ed 2d byte array : buf[dgram_slots][dgram_max_size] */
};

void _sigalrm_handler(int signum);
int  _compare_uint_pair(const void *pa, const void *pb);
void _update_ordered_seq_numbers(dgrambuf_t dbuf);

#ifndef HAVE_MIN_SIZE_T
size_t min_size_t(size_t a, size_t b) { return (a<b)?a:b; }
#endif /*HAVE_MIN_SIZE_T*/

void dgrambuf_set_validate_func(dgrambuf_t dbuf, int (*validate_func)(unsigned int, void *, unsigned int *)) {
	dbuf->validate_func = validate_func;
}

size_t dgrambuf_get_free_count(const dgrambuf_t dbuf) {
	return dbuf->dgram_free_count;
}

int dgrambuf_everything_was_received(dgrambuf_t dbuf) {
	/*TODO really implement this */
	return dbuf->dgram_seq_last && ( dbuf->dgram_seq_base - 1 == dbuf->dgram_seq_last );
}

ssize_t dgrambuf_recvmmsg(dgrambuf_t dbuf, int sockfd, int timeout, int *info) {
	uint8_t *dgram_base;
	ssize_t recv_byte;
	size_t i, vlen, dgram_index, recv_msg_count;
	int res;
	unsigned int seq, dgram_len;
	struct sigaction sa_old;

	/* Info ptr is mandatory */
	*info = 0;

	/* Validate function is mandatory */
	if ( !dbuf->validate_func ) {
		return -3;
	}

	/* Buffer is full, can't receive */
	if ( dbuf->dgram_free_count == 0 ) {
		dbuf->stats.dgrambuf_read_on_full++;
		*info |= DGRAMBUF_RECV_OVERWRITE;
		/*FIXME : this locks everything if buf full + next seq missing*/
		return 0;
	}

	/* Initialize recvmmsg() syscall arguments */
	for (i=0, vlen=0; i < dbuf->dgram_slots; i++) {
		/*XXX linear search is not optimal, notably if is_dirty == 0*/
		if ( dbuf->dgram_seq_numbers[i] == 0 ) {
			dbuf->iov_recv[vlen].iov_base = dbuf->buf + i*dbuf->dgram_max_size;
			dbuf->iov_recv[vlen].iov_len = dbuf->dgram_max_size;
			memset(dbuf->msgs + vlen, 0, sizeof(struct mmsghdr));
			dbuf->msgs[vlen].msg_hdr.msg_iov = dbuf->iov_recv + vlen;
			dbuf->msgs[vlen].msg_hdr.msg_iovlen = 1;
			vlen++;
			if ( vlen == dbuf->iovec_slots )
				break;
		}
	}

	/* Do the syscall with alarm() to circumvent bad behavior in recvmmsg(2) timeout */
	if (timeout) {
		sigaction(SIGALRM, &(dbuf->sa_sigalrm), &sa_old);
		alarm(timeout);
	}
	res = recvmmsg(sockfd, dbuf->msgs, vlen, MSG_WAITFORONE, NULL);
	if (timeout) {
		alarm(0);
		sigaction(SIGALRM, &sa_old, NULL);
	}
	dbuf->stats.recvmmsg_calls++;

	if (res < 0) {
		if ( errno == EINTR ) {
			recv_msg_count = 0;
			*info |= DGRAMBUF_RECV_EINTR;
		} else {
			perror("recvmmsg()");
			return -1;
		}
	} else {
		recv_msg_count = res;
	}

	if (recv_msg_count > 0) {
		dbuf->dgram_ordered_seq_numbers_is_dirty = 1;
		dbuf->stats.recv_dgrams += recv_msg_count;
		if ( recv_msg_count == vlen ) {
			*info |= DGRAMBUF_RECV_IOVEC_FULL;
		}
	}

	/* Check all received messages */
	for (i=0, recv_byte=0; i<recv_msg_count; i++) {
		dgram_base = dbuf->iov_recv[i].iov_base;
		dgram_index = (dgram_base - dbuf->buf) / dbuf->dgram_max_size;
		dgram_len = dbuf->msgs[i].msg_len;

		/* dgrambuf_new() adjust iovec_len to prevent overflows on ssize_t*/
		recv_byte += dgram_len;

		res = dbuf->validate_func(dgram_len, dgram_base, &seq);
		switch (res) {
			case 1:
				if ( seq < dbuf->dgram_seq_base ) {
					fprintf(stderr, "dgrambuf_recvmmsg(): #%zu past (%u)\n", i, seq);
					dbuf->stats.dgram_past++;
				} else if ( seq >= dbuf->dgram_seq_base + dbuf->dgram_slots ) {
					fprintf(stderr, "dgrambuf_recvmmsg(): #%zu future (%u)\n", i, seq);
					dbuf->stats.dgram_future++;
					*info |= DGRAMBUF_RECV_FUTURE_DGRAM;
				} else {
					/*fprintf(stderr, "dgrambuf_recvmmsg(): #%zu valid (%u)\n", i, seq);*/
					*info |= DGRAMBUF_RECV_VALID_DGRAM;
					dbuf->dgram_seq_numbers[dgram_index] = seq;
					dbuf->dgram_ordered_seq_numbers_is_dirty = 1;
					dbuf->dgram_len[dgram_index] = dgram_len;
					dbuf->dgram_free_count--;
				}
				break;
			case 2:
				fprintf(stderr, "dgrambuf_recvmmsg(): #%zu finalize (%u)\n", i, seq);
				dbuf->stats.dgram_end_marker++;
				dbuf->dgram_seq_last = seq;
				*info |= DGRAMBUF_RECV_FINALIZE;
				break;
			default:
				fprintf(stderr, "dgrambuf_recvmmsg(): #%zu invalid\n", i);
				dbuf->stats.dgram_invalid++;
				break;
		}
	}

	dbuf->stats.recv_byte += recv_byte;

	return recv_byte;
}

int dgrambuf_have_data_ready_to_write(dgrambuf_t dbuf) {
	unsigned int next_dgram_seq;

	/* Last write was partial, so there is more to write */
	if ( dbuf->partial_write_remaining_bytes > 0 ) {
		return 1;
	}

	/* Buffer is empty, nothing to write */
	if ( dbuf->dgram_free_count == dbuf->dgram_slots ) {
		return 0;
	}

	/* Other test cases needs ordered_seq_numbers */
	if ( dbuf->dgram_ordered_seq_numbers_is_dirty ) {
		_update_ordered_seq_numbers(dbuf);
	}

	/* Nothing to write if next dgram is not in buffer at all */
	next_dgram_seq = dbuf->dgram_ordered_seq_numbers[dbuf->dgram_free_count].value;
	if ( next_dgram_seq != dbuf->dgram_seq_base ) {
		return 0;
	}
	/* At least some data of one dgram is availble for writing out */
	return 1;
}

int dgrambuf_have_received_everything(dgrambuf_t dbuf) {
	if (dbuf) {};
	return 0; /*FIXME to be implemented*/
}

ssize_t dgrambuf_write(dgrambuf_t dbuf, int fd, int *info) {
	size_t dgram_index, i, vlen;
	unsigned int curr_seq, prev_seq, dgram_len;
	ssize_t nwrite, total, remain, len;
	struct iovec *iov;

	/* FIXME Info ptr is mandatory */
	*info = 0;

	if ( dbuf->partial_write_remaining_bytes > 0 ) {
		/* Previous writev() was partial, continue it */
		iov = dbuf->partial_write_iov;
		vlen = dbuf->partial_write_remaining_iovcnt;
		total = dbuf->partial_write_remaining_bytes;
	} else if ( ! dgrambuf_have_data_ready_to_write(dbuf) ) {
		return 0; /* XXX Inline code ? */
	} else {
		/* Prepare a write batch, buffer state is in dgram_seq_numbers */
		iov = dbuf->iov_write;
		vlen = 0;
		total = 0;
		/* Write needs up to date ordered_seq_numbers (dgrams could be unsorted or some are lost)*/
		if ( dbuf->dgram_ordered_seq_numbers_is_dirty ) {
			_update_ordered_seq_numbers(dbuf);
		}
		/* Initialize iovecs for writev, take dgram payloads following the sequence numbers */
		prev_seq = 0;
		for (i = dbuf->dgram_free_count; i < dbuf->dgram_slots && vlen < dbuf->iovec_slots; i++) {
			curr_seq = dbuf->dgram_ordered_seq_numbers[i].value;

			/* Skip empty dgram slot */
			if ( curr_seq == 0 ) {
				fprintf(stderr, "Oops : found empty slot (i==%zu)\n", i);
				continue;
			}
			/* Skip if current dgram is a dup of the previous */
			if ( curr_seq == prev_seq ) {
				dbuf->stats.dgram_dup++;
				goto mark_empty;
			}
			/* Skip dgram comming from the past */
			if ( curr_seq < dbuf->dgram_seq_base ) {
				fprintf(stderr, "Oops : found dgram from past in buffer (%u)\n", curr_seq);
				goto mark_empty;
			}
			/* Stop if first dgram to write is not in buffer at all */
			if ( ( vlen==0 ) && (curr_seq != dbuf->dgram_seq_base) ) {
				fprintf(stderr, "Oops : nothing to write, missing %u seq\n", dbuf->dgram_seq_base);
				break;
			}
			/* Stop if current seq dgram is missing */
			if ( ( vlen > 0 ) && (curr_seq > prev_seq+1 ) ) {
				break;
			}

			/* Normal case : curr_seq is the next dgram to write */
			dgram_index = dbuf->dgram_ordered_seq_numbers[i].index;
			dgram_len = dbuf->dgram_len[dgram_index] - dbuf->dgram_header_size;

			/* Setup iovecs */
			dbuf->iov_write[vlen].iov_len = dgram_len;
			dbuf->iov_write[vlen].iov_base = dbuf->buf
				+ dgram_index*dbuf->dgram_max_size + dbuf->dgram_header_size;

			/* Update counters */
			total += dgram_len;
			prev_seq = curr_seq;
			vlen++;

			/* Mark dgram slot about to be written out as empty for next read */
			/*FIXME These cause harm if writev() is incomplete*/
			dbuf->dgram_seq_base = curr_seq + 1;
mark_empty:
			/* Mark slot as empty */
			dgram_index = dbuf->dgram_ordered_seq_numbers[i].index;
			dbuf->dgram_seq_numbers[dgram_index] = 0;
			dbuf->dgram_free_count++;
		}

		/* Nothing valid to write out (but buffer not empty, missing the next dgram) */
		if ( vlen == 0 ) {
			fprintf(stderr, "Oops : nothing to write at all\n");
			return -2;
		}

		if ( vlen == dbuf->iovec_slots ) {
			*info |= DGRAMBUF_WRITE_IOVEC_FULL;
		}
	}

	nwrite = writev(fd, iov, vlen);
	dbuf->stats.writev_calls++;
	if ( nwrite < 0 ) {
		/* Treat non fatal errors */
		if ( errno == EAGAIN || errno == EWOULDBLOCK || errno == EINTR) {
			/* Keeps some state informations for retry */
			dbuf->partial_write_remaining_bytes = total;
			dbuf->partial_write_remaining_iovcnt = vlen;
			dbuf->partial_write_iov = iov;
			*info |= DGRAMBUF_WRITE_EWOULDBLOCK_OR_EINTR;
			return 0;
		}
		/* Print fatal errors and bail out */
		perror("writev()");
		return -1;
	}

	if ( nwrite > 0 ) {
		dbuf->stats.write_byte += nwrite;
		*info |= DGRAMBUF_WRITE_SUCCESS;
	}

	/* Check if the write was partially done */
	dbuf->partial_write_remaining_bytes = total - nwrite;
	if ( dbuf->partial_write_remaining_bytes > 0 ) {
		*info |= DGRAMBUF_WRITE_PARTIAL;
		dbuf->stats.write_partial++;
		/* Find the partially written iov and update it */
		remain = nwrite;
		for (i=0; i<vlen; i++) {
			len = dbuf->iov_write[i].iov_len;
			if ( remain < len ) {
				dbuf->partial_write_remaining_iovcnt = vlen - i;
				dbuf->partial_write_iov = dbuf->iov_write + i;

				dbuf->iov_write[i].iov_base = 
					(uint8_t *) dbuf->iov_write[i].iov_base + remain;
				dbuf->iov_write[i].iov_len -= remain;
				break;
			}
			remain -= len;
		}
		if ( i == vlen ) {
			/* FIXME : this happens */
			fprintf(stderr, "Fatal : failed to find partial iov after partial write\n");
			return -3;
		}

	} else {
		/* Wipe outdated values for clarity in debug mode (only _bytes is use on branching) */
		dbuf->partial_write_iov = NULL;
		dbuf->partial_write_remaining_iovcnt = 0;
	}

	return nwrite;
}

dgrambuf_t dgrambuf_new(size_t dgram_slots, size_t dgram_max_size, size_t dgram_header_size, size_t iovec_slots) {

	dgrambuf_t dbuf = calloc(1, sizeof(struct dgrambuf_t));
	if (!dbuf) goto fail0;

	dbuf->validate_func = NULL;
	/* Implicit with dbuf = calloc(...)
	memset(&(dbuf->stats), 0, sizeof(struct dgrambuf_stats_t));
	memset(&(dbuf->sa_sigalrm), 0, sizeof(struct sigaction));
	*/
	dbuf->sa_sigalrm.sa_handler = _sigalrm_handler;

	dbuf->dgram_slots = dgram_slots;
	dbuf->dgram_free_count = dgram_slots;
	dbuf->dgram_max_size = dgram_max_size;
	dbuf->dgram_header_size = dgram_header_size;

	/* writev() and dgrambuf_recvmmsg accumulates read/write bytes in ssize_t */
	iovec_slots = min_size_t(iovec_slots, SSIZE_MAX/dgram_max_size);
	dbuf->iovec_slots = iovec_slots;

	dbuf->msgs = calloc(iovec_slots, sizeof(struct mmsghdr));
	if (!dbuf->msgs) goto fail1;

	dbuf->iov_recv = calloc(iovec_slots, sizeof(struct iovec));
	if (!dbuf->iov_recv) goto fail2;

	dbuf->iov_write = calloc(iovec_slots, sizeof(struct iovec));
	if (!dbuf->iov_write) goto fail3;

	/* Implicit with dbuf = calloc(...)
	dbuf->partial_write_iov = NULL;
	dbuf->partial_write_remaining_iovcnt = 0;
	dbuf->partial_write_remaining_bytes = 0;

	dbuf->dgram_seq_last = 0;
	*/
	dbuf->dgram_seq_base = 1;
	dbuf->dgram_len = calloc(dgram_slots, sizeof(unsigned int));
	if (!dbuf->dgram_len) goto fail4;

	dbuf->dgram_seq_numbers = calloc(dgram_slots, sizeof(unsigned int));
	if (!dbuf->dgram_seq_numbers) goto fail5;

	dbuf->dgram_ordered_seq_numbers = calloc(dgram_slots, sizeof(struct uint_pair));
	if (!dbuf->dgram_ordered_seq_numbers) goto fail6;
	dbuf->dgram_ordered_seq_numbers_is_dirty = 1;

	dbuf->buf = calloc(dgram_slots, dgram_max_size);
	if (!dbuf->buf) goto fail7;

	return dbuf;

fail7:  free(dbuf->dgram_ordered_seq_numbers);
fail6:	free(dbuf->dgram_seq_numbers);
fail5:	free(dbuf->dgram_len);
fail4:	free(dbuf->iov_write);
fail3:	free(dbuf->iov_recv);
fail2:	free(dbuf->msgs);
fail1:	free(dbuf);
fail0:	return NULL;
}

void dgrambuf_free(dgrambuf_t *dbuf) {
	if (dbuf && *dbuf) {
		free((*dbuf)->buf);
		free((*dbuf)->dgram_ordered_seq_numbers);
		free((*dbuf)->dgram_seq_numbers);
		free((*dbuf)->dgram_len);
		free((*dbuf)->iov_write);
		free((*dbuf)->iov_recv);
		free((*dbuf)->msgs);
		free(*dbuf);
		*dbuf = NULL;
	}
}

void _sigalrm_handler(int signum) {
	/* Nothing to do except interrupting the pending syscall */
	if (signum) {} /* Avoid compiler warning */
}

int _compare_uint_pair(const void *pa, const void *pb) {
	const struct uint_pair *a = pa;
	const struct uint_pair *b = pb;
	if (a->value < b->value)
		return -1;
	else if ( a->value > b->value )
		return 1;
	else
		return 0;
}

void _update_ordered_seq_numbers(dgrambuf_t dbuf) {
	size_t i;
	/* Initialize dgram_ordered_seq_numbers from dgram_seq_numbers */
	for (i=0; i < dbuf->dgram_slots; i++) {
		dbuf->dgram_ordered_seq_numbers[i].index = i;
		dbuf->dgram_ordered_seq_numbers[i].value = dbuf->dgram_seq_numbers[i];
	}
	/* Inplace sorting of dgram_ordered_seq_numbers */
	qsort(dbuf->dgram_ordered_seq_numbers, dbuf->dgram_slots, sizeof(struct uint_pair), _compare_uint_pair);
	dbuf->stats.qsort_calls++;

	dbuf->dgram_ordered_seq_numbers_is_dirty = 0;
}