path: root/ext/socket_dontwait/socket_dontwait_ext.c

#include <ruby.h>
#include <ruby/io.h>
#include <errno.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <fcntl.h>
#ifndef MSG_DONTWAIT
#  error MSG_DONTWAIT not defined!
#endif

struct io_args {
	VALUE buf;
	char *ptr;
	long len;
	int fd;
};

#ifndef HAVE_RB_THREAD_IO_BLOCKING_REGION
#  define rb_thread_io_blocking_region(fn,data,fd) \
           rb_thread_blocking_region((fn),(data), RUBY_UBF_IO, 0)
#endif
static long my_tbr(rb_blocking_function_t *fn, void *ptr)
{
	struct io_args *args = (struct io_args *)ptr;

	return (long)rb_thread_io_blocking_region(fn, ptr, args->fd);
}

static int set_blocking(int fd)
{
	int flags = fcntl(fd, F_GETFL);

	if (flags == -1)
		return flags;

	if ((flags & O_NONBLOCK))
		flags = fcntl(fd, F_SETFL, flags & ~O_NONBLOCK);

	return flags;
}

static int can_retry(int fd)
{
	if (errno == EINTR)
		return 1;
	if (errno == EAGAIN)
		return set_blocking(fd) == -1 ? 0 : 1;
	return 0;
}

static void prepare_read_buf(struct io_args *a, VALUE length)
{
	a->len = NUM2LONG(length);
	if (NIL_P(a->buf)) {
		a->buf = rb_str_new(NULL, a->len);
	} else {
		StringValue(a->buf);
		rb_str_modify(a->buf);
		rb_str_resize(a->buf, a->len);
	}
	a->ptr = RSTRING_PTR(a->buf);
}

static void prepare_read_io(struct io_args *a, VALUE io)
{
	rb_io_t *fptr;

	GetOpenFile(io, fptr);
	rb_io_check_readable(fptr);
	a->fd = fptr->fd;
}

static void read_args(struct io_args *a, int argc, VALUE *argv, VALUE io)
{
	VALUE length;

	prepare_read_io(a, io);
	rb_scan_args(argc, argv, "11", &length, &a->buf);
	prepare_read_buf(a, length);
}

static VALUE read_retval(struct io_args *a, long n, const char *msg)
{
	if (n == -1) {
		rb_str_set_len(a->buf, 0);
		rb_sys_fail(msg);
	}
	rb_str_set_len(a->buf, n);
	if (n == 0)
		rb_eof_error();

	return a->buf;
}

/*
 *  call-seq:
 *     ios.read_nonblock(maxlen)              -> string
 *     ios.read_nonblock(maxlen, outbuf)      -> outbuf
 *
 * This behaves like IO#read_nonblock in Ruby core.
 * Unlike IO#read_nonblock, this does not have the side effect of
 * setting the O_NONBLOCK flag on the file descriptor.  It should
 * otherwise behave exactly like IO#read_nonblock when dealing with
 * sockets.
 *
 * Unlike BasicSocket#recv_nonblock, this allows +outbuf+ to be
 * specified and reused to reduce impact on GC.  This method never
 * releases the GVL.
 */
static VALUE read_nonblock(int argc, VALUE *argv, VALUE io)
{
	struct io_args a;
	long n;

	read_args(&a, argc, argv, io);
	if (a.len == 0)
		return a.buf;

	n = (long)recv(a.fd, a.ptr, a.len, MSG_DONTWAIT);

	return read_retval(&a, n, "recv");
}

/* used to implement readpartial if the initial recv() fails with EAGAIN */
static VALUE recv_once(void *arg)
{
	struct io_args *a = arg;
	long n = (long)recv(a->fd, a->ptr, a->len, 0);

	return (VALUE)n;
}

/*
 *  call-seq:
 *     ios.readpartial(maxlen)              -> string
 *     ios.readpartial(maxlen, outbuf)      -> outbuf
 *
 *  This behaves like IO#readpartial from Ruby core.
 *  If data is immediately not available, this will _unset_ the
 *  O_NONBLOCK flag, release the GVL, and block on the socket.
 */
static VALUE readpartial(int argc, VALUE *argv, VALUE io)
{
	struct io_args a;
	long n;

	read_args(&a, argc, argv, io);
	if (a.len == 0)
		return a.buf;

	/* try optimistically first */
	n = (long)recv(a.fd, a.ptr, a.len, MSG_DONTWAIT);

	while (n < 0 && can_retry(a.fd)) {
		rb_str_locktmp(a.buf);
		/* ugh, nothing available: block on the socket */
		n = my_tbr(recv_once, &a);
		rb_str_unlocktmp(a.buf);
	}

	return read_retval(&a, n, "recv");
}

/* used to implement BasicSocket#read */
static VALUE recv_all(void *arg)
{
	struct io_args *a = arg;
	long n;

	errno = 0;
	n = (long)recv(a->fd, a->ptr, a->len, MSG_WAITALL);
	if (n >= 0) {
		a->ptr += n;
		a->len -= n;
	}
	return (VALUE)n;
}

/* used to implement BasicSocket#read() */
static VALUE read_all(struct io_args *a)
{
	int rd_size = 16384;
	long cur_len, n;

	prepare_read_buf(a, INT2FIX(rd_size));

	rb_str_locktmp(a->buf);
	for (;;) {
		do
			n = my_tbr(recv_all, a);
		while ((n < 0 && can_retry(a->fd)) || (n > 0 && a->len > 0));

		cur_len = RSTRING_LEN(a->buf);

		if (n == 0 || a->len > 0) {
			rb_str_unlocktmp(a->buf);
			rb_str_set_len(a->buf, cur_len - a->len);
			return a->buf;
		}

		/* everything was fully read, allocate more */
		rb_str_unlocktmp(a->buf);
		rb_str_resize(a->buf, cur_len + rd_size);
		rb_str_locktmp(a->buf);
		a->ptr = RSTRING_PTR(a->buf) + cur_len;
		a->len = rd_size;
	}
}

