Project

General

Profile

Feature #9323 ยป patch.diff

IO#write accepts multiple arguments - Glass_saga (Masaki Matsushita), 10/21/2017 12:20 PM

View differences:

io.c
1482 1482
    return LONG2FIX(n);
1483 1483
}
1484 1484

  
1485
#ifdef HAVE_WRITEV
1486
struct binwritev_arg {
1487
    rb_io_t *fptr;
1488
    const struct iovec *iov;
1489
    int iovcnt;
1490
};
1491

  
1492
static VALUE
1493
call_writev_internal(VALUE arg)
1494
{
1495
    struct binwritev_arg *p = (struct binwritev_arg *)arg;
1496
    return rb_writev_internal(p->fptr->fd, p->iov, p->iovcnt);
1497
}
1498

  
1499
static long
1500
io_binwritev(struct iovec *iov, int iovcnt, rb_io_t *fptr)
1501
{
1502
    int i;
1503
    long r, total = 0, written_len = 0;
1504

  
1505
    /* don't write anything if current thread has a pending interrupt. */
1506
    rb_thread_check_ints();
1507

  
1508
    if (iovcnt == 0) return 0;
1509
    for (i = 1; i < iovcnt; i++) total += iov[i].iov_len;
1510

  
1511
    if (fptr->wbuf.ptr == NULL && !(fptr->mode & FMODE_SYNC)) {
1512
	fptr->wbuf.off = 0;
1513
	fptr->wbuf.len = 0;
1514
	fptr->wbuf.capa = IO_WBUF_CAPA_MIN;
1515
	fptr->wbuf.ptr = ALLOC_N(char, fptr->wbuf.capa);
1516
	fptr->write_lock = rb_mutex_new();
1517
	rb_mutex_allow_trap(fptr->write_lock, 1);
1518
    }
1519

  
1520
    if (fptr->wbuf.ptr && fptr->wbuf.len) {
1521
	if (fptr->wbuf.off + fptr->wbuf.len + total <= fptr->wbuf.capa) {
1522
	    long offset = fptr->wbuf.off;
1523
	    for (i = 1; i < iovcnt; i++) {
1524
		memcpy(fptr->wbuf.ptr+offset, iov[i].iov_base, iov[i].iov_len);
1525
		offset += iov[i].iov_len;
1526
	    }
1527
	    fptr->wbuf.len += total;
1528
	    return total;
1529
	}
1530
	else {
1531
	    iov[0].iov_base = fptr->wbuf.ptr + fptr->wbuf.off;
1532
	    iov[0].iov_len  = fptr->wbuf.len;
1533
	}
1534
    }
1535
    else {
1536
	iov++;
1537
	iovcnt--;
1538
    }
1539

  
1540
  retry:
1541
    if (fptr->write_lock) {
1542
	struct binwritev_arg arg;
1543
	arg.fptr = fptr;
1544
	arg.iov  = iov;
1545
	arg.iovcnt = iovcnt;
1546
	r = rb_mutex_synchronize(fptr->write_lock, call_writev_internal, (VALUE)&arg);
1547
    }
1548
    else {
1549
	r = rb_writev_internal(fptr->fd, iov, iovcnt);
1550
    }
1551

  
1552
    if (r >= 0) {
1553
	written_len += r;
1554
	if (fptr->wbuf.ptr && fptr->wbuf.len) {
1555
	    if (written_len < fptr->wbuf.len) {
1556
		fptr->wbuf.off += r;
1557
		fptr->wbuf.len -= r;
1558
	    }
1559
	    else {
1560
		fptr->wbuf.off = 0;
1561
		fptr->wbuf.len = 0;
1562
	    }
1563
	}
1564
	if (written_len == total) return written_len;
1565

  
1566
	for (i = 0; i < iovcnt; i++) {
1567
	    if (r > (ssize_t)iov[i].iov_len) {
1568
		r -= iov[i].iov_len;
1569
		iov[i].iov_len = 0;
1570
            }
1571
	    else {
1572
		iov[i].iov_base = (char *)iov[i].iov_base + r;
1573
		iov[i].iov_len  -= r;
1574
		break;
1575
	    }
1576
	}
1577

  
1578
        errno = EAGAIN;
1579
    }
1580
    if (rb_io_wait_writable(fptr->fd)) {
1581
        rb_io_check_closed(fptr);
1582
	goto retry;
1583
    }
1584

  
1585
    return -1L;
1586
}
1587

  
1588
static long
1589
io_fwritev(int argc, VALUE *argv, rb_io_t *fptr)
1590
{
1591
    int i, converted, iovcnt = argc + 1;
1592
    long n;
1593
    VALUE v1, v2, str, tmp, *tmp_array;
1594
    struct iovec *iov;
1595

  
1596
    if (iovcnt > IOV_MAX) {
1597
	rb_raise(rb_eArgError, "too many items (IOV_MAX: %d)", IOV_MAX);
1598
    }
1599

  
1600
    iov = ALLOCV_N(struct iovec, v1, iovcnt);
1601
    tmp_array = ALLOCV_N(VALUE, v2, argc);
1602

  
1603
    for (i = 0; i < argc; i++) {
1604
	str = argv[i];
1605
	converted = 0;
1606
	str = do_writeconv(str, fptr, &converted);
1607
	if (converted)
1608
	    OBJ_FREEZE(str);
1609

  
1610
	tmp = rb_str_tmp_frozen_acquire(str);
1611
	tmp_array[i] = tmp;
1612
	/* iov[0] is reserved for buffer of fptr */
1613
	iov[i+1].iov_base = RSTRING_PTR(tmp);
1614
	iov[i+1].iov_len = RSTRING_LEN(tmp);
1615
    }
1616

  
1617
    n = io_binwritev(iov, iovcnt, fptr);
1618
    if (v1) ALLOCV_END(v1);
1619

  
1620
    for (i = 0; i < argc; i++) {
1621
	rb_str_tmp_frozen_release(argv[i], tmp_array[i]);
1622
    }
1623

  
1624
    if (v2) ALLOCV_END(v2);
1625

  
1626
    return n;
1627
}
1628

  
1629
static VALUE
1630
io_writev(int argc, VALUE *argv, VALUE io)
1631
{
1632
    rb_io_t *fptr;
1633
    long n;
1634
    VALUE tmp;
1635

  
1636
    io = GetWriteIO(io);
1637
    tmp = rb_io_check_io(io);
1638
    if (NIL_P(tmp)) {
1639
	/* port is not IO, call writev method for it. */
1640
	return rb_funcallv(io, id_write, argc, argv);
1641
    }
1642
    io = tmp;
1643

  
1644
    GetOpenFile(io, fptr);
1645
    rb_io_check_writable(fptr);
1646

  
1647
    n = io_fwritev(argc, argv, fptr);
1648
    if (n == -1L) rb_sys_fail_path(fptr->pathv);
1649

  
1650
    return LONG2FIX(n);
1651
}
1652
#else
1653
static VALUE
1654
io_writev(int argc, VALUE *argv, VALUE io)
1655
{
1656
    rb_io_t *fptr;
1657
    long n, total;
1658
    VALUE str, tmp, total = INT2FIX(0);
1659
    int nosync;
1660

  
1661
    io = GetWriteIO(io);
1662
    tmp = rb_io_check_io(io);
1663
    if (NIL_P(tmp)) {
1664
	/* port is not IO, call writev method for it. */
1665
	return rb_funcallv(io, id_write, argc, argv);
1666
    }
1667
    io = tmp;
1668

  
1669
    GetOpenFile(io, fptr);
1670
    rb_io_check_writable(fptr);
1671

  
1672
    for (i = 0; i < argc; i++) {
1673
	/* sync at last item */
1674
	if (i == argc-1)
1675
	    nosync = 0;
1676
	else
1677
	    nosync = 1;
1678

  
1679
	str = argv[i];
1680
	n = io_fwrite(str, fptr, nosync);
1681
	if (n == -1L) rb_sys_fail_path(fptr->pathv);
1682
	total = rb_fix_plus_fix(LONG2FIX(n), total);
1683
    }
1684

  
1685
    return total;
1686
}
1687
#endif /* HAVE_WRITEV */
1688

  
1485 1689
/*
1486 1690
 *  call-seq:
1487
 *     ios.write(string)    -> integer
1691
 *     ios.write(string, ...)    -> integer
1488 1692
 *
1489
 *  Writes the given string to <em>ios</em>. The stream must be opened
1490
 *  for writing. If the argument is not a string, it will be converted
1491
 *  to a string using <code>to_s</code>. Returns the number of bytes
1492
 *  written.
1693
 *  Writes the given strings to <em>ios</em>. The stream must be opened
1694
 *  for writing. If the arguments are not a string, they will be converted
1695
 *  to String using <code>to_s</code>. Returns the number of bytes
1696
 *  written in total.
1493 1697
 *
1494
 *     count = $stdout.write("This is a test\n")
1698
 *     count = $stdout.write("This is", "a test\n")
1495 1699
 *     puts "That was #{count} bytes of data"
1496 1700
 *
1497 1701
 *  <em>produces:</em>
......
1501 1705
 */
