libbsock/ringbuf.c

#include <sys/uio.h>

#include <errno.h>
#include <unistd.h>

#include "ringbuf.h"

#ifdef min
#undef min
#endif
#define min(a, b)                       \
	({                              \
		__typeof__(a) _a = (a); \
		__typeof__(b) _b = (b); \
		_a < _b ? _a : _b;      \
	})

static inline char *
bsock_ringbuf_wrap_ptr(struct bsock_ringbuf *rb, char *ptr)
{
	while ((uintptr_t)(ptr - rb->buf) > rb->max_sz) {
		ptr = ptr - rb->max_sz;
	}
	return ptr;
}

static void
bsock_ringbuf_peek_nochk(struct bsock_ringbuf *rb, char *buf, size_t len)
{
	assert(len <= rb->sz && len > 0);

	if (rb->end >= rb->start) {
		// no wrapping
		memcpy(buf, rb->start, len);
	} else {
		// wrapping
		size_t rem_size = rb->buf + rb->max_sz - rb->start;
		if (len > rem_size) {
			memcpy(buf, rb->start, rem_size);
			memcpy(buf + rem_size, rb->buf, len - rem_size);
		} else {
			memcpy(buf, rb->start, len);
		}
	}
}

int
bsock_ringbuf_peek(struct bsock_ringbuf *rb, char *buf, size_t len)
{
	if (len == 0) {
		return 0;
	}

	if (len > rb->sz) {
		len = rb->sz;
	}

	if (len > 0) {
		bsock_ringbuf_peek_nochk(rb, buf, len);
	}
	return len;
}

int
bsock_ringbuf_read(struct bsock_ringbuf *rb, char *buf, size_t len)
{
	if (len == 0) {
		return 0;
	}

	if (len > rb->sz) {
		len = rb->sz;
	}

	if (len > 0) {
		bsock_ringbuf_peek_nochk(rb, buf, len);
		rb->start = bsock_ringbuf_wrap_ptr(rb, rb->start + len);
		rb->sz -= len;
	}
	return len;
}

int
bsock_ringbuf_write(struct bsock_ringbuf *rb, void *ctx, struct bsock_ringbuf_io *io, char *buf,
    size_t len)
{
	if (len == 0) {
		return 0;
	}

	size_t free_sz = rb->max_sz - rb->sz;
	if (len > free_sz) {
		int ret = bsock_ringbuf_flush(rb, ctx, io, rb->sz);
		if (ret < 0) {
			return ret;
		}

		free_sz = rb->max_sz - rb->sz;
		if (len > rb->max_sz) {
			return io->write(ctx, buf, len);
		}
	}

	if (rb->end >= rb->start) {
		// no wrapping
		size_t rem_size = rb->buf + rb->max_sz - rb->end;
		if (len > rem_size) {
			memcpy(rb->end, buf, rem_size);
			memcpy(rb->buf, buf + rem_size, len - rem_size);
		} else {
			memcpy(rb->end, buf, len);
		}
	} else {
		memcpy(rb->end, buf, len);
	}

	rb->sz += len;
	rb->end = bsock_ringbuf_wrap_ptr(rb, rb->end + len);

	return len;
}

int
bsock_ringbuf_poll(struct bsock_ringbuf *rb, void *ctx, struct bsock_ringbuf_io *io, size_t len)
{
	size_t free_sz = rb->max_sz - rb->sz;

	if (len == 0 || free_sz == 0) {
		return 0;
	}

	if (len > free_sz) {
		len = free_sz;
	}

	int ret;
	if (rb->end >= rb->start) {
		// no wrapping
		size_t rem_size = rb->buf + rb->max_sz - rb->end;
		if (len > rem_size) {
			struct iovec iov[2];
			iov[0].iov_base = rb->end;
			iov[0].iov_len = rem_size;

			iov[1].iov_base = rb->buf;
			iov[1].iov_len = len - rem_size;
			ret = io->readv(ctx, iov, 2);
		} else {
			ret = io->read(ctx, rb->end, len);
		}
	} else {
		ret = io->read(ctx, rb->end, len);
	}

	if (ret > 0) {
		rb->sz += ret;
		rb->end = bsock_ringbuf_wrap_ptr(rb, rb->end + ret);
	}
	return ret;
}

int
bsock_ringbuf_flush(struct bsock_ringbuf *rb, void *ctx, struct bsock_ringbuf_io *io, size_t len)
{
	if (rb->sz == 0 || len == 0) {
		return 0;
	}

	if (len > rb->sz) {
		len = rb->sz;
	}

	int ret;
	if (rb->end >= rb->start) {
		// no wrapping
		ret = io->write(ctx, rb->start, len);
	} else {
		// there is wrapping
		size_t rem_size = rb->buf + rb->max_sz - rb->start;
		if (len > rem_size) {
			struct iovec iov[2];
			iov[0].iov_base = rb->start;
			iov[0].iov_len = rem_size;

			iov[1].iov_base = rb->buf;
			iov[1].iov_len = len - rem_size;
			ret = io->writev(ctx, iov, 2);
		} else {
			ret = io->write(ctx, rb->start, len);
		}
	}

	if (ret > 0) {
		rb->sz -= ret;
		rb->start = bsock_ringbuf_wrap_ptr(rb, rb->start + ret);
	}
	return ret;
}