/*
 *  call-seq:
 *     ios.read([length [, buffer]])    -> string, buffer, or nil
 *
 *  This behaves like IO#read from Ruby core.
 *  If data is not immediately available, this will _unset_ the
 *  O_NONBLOCK flag, release the GVL, and block on the socket.
 *  This will use the  MSG_WAITALL flag for the recv(2) syscall to
 *  reduce context switching.
 */
static VALUE xread(int argc, VALUE *argv, VALUE io)
{
	struct io_args a;
	VALUE length;
	long n;

	prepare_read_io(&a, io);
	rb_scan_args(argc, argv, "02", &length, &a.buf);

	if (NIL_P(length))
		return read_all(&a);

	prepare_read_buf(&a, length);
	if (a.len == 0)
		return a.buf;

	/* try to read as much as possible without blocking */
	errno = 0;
	do
		n = (long)recv(a.fd, a.ptr, a.len, MSG_DONTWAIT);
	while (n > 0 && (a.ptr += n) && (a.len -= n) > 0);

	/* release the GVL to block on whatever's left */
	rb_str_locktmp(a.buf);
	while (a.len > 0 && n != 0) {
		n = my_tbr(recv_all, &a);
		if (n < 0) {
			if (!can_retry(a.fd))
				break;
		}
	}
	rb_str_unlocktmp(a.buf);
	n = RSTRING_LEN(a.buf) - a.len;
	rb_str_set_len(a.buf, n);
	if (n == 0) {
		if (errno)
			rb_sys_fail("recv");
		return Qnil;
	}

	return a.buf;
}

static void prepare_write_args(struct io_args *a, VALUE io, VALUE str)
{
	rb_io_t *fptr;

	a->buf = (TYPE(str) == T_STRING) ? str : rb_obj_as_string(str);
	a->ptr = RSTRING_PTR(a->buf);
	a->len = RSTRING_LEN(a->buf);
	GetOpenFile(io, fptr);
	rb_io_check_writable(fptr);
	a->fd = fptr->fd;
}

/*
 *  call-seq:
 *     ios.write_nonblock(string)   -> integer
 *
 * This behaves like IO#write_nonblock in Ruby core.
 * Unlike IO#write_nonblock, this does not have the side effect of
 * setting the O_NONBLOCK flag on the file descriptor.  It should
 * otherwise behave exactly like IO#write_nonblock when dealing with
 * sockets.
 *
 * This method never releases the GVL.
 */
static VALUE write_nonblock(VALUE io, VALUE str)
{
	struct io_args a;
	long n;

	prepare_write_args(&a, io, str);
	n = (long)send(a.fd, a.ptr, a.len, MSG_DONTWAIT);
	if (n == -1)
		rb_sys_fail("send");

	return LONG2FIX(n);
}

/* used to implement BasicSocket#write */
static VALUE send_once(void *args)
{
	struct io_args *a = args;
	long n = (long)send(a->fd, a->ptr, a->len, 0);

	if (n >= 0) {
		a->ptr += n;
		a->len -= n;
	}
	return (VALUE)n;
}

/*
 *  call-seq:
 *     ios.write(string)    -> integer
 *
 * This behaves like IO#write in Ruby core.
 *
 * If socket buffer space is not immediately available in the kernel,
 * this will _unset_ the O_NONBLOCK flag, release the GVL, and block
 * on the socket until data is written.
 */
static VALUE xwrite(VALUE io, VALUE str)
{
	struct io_args a;
	long n;

	prepare_write_args(&a, io, str);

	/* optimistically try to send everything w/o releasing GVL */
	n = (long)send(a.fd, a.ptr, a.len, MSG_DONTWAIT);
	if (n == a.len)
		return LONG2FIX(n);

	/* buffer may be expanded in the kernel, keep trying w/o blocking */
	while (n >= 0 && (a.ptr += n) && (a.len -= n) > 0)
		n = (long)send(a.fd, a.ptr, a.len, MSG_DONTWAIT);

	/* all done, we managed to finish without releasing the GVL */
	if (a.len == 0)
		return LONG2FIX(RSTRING_LEN(a.buf));

	if (n < 0 && !can_retry(a.fd))
		rb_sys_fail("send");

	rb_str_locktmp(a.buf);
	while (a.len > 0) {
		n = my_tbr(send_once, &a);
		if (n < 0) {
			if (!can_retry(a.fd))
				break;
		}
	}
	rb_str_unlocktmp(a.buf);
	n = RSTRING_LEN(a.buf) - a.len;
	rb_str_set_len(a.buf, n);

	if (a.len > 0)
		rb_sys_fail("send");
	return LONG2FIX(n);
}

/*
 *  call-seq:
 *     ios.sync = boolean   -> boolean
 *
 * socket_dontwait makes BasicSocket#sync= a no-op.
 * Ruby sockets already default to synchronized operation,
 * and socket_dontwait prevents users from changing this default
 * as it increases complexity.
 */
static VALUE set_sync(VALUE io, VALUE boolean)
{
	return boolean;
}

void Init_socket_dontwait_ext(void)
{
	VALUE mod = rb_define_module("SocketDontwait");

	rb_define_method(mod, "sync=", set_sync, 1);
	rb_define_method(mod, "read", xread, -1);
	rb_define_method(mod, "read_nonblock", read_nonblock, -1);
	rb_define_method(mod, "readpartial", readpartial, -1);
	rb_define_method(mod, "write_nonblock", write_nonblock, 1);
	rb_define_method(mod, "write", xwrite, 1);
}