1502 1706

  
1503 1707
static VALUE
1504
io_write_m(VALUE io, VALUE str)
1708
io_write_m(int argc, VALUE *argv, VALUE io)
1505 1709
{
1506
    return io_write(io, str, 0);
1710
#ifdef HAVE_WRITEV
1711
    rb_check_arity(argc, 1, IOV_MAX-1);
1712
#else
1713
    rb_check_arity(argc, 1, UNLIMITED_ARGUMENTS);
1714
#endif
1715

  
1716
    if (argc > 1) {
1717
	return io_writev(argc, argv, io);
1718
    }
1719
    else {
1720
	VALUE str = argv[0];
1721
	return io_write(io, str, 0);
1722
    }
1507 1723
}
1508 1724

  
1509 1725
VALUE
......
12674 12890

  
12675 12891
    rb_define_method(rb_cIO, "readpartial",  io_readpartial, -1);
12676 12892
    rb_define_method(rb_cIO, "read",  io_read, -1);
12677
    rb_define_method(rb_cIO, "write", io_write_m, 1);
12893
    rb_define_method(rb_cIO, "write", io_write_m, -1);
12678 12894
    rb_define_method(rb_cIO, "gets",  rb_io_gets_m, -1);
12679 12895
    rb_define_method(rb_cIO, "readline",  rb_io_readline, -1);
12680 12896
    rb_define_method(rb_cIO, "getc",  rb_io_getc, 0);
test/ruby/test_io.rb
1216 1216
    end)
1217 1217
  end
1218 1218

  
1219
  def test_write_with_multiple_arguments
1220
    pipe(proc do |w|
1221
      w.write("foo", "bar")
1222
      w.close
1223
    end, proc do |r|
1224
      assert_equal("foobar", r.read)
1225
    end)
1226
  end
1227

  
1219 1228
  def test_write_non_writable
1220 1229
    with_pipe do |r, w|
1221 1230
      assert_raise(IOError) do