1 files changed, 260 insertions, 0 deletions
diff --git a/lib/pullwr.c b/lib/pullwr.c
new file mode 100644
index 0000000..3967eb5
--- /dev/null
+++ b/lib/pullwr.c
@@ -0,0 +1,260 @@
+// SPDX-License-Identifier: GPL-2.0-or-later
+/*
+ * Pull-driven write event handler
+ * Copyright (C) 2019  David Lamparter
+ */
+
+#include "zebra.h"
+
+#include "pullwr.h"
+#include "memory.h"
+#include "monotime.h"
+
+/* defaults */
+#define PULLWR_THRESH	16384	/* size at which we start to call write() */
+#define PULLWR_MAXSPIN	2500	/* max µs to spend grabbing more data */
+
+struct pullwr {
+	int fd;
+	struct event_loop *tm;
+	/* writer == NULL <=> we're idle */
+	struct event *writer;
+
+	void *arg;
+	void (*fill)(void *, struct pullwr *);
+	void (*err)(void *, struct pullwr *, bool);
+
+	/* ring buffer (although it's "un-ringed" on resizing, it WILL wrap
+	 * around if data is trickling in while keeping it at a constant size)
+	 */
+	size_t bufsz, valid, pos;
+	uint64_t total_written;
+	char *buffer;
+
+	size_t thresh;		/* PULLWR_THRESH */
+	int64_t maxspin;	/* PULLWR_MAXSPIN */
+};
+
+DEFINE_MTYPE_STATIC(LIB, PULLWR_HEAD, "pull-driven write controller");
+DEFINE_MTYPE_STATIC(LIB, PULLWR_BUF,  "pull-driven write buffer");
+
+static void pullwr_run(struct event *t);
+
+struct pullwr *_pullwr_new(struct event_loop *tm, int fd, void *arg,
+			   void (*fill)(void *, struct pullwr *),
+			   void (*err)(void *, struct pullwr *, bool))
+{
+	struct pullwr *pullwr;
+
+	pullwr = XCALLOC(MTYPE_PULLWR_HEAD, sizeof(*pullwr));
+	pullwr->fd = fd;
+	pullwr->tm = tm;
+	pullwr->arg = arg;
+	pullwr->fill = fill;
+	pullwr->err = err;
+
+	pullwr->thresh = PULLWR_THRESH;
+	pullwr->maxspin = PULLWR_MAXSPIN;
+
+	return pullwr;
+}
+
+void pullwr_del(struct pullwr *pullwr)
+{
+	EVENT_OFF(pullwr->writer);
+
+	XFREE(MTYPE_PULLWR_BUF, pullwr->buffer);
+	XFREE(MTYPE_PULLWR_HEAD, pullwr);
+}
+
+void pullwr_cfg(struct pullwr *pullwr, int64_t max_spin_usec,
+		size_t write_threshold)
+{
+	pullwr->maxspin = max_spin_usec ?: PULLWR_MAXSPIN;
+	pullwr->thresh = write_threshold ?: PULLWR_THRESH;
+}
+
+void pullwr_bump(struct pullwr *pullwr)
+{
+	if (pullwr->writer)
+		return;
+
+	event_add_timer(pullwr->tm, pullwr_run, pullwr, 0, &pullwr->writer);
+}
+
+static size_t pullwr_iov(struct pullwr *pullwr, struct iovec *iov)
+{
+	size_t len1;
+
+	if (pullwr->valid == 0)
+		return 0;
+
+	if (pullwr->pos + pullwr->valid <= pullwr->bufsz) {
+		iov[0].iov_base = pullwr->buffer + pullwr->pos;
+		iov[0].iov_len = pullwr->valid;
+		return 1;
+	}
+
+	len1 = pullwr->bufsz - pullwr->pos;
+
+	iov[0].iov_base = pullwr->buffer + pullwr->pos;
+	iov[0].iov_len = len1;
+	iov[1].iov_base = pullwr->buffer;
+	iov[1].iov_len = pullwr->valid - len1;
+	return 2;
+}
+
+static void pullwr_resize(struct pullwr *pullwr, size_t need)
+{
+	struct iovec iov[2];
+	size_t niov, newsize;
+	char *newbuf;
+
+	/* the buffer is maintained at pullwr->thresh * 2 since we'll be
+	 * trying to fill it as long as it's anywhere below pullwr->thresh.
+	 * That means we frequently end up a little short of it and then write
+	 * something that goes over the threshold.  So, just use double.
+	 */
+	if (need) {
+		/* resize up */
+		if (pullwr->bufsz - pullwr->valid >= need)
+			return;
+
+		newsize = MAX((pullwr->valid + need) * 2, pullwr->thresh * 2);
+		newbuf = XMALLOC(MTYPE_PULLWR_BUF, newsize);
+	} else if (!pullwr->valid) {
+		/* resize down, buffer empty */
+		newsize = 0;
+		newbuf = NULL;
+	} else {
+		/* resize down */
+		if (pullwr->bufsz - pullwr->valid < pullwr->thresh)
+			return;
+		newsize = MAX(pullwr->valid, pullwr->thresh * 2);
+		newbuf = XMALLOC(MTYPE_PULLWR_BUF, newsize);
+	}
+
+	niov = pullwr_iov(pullwr, iov);
+	if (niov >= 1) {
+		memcpy(newbuf, iov[0].iov_base, iov[0].iov_len);
+		if (niov >= 2)
+			memcpy(newbuf + iov[0].iov_len,
+				iov[1].iov_base, iov[1].iov_len);
+	}
+
+	XFREE(MTYPE_PULLWR_BUF, pullwr->buffer);
+	pullwr->buffer = newbuf;
+	pullwr->bufsz = newsize;
+	pullwr->pos = 0;
+}
+
+void pullwr_write(struct pullwr *pullwr, const void *data, size_t len)
+{
+	pullwr_resize(pullwr, len);
+
+	if (pullwr->pos + pullwr->valid > pullwr->bufsz) {
+		size_t pos;
+
+		pos = (pullwr->pos + pullwr->valid) % pullwr->bufsz;
+		memcpy(pullwr->buffer + pos, data, len);
+	} else {
+		size_t max1, len1;
+		max1 = pullwr->bufsz - (pullwr->pos + pullwr->valid);
+		max1 = MIN(max1, len);
+
+		memcpy(pullwr->buffer + pullwr->pos + pullwr->valid,
+				data, max1);
+		len1 = len - max1;
+
+		if (len1)
+			memcpy(pullwr->buffer, (char *)data + max1, len1);
+
+	}
+	pullwr->valid += len;
+
+	pullwr_bump(pullwr);
+}
+
+static void pullwr_run(struct event *t)
+{
+	struct pullwr *pullwr = EVENT_ARG(t);
+	struct iovec iov[2];
+	size_t niov, lastvalid;
+	ssize_t nwr;
+	struct timeval t0;
+	bool maxspun = false;
+
+	monotime(&t0);
+
+	do {
+		lastvalid = pullwr->valid - 1;
+		while (pullwr->valid < pullwr->thresh
+				&& pullwr->valid != lastvalid
+				&& !maxspun) {
+			lastvalid = pullwr->valid;
+			pullwr->fill(pullwr->arg, pullwr);
+
+			/* check after doing at least one fill() call so we
+			 * don't spin without making progress on slow boxes
+			 */
+			if (!maxspun && monotime_since(&t0, NULL)
+					>= pullwr->maxspin)
+				maxspun = true;
+		}
+
+		if (pullwr->valid == 0) {
+			/* we made a fill() call above that didn't feed any
+			 * data in, and we have nothing more queued, so we go
+			 * into idle, i.e. no calling event_add_write()
+			 */
+			pullwr_resize(pullwr, 0);
+			return;
+		}
+
+		niov = pullwr_iov(pullwr, iov);
+		assert(niov);
+
+		nwr = writev(pullwr->fd, iov, niov);
+		if (nwr < 0) {
+			if (errno == EAGAIN || errno == EWOULDBLOCK)
+				break;
+			pullwr->err(pullwr->arg, pullwr, false);
+			return;
+		}
+
+		if (nwr == 0) {
+			pullwr->err(pullwr->arg, pullwr, true);
+			return;
+		}
+
+		pullwr->total_written += nwr;
+		pullwr->valid -= nwr;
+		pullwr->pos += nwr;
+		pullwr->pos %= pullwr->bufsz;
+	} while (pullwr->valid == 0 && !maxspun);
+	/* pullwr->valid != 0 implies we did an incomplete write, i.e. socket
+	 * is full and we go wait until it's available for writing again.
+	 */
+
+	event_add_write(pullwr->tm, pullwr_run, pullwr, pullwr->fd,
+			&pullwr->writer);
+
+	/* if we hit the time limit, just keep the buffer, we'll probably need
+	 * it anyway & another run is already coming up.
+	 */
+	if (!maxspun)
+		pullwr_resize(pullwr, 0);
+}
+
+void pullwr_stats(struct pullwr *pullwr, uint64_t *total_written,
+		  size_t *pending, size_t *kernel_pending)
+{
+	int tmp;
+
+	*total_written = pullwr->total_written;
+	*pending = pullwr->valid;
+
+	if (ioctl(pullwr->fd, TIOCOUTQ, &tmp) != 0)
+		tmp = 0;
+	*kernel_pending = tmp;
+}