[PATCH 1/2] add `#kgio_writev` and `#kgio_trywritev`

[Sokolov Yura 'funny-falcon <funny.falcon@gmail.com>]

Add methods for using writev(2) syscall for sending array of string in
a single syscall. This is more efficient than concatenating strings on
Ruby side.
`#kgio_trywritev` returns array of strings which are not sent to the
socket.

Since both methods dups array and strings in it, `#kgio_writev` semantic
a bit different from `#kgio_write`: it does not react on changes to
array/strings that made in other thread. But I think, this way is more
correct.
---
 ext/kgio/extconf.rb   |    2 +
 ext/kgio/read_write.c |  179 +++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 181 insertions(+)

diff --git a/ext/kgio/extconf.rb b/ext/kgio/extconf.rb
index fb680f7..eb5d443 100644
--- a/ext/kgio/extconf.rb
+++ b/ext/kgio/extconf.rb
@@ -23,6 +23,8 @@ have_type("struct sockaddr_storage", %w(sys/types.h sys/socket.h)) or
 have_func('accept4', %w(sys/socket.h))
 have_header("sys/select.h")
 
+have_func("writev", "sys/uio.h")
+
 if have_header('ruby/io.h')
   rubyio = %w(ruby.h ruby/io.h)
   have_struct_member("rb_io_t", "fd", rubyio)
diff --git a/ext/kgio/read_write.c b/ext/kgio/read_write.c
index 51d2d16..3638655 100644
--- a/ext/kgio/read_write.c
+++ b/ext/kgio/read_write.c
@@ -1,6 +1,9 @@
 #include "kgio.h"
 #include "my_fileno.h"
 #include "nonblock.h"
+#ifdef HAVE_WRITEV
+#  include "sys/uio.h"
+#endif
 static VALUE sym_wait_readable, sym_wait_writable;
 static VALUE eErrno_EPIPE, eErrno_ECONNRESET;
 static ID id_set_backtrace;
@@ -403,6 +406,158 @@ static VALUE kgio_trywrite(VALUE io, VALUE str)
 	return my_write(io, str, 0);
 }
 
