Move IO write buffer allocation into function.

1 files changed, 119 insertions, 99 deletions

diff --git a/io.c b/io.c
index 095013d39e..179ccc0644 100644
--- a/io.c
+++ b/io.c
@@ -1646,6 +1646,19 @@ io_binwrite_string(VALUE arg)
 }
 #endif
 
+inline static void
+io_allocate_write_buffer(rb_io_t *fptr, int sync)
+{
+    if (fptr->wbuf.ptr == NULL && !(sync && (fptr->mode & FMODE_SYNC))) {
+        fptr->wbuf.off = 0;
+        fptr->wbuf.len = 0;
+        fptr->wbuf.capa = IO_WBUF_CAPA_MIN;
+        fptr->wbuf.ptr = ALLOC_N(char, fptr->wbuf.capa);
+        fptr->write_lock = rb_mutex_new();
+        rb_mutex_allow_trap(fptr->write_lock, 1);
+    }
+}
+
 static long
 io_binwrite(VALUE str, const char *ptr, long len, rb_io_t *fptr, int nosync)
 {
@@ -1656,14 +1669,7 @@ io_binwrite(VALUE str, const char *ptr, long len, rb_io_t *fptr, int nosync)
 
     if ((n = len) <= 0) return n;
 
-    if (fptr->wbuf.ptr == NULL && !(!nosync && (fptr->mode & FMODE_SYNC))) {
-        fptr->wbuf.off = 0;
-        fptr->wbuf.len = 0;
-        fptr->wbuf.capa = IO_WBUF_CAPA_MIN;
-        fptr->wbuf.ptr = ALLOC_N(char, fptr->wbuf.capa);
-        fptr->write_lock = rb_mutex_new();
-        rb_mutex_allow_trap(fptr->write_lock, 1);
-    }
+    io_allocate_write_buffer(fptr, !nosync);
 
     if ((!nosync && (fptr->mode & (FMODE_SYNC|FMODE_TTY))) ||
         (fptr->wbuf.ptr && fptr->wbuf.capa <= fptr->wbuf.len + len)) {
@@ -1865,78 +1871,79 @@ io_binwritev(struct iovec *iov, int iovcnt, rb_io_t *fptr)
     if (iovcnt == 0) return 0;
     for (i = 1; i < iovcnt; i++) total += iov[i].iov_len;
 
-    if (fptr->wbuf.ptr == NULL && !(fptr->mode & FMODE_SYNC)) {
-	fptr->wbuf.off = 0;
-	fptr->wbuf.len = 0;
-	fptr->wbuf.capa = IO_WBUF_CAPA_MIN;
-	fptr->wbuf.ptr = ALLOC_N(char, fptr->wbuf.capa);
-	fptr->write_lock = rb_mutex_new();
-	rb_mutex_allow_trap(fptr->write_lock, 1);
-    }
+    io_allocate_write_buffer(fptr, 1);
 
     if (fptr->wbuf.ptr && fptr->wbuf.len) {
-	long offset = fptr->wbuf.off + fptr->wbuf.len;
-	if (offset + total <= fptr->wbuf.capa) {
-	    for (i = 1; i < iovcnt; i++) {
-		memcpy(fptr->wbuf.ptr+offset, iov[i].iov_base, iov[i].iov_len);
-		offset += iov[i].iov_len;
-	    }
-	    fptr->wbuf.len += total;
-	    return total;
-	}
-	else {
-	    iov[0].iov_base = fptr->wbuf.ptr + fptr->wbuf.off;
-	    iov[0].iov_len  = fptr->wbuf.len;
-	}
+        long offset = fptr->wbuf.off + fptr->wbuf.len;
+        if (offset + total <= fptr->wbuf.capa) {
+            for (i = 1; i < iovcnt; i++) {
+                memcpy(fptr->wbuf.ptr+offset, iov[i].iov_base, iov[i].iov_len);
+                offset += iov[i].iov_len;
+            }
+
+            fptr->wbuf.len += total;
+            return total;
+        }
+        else {
+            iov[0].iov_base = fptr->wbuf.ptr + fptr->wbuf.off;
+            iov[0].iov_len  = fptr->wbuf.len;
+        }
     }
     else {
-	iov++;
-	if (!--iovcnt) return 0;
+        iov++;
+        if (!--iovcnt) return 0;
     }
 
   retry:
     if (fptr->write_lock) {
-	struct binwritev_arg arg;
-	arg.fptr = fptr;
-	arg.iov  = iov;
-	arg.iovcnt = iovcnt;
-	r = rb_mutex_synchronize(fptr->write_lock, call_writev_internal, (VALUE)&arg);
+        struct binwritev_arg arg;
+        arg.fptr = fptr;
+        arg.iov  = iov;
+        arg.iovcnt = iovcnt;
+        r = rb_mutex_synchronize(fptr->write_lock, call_writev_internal, (VALUE)&arg);
     }
     else {
-	r = rb_writev_internal(fptr, iov, iovcnt);
+        r = rb_writev_internal(fptr, iov, iovcnt);
     }
 
     if (r >= 0) {
-	written_len += r;
-	if (fptr->wbuf.ptr && fptr->wbuf.len) {
-	    if (written_len < fptr->wbuf.len) {
-		fptr->wbuf.off += r;
-		fptr->wbuf.len -= r;
-	    }
-	    else {
-		written_len -= fptr->wbuf.len;
-		fptr->wbuf.off = 0;
-		fptr->wbuf.len = 0;
-	    }
-	}
-	if (written_len == total) return total;
+        written_len += r;
+        if (fptr->wbuf.ptr && fptr->wbuf.len) {
+            if (written_len < fptr->wbuf.len) {
+                fptr->wbuf.off += r;
+                fptr->wbuf.len -= r;
+            }
+            else {
+                written_len -= fptr->wbuf.len;
+                fptr->wbuf.off = 0;
+                fptr->wbuf.len = 0;
+            }
+        }
 
-	while (r >= (ssize_t)iov->iov_len) {
-	    /* iovcnt > 0 */
-	    r -= iov->iov_len;
-	    iov->iov_len = 0;
-	    iov++;
-	    if (!--iovcnt) return total;
-	    /* defensive check: written_len should == total */
-	}
-	iov->iov_base = (char *)iov->iov_base + r;
-	iov->iov_len -= r;
+        if (written_len == total) return total;
 
-	errno = EAGAIN;
+        while (r >= (ssize_t)iov->iov_len) {
+            /* iovcnt > 0 */
+            r -= iov->iov_len;
+            iov->iov_len = 0;
+            iov++;
+
+            if (!--iovcnt) {
+                // assert(written_len == total);
+
+                return total;
+            }
+        }
+
+        iov->iov_base = (char *)iov->iov_base + r;
+        iov->iov_len -= r;
+
+        errno = EAGAIN;
     }
+
     if (rb_io_maybe_wait_writable(errno, fptr->self, Qnil)) {
-	rb_io_check_closed(fptr);
-	goto retry;
+        rb_io_check_closed(fptr);
+        goto retry;
     }
 
     return -1L;
@@ -1954,24 +1961,26 @@ io_fwritev(int argc, const VALUE *argv, rb_io_t *fptr)
     tmp_array = ALLOCV_N(VALUE, v2, argc);
 
     for (i = 0; i < argc; i++) {
-	str = rb_obj_as_string(argv[i]);
-	converted = 0;
-	str = do_writeconv(str, fptr, &converted);
-	if (converted)
-	    OBJ_FREEZE(str);
+        str = rb_obj_as_string(argv[i]);
+        converted = 0;
+        str = do_writeconv(str, fptr, &converted);
 
-	tmp = rb_str_tmp_frozen_acquire(str);
-	tmp_array[i] = tmp;
-	/* iov[0] is reserved for buffer of fptr */
-	iov[i+1].iov_base = RSTRING_PTR(tmp);
-	iov[i+1].iov_len = RSTRING_LEN(tmp);
+        if (converted)
+            OBJ_FREEZE(str);
+
+        tmp = rb_str_tmp_frozen_acquire(str);
+        tmp_array[i] = tmp;
+
+        /* iov[0] is reserved for buffer of fptr */
+        iov[i+1].iov_base = RSTRING_PTR(tmp);
+        iov[i+1].iov_len = RSTRING_LEN(tmp);
     }
 
     n = io_binwritev(iov, iovcnt, fptr);
     if (v1) ALLOCV_END(v1);
 
     for (i = 0; i < argc; i++) {
-	rb_str_tmp_frozen_release(argv[i], tmp_array[i]);
+        rb_str_tmp_frozen_release(argv[i], tmp_array[i]);
     }
 
     if (v2) ALLOCV_END(v2);
@@ -2000,10 +2009,12 @@ io_writev(int argc, const VALUE *argv, VALUE io)
 
     io = GetWriteIO(io);
     tmp = rb_io_check_io(io);
+
     if (NIL_P(tmp)) {
-	/* port is not IO, call write method for it. */
-	return rb_funcallv(io, id_write, argc, argv);
+        /* port is not IO, call write method for it. */
+        return rb_funcallv(io, id_write, argc, argv);
     }
+
     io = tmp;
 
     GetOpenFile(io, fptr);
@@ -2011,18 +2022,21 @@ io_writev(int argc, const VALUE *argv, VALUE io)
 
     for (i = 0; i < argc; i += cnt) {
 #ifdef HAVE_WRITEV
-	if ((fptr->mode & (FMODE_SYNC|FMODE_TTY)) && iovcnt_ok(cnt = argc - i)) {
-	    n = io_fwritev(cnt, &argv[i], fptr);
-	}
-	else
+        if ((fptr->mode & (FMODE_SYNC|FMODE_TTY)) && iovcnt_ok(cnt = argc - i)) {
+            n = io_fwritev(cnt, &argv[i], fptr);
+        }
+        else
 #endif
-	{
-	    cnt = 1;
-	    /* sync at last item */
-	    n = io_fwrite(rb_obj_as_string(argv[i]), fptr, (i < argc-1));
-	}
-        if (n < 0L) rb_sys_fail_on_write(fptr);
-	total = rb_fix_plus(LONG2FIX(n), total);
+        {
+            cnt = 1;
+            /* sync at last item */
+            n = io_fwrite(rb_obj_as_string(argv[i]), fptr, (i < argc-1));
+        }
+
+        if (n < 0L)
+            rb_sys_fail_on_write(fptr);
+
+        total = rb_fix_plus(LONG2FIX(n), total);
     }
 
     return total;
@@ -2051,11 +2065,11 @@ static VALUE
 io_write_m(int argc, VALUE *argv, VALUE io)
 {
     if (argc != 1) {
-	return io_writev(argc, argv, io);
+        return io_writev(argc, argv, io);
     }
     else {
-	VALUE str = argv[0];
-	return io_write(io, str, 0);
+        VALUE str = argv[0];
+        return io_write(io, str, 0);
     }
 }
 
@@ -2069,16 +2083,22 @@ static VALUE
 rb_io_writev(VALUE io, int argc, const VALUE *argv)
 {
     if (argc > 1 && rb_obj_method_arity(io, id_write) == 1) {
-	if (io != rb_ractor_stderr() && RTEST(ruby_verbose)) {
-	    VALUE klass = CLASS_OF(io);
-	    char sep = FL_TEST(klass, FL_SINGLETON) ? (klass = io, '.') : '#';
-            rb_category_warning(RB_WARN_CATEGORY_DEPRECATED, "%+"PRIsVALUE"%c""write is outdated interface"
-		       " which accepts just one argument",
-		       klass, sep);
-	}
-	do rb_io_write(io, *argv++); while (--argc);
-	return argv[0];		/* unused right now */
+        if (io != rb_ractor_stderr() && RTEST(ruby_verbose)) {
+            VALUE klass = CLASS_OF(io);
+            char sep = FL_TEST(klass, FL_SINGLETON) ? (klass = io, '.') : '#';
+            rb_category_warning(
+                RB_WARN_CATEGORY_DEPRECATED, "%+"PRIsVALUE"%c""write is outdated interface"
+                " which accepts just one argument",
+                klass, sep
+            );
+        }
+
+        do rb_io_write(io, *argv++); while (--argc);
+
+        /* unused right now */
+        return argv[0];
     }
+
     return rb_funcallv(io, id_write, argc, argv);
 }