+struct io_args_v {
+    VALUE io;
+    VALUE buf;
+    VALUE vec_buf;
+    struct iovec *vec;
+    int rest_len;
+    int total_len;
+    int fd;
+};
+
+#ifdef HAVE_WRITEV
+static void prepare_writev(struct io_args_v *a, VALUE io, VALUE ary)
+{
+	long i;
+	a->io = io;
+	a->fd = my_fileno(io);
+	a->buf = rb_ary_dup(ary);
+	a->vec_buf = rb_str_tmp_new(sizeof(struct iovec) * RARRAY_LEN(a->buf));
+	a->vec = (struct iovec*)RSTRING_PTR(a->vec_buf);
+	a->total_len = 0;
+	for(i=0; i < RARRAY_LEN(a->buf); i++) {
+		VALUE str = RARRAY_PTR(a->buf)[i];
+		if (TYPE(str) != T_STRING) {
+			str = rb_obj_as_string(str);
+		} else {
+			str = rb_str_dup_frozen(str);
+		}
+		RARRAY_PTR(a->buf)[i] = str;
+		a->vec[i].iov_base = RSTRING_PTR(str);
+		a->vec[i].iov_len = RSTRING_LEN(str);
+		a->total_len += RSTRING_LEN(str);
+	}
+	a->rest_len = a->total_len;
+}
+
+static int writev_check(struct io_args_v *a, long n, const char *msg, int io_wait)
+{
+	if (a->rest_len == n) {
+done:
+		a->buf = Qnil;
+	} else if (n == -1) {
+		if (errno == EINTR) {
+			a->fd = my_fileno(a->io);
+			return -1;
+		}
+		if (errno == EAGAIN) {
+			if (io_wait) {
+				(void)kgio_call_wait_writable(a->io);
+				return -1;
+			} else if (a->total_len == a->rest_len) {
+				a->buf = sym_wait_writable;
+			}
+			return 0;
+		}
+		wr_sys_fail(msg);
+	} else {
+		assert(n >= 0 && n < a->rest_len && "writev syscall broken?");
+		a->rest_len -= n;
+		while (n > 0) {
+			VALUE str = RARRAY_PTR(a->buf)[0];
+			if (RSTRING_LEN(str) > n) {
+				str = rb_str_subseq(str, n, RSTRING_LEN(str) - n);
+				RARRAY_PTR(a->buf)[0] = str;
+				a->vec->iov_base = RSTRING_PTR(str);
+				a->vec->iov_len = RSTRING_LEN(str);
+				n = 0;
+			} else {
+				n -= RSTRING_LEN(str);
+				rb_ary_shift(a->buf);
+				a->vec++;
+			}
+		}
+		return -1;
+	}
+	return 0;
+}
+
+static VALUE my_writev(VALUE io, VALUE str, int io_wait)
+{
+	struct io_args_v a;
+	long n, iov_cnt, iov_max;
+
+	prepare_writev(&a, io, str);
+	set_nonblocking(a.fd);
+	iov_max = sysconf(_SC_IOV_MAX);
+retry:
+	iov_cnt = RARRAY_LEN(a.buf);
+	if (iov_cnt > iov_max) iov_cnt = iov_max;
+	n = (long)writev(a.fd, a.vec, iov_cnt);
+	if (writev_check(&a, n, "writev", io_wait) != 0)
+		goto retry;
+	if (TYPE(a.buf) != T_SYMBOL)
+		kgio_autopush_write(io);
+	return a.buf;
+}
+#endif
+
+/*
+ * call-seq:
+ *
+ *	io.kgio_writev(array)	-> nil
+ *
+ * Returns nil when the write completes.
+ *
+ * This may block and call any method defined to +kgio_wait_writable+
+ * for the class.
+ *
+ * It fallbacks to kgio_write when writev(2) syscall is missing
+ */
+static VALUE kgio_writev(VALUE io, VALUE ary)
+{
+	VALUE array = rb_check_array_type(ary);
+#ifdef HAVE_WRITEV
+	return my_writev(io, array, 1);
+#else
+	VALUE str = rb_ary_join(array, Qnil);
+	return my_write(io, str, 1);
+#endif
+}
+
+/*
+ * call-seq:
+ *
+ *	io.kgio_trywritev(array)	-> nil, Array or :wait_writable
+ *
+ * Returns nil if the write was completed in full.
+ *
+ * Returns an Array of strings containing the unwritten portion
+ * if EAGAIN was encountered, but some portion was successfully written.
+ *
+ * Returns :wait_writable if EAGAIN is encountered and nothing
+ * was written.
+ *
+ * It fallbacks to kgio_trywrite on joined string when writev(2) syscall
+ * is missing. In this case returned array could contain single joined
+ * string
+ */
+static VALUE kgio_trywritev(VALUE io, VALUE ary)
+{
+	VALUE array = rb_check_array_type(ary);
+#ifdef HAVE_WRITEV
+	return my_writev(io, array, 0);
+#else
+	VALUE str = rb_ary_join(array, Qnil);
+	VALUE result = my_write(io, str, 0);
+	if (TYPE(result) == T_STRING) {
+		return rb_ary_new4(1, &result);
+	}
+	return result;
+#endif
+}
+
 #ifdef USE_MSG_DONTWAIT
 /*
  * This method behaves like Kgio::PipeMethods#kgio_write, except
@@ -485,6 +640,25 @@ static VALUE s_trywrite(VALUE mod, VALUE io, VALUE str)
 {
 	return my_write(io, str, 0);
 }
+/*
+ * call-seq:
+ *
+ *	Kgio.trywritev(io, array)    -> nil, Array or :wait_writable
+ *
+ * Returns nil if the write was completed in full.
+ *
+ * Returns a Array of strings containing the unwritten portion if EAGAIN
+ * was encountered, but some portion was successfully written.
+ *
+ * Returns :wait_writable if EAGAIN is encountered and nothing
+ * was written.
+ *
+ * Maybe used in place of PipeMethods#kgio_trywritev for non-Kgio objects
+ */
+static VALUE s_trywritev(VALUE mod, VALUE io, VALUE ary)
+{
+	return kgio_trywritev(io, ary);
+}
 
 void init_kgio_read_write(void)
 {
@@ -497,6 +671,7 @@ void init_kgio_read_write(void)
 
 	rb_define_singleton_method(mKgio, "tryread", s_tryread, -1);
 	rb_define_singleton_method(mKgio, "trywrite", s_trywrite, 2);
+	rb_define_singleton_method(mKgio, "trywritev", s_trywritev, 2);
 	rb_define_singleton_method(mKgio, "trypeek", s_trypeek, -1);
 
 	/*
@@ -510,8 +685,10 @@ void init_kgio_read_write(void)
 	rb_define_method(mPipeMethods, "kgio_read", kgio_read, -1);
 	rb_define_method(mPipeMethods, "kgio_read!", kgio_read_bang, -1);
 	rb_define_method(mPipeMethods, "kgio_write", kgio_write, 1);
+	rb_define_method(mPipeMethods, "kgio_writev", kgio_writev, 1);
 	rb_define_method(mPipeMethods, "kgio_tryread", kgio_tryread, -1);
 	rb_define_method(mPipeMethods, "kgio_trywrite", kgio_trywrite, 1);
+	rb_define_method(mPipeMethods, "kgio_trywritev", kgio_trywritev, 1);
 
 	/*
 	 * Document-module: Kgio::SocketMethods
@@ -524,8 +701,10 @@ void init_kgio_read_write(void)
 	rb_define_method(mSocketMethods, "kgio_read", kgio_recv, -1);
 	rb_define_method(mSocketMethods, "kgio_read!", kgio_recv_bang, -1);
 	rb_define_method(mSocketMethods, "kgio_write", kgio_send, 1);
+	rb_define_method(mSocketMethods, "kgio_writev", kgio_writev, 1);
 	rb_define_method(mSocketMethods, "kgio_tryread", kgio_tryrecv, -1);
 	rb_define_method(mSocketMethods, "kgio_trywrite", kgio_trysend, 1);
+	rb_define_method(mSocketMethods, "kgio_trywritev", kgio_trywritev, 1);
 	rb_define_method(mSocketMethods, "kgio_trypeek", kgio_trypeek, -1);
 	rb_define_method(mSocketMethods, "kgio_peek", kgio_peek, -1);
 
-- 
1.7.9.5