cool.io 1.0.0 → 1.9.0

data/ext/cool.io/buffer.c

@@ -0,0 +1,781 @@
+/*
+ * Copyright (C) 2007-12 Tony Arcieri
+ * You may redistribute this under the terms of the MIT license.
+ * See LICENSE for details
+ */
+
+#include "ruby.h"
+#include "ruby/io.h"
+
+#include <assert.h>
+
+#include <string.h>
+#include <time.h>
+#ifndef _MSC_VER
+#include <unistd.h>
+#endif
+#include <errno.h>
+
+/* 1 GiB maximum buffer size */
+#define MAX_BUFFER_SIZE 0x40000000
+
+/* Macro for retrieving the file descriptor from an FPTR */
+#if !HAVE_RB_IO_T_FD
+#define FPTR_TO_FD(fptr) fileno(fptr->f)
+#else
+#define FPTR_TO_FD(fptr) fptr->fd
+#endif
+
+/* Default number of bytes in each node's buffer.  Should be >= MTU */
+#define DEFAULT_NODE_SIZE 16384
+static unsigned default_node_size = DEFAULT_NODE_SIZE;
+
+struct buffer {
+    unsigned size, node_size;
+    struct buffer_node *head, *tail;
+    struct buffer_node *pool_head, *pool_tail;
+
+};
+
+struct buffer_node {
+    unsigned start, end;
+    struct buffer_node *next;
+    unsigned char   data[0];
+};
+
+static VALUE    mCoolio = Qnil;
+static VALUE    cCoolio_Buffer = Qnil;
+
+static VALUE    Coolio_Buffer_allocate(VALUE klass);
+static void     Coolio_Buffer_mark(struct buffer *);
+static void     Coolio_Buffer_free(struct buffer *);
+
+static VALUE    Coolio_Buffer_default_node_size(VALUE klass);
+static VALUE    Coolio_Buffer_set_default_node_size(VALUE klass, VALUE size);
+static VALUE    Coolio_Buffer_initialize(int argc, VALUE * argv, VALUE self);
+static VALUE    Coolio_Buffer_clear(VALUE self);
+static VALUE    Coolio_Buffer_size(VALUE self);
+static VALUE    Coolio_Buffer_empty(VALUE self);
+static VALUE    Coolio_Buffer_append(VALUE self, VALUE data);
+static VALUE    Coolio_Buffer_prepend(VALUE self, VALUE data);
+static VALUE    Coolio_Buffer_read(int argc, VALUE * argv, VALUE self);
+static VALUE    Coolio_Buffer_read_frame(VALUE self, VALUE data, VALUE mark);
+static VALUE    Coolio_Buffer_to_str(VALUE self);
+static VALUE    Coolio_Buffer_read_from(VALUE self, VALUE io);
+static VALUE    Coolio_Buffer_write_to(VALUE self, VALUE io);
+
+static struct buffer *buffer_new(void);
+static void     buffer_clear(struct buffer * buf);
+static void     buffer_free(struct buffer * buf);
+static void     buffer_free_pool(struct buffer * buf);
+static void     buffer_prepend(struct buffer * buf, char *str, unsigned len);
+static void     buffer_append(struct buffer * buf, char *str, unsigned len);
+static void     buffer_read(struct buffer * buf, char *str, unsigned len);
+static int      buffer_read_frame(struct buffer * buf, VALUE str, char frame_mark);
+static void     buffer_copy(struct buffer * buf, char *str, unsigned len);
+static int      buffer_read_from(struct buffer * buf, int fd);
+static int      buffer_write_to(struct buffer * buf, int fd);
+
+/*
+ * High-performance I/O buffer intended for use in non-blocking programs
+ *
+ * Data is stored in as a memory-pooled linked list of equally sized chunks.
+ * Routines are provided for high speed non-blocking reads and writes from
+ * Ruby IO objects.
+ */
+void
+Init_coolio_buffer()
+{
+    VALUE cCoolio_IO;
+
+    mCoolio = rb_define_module("Coolio");
+    cCoolio_Buffer = rb_define_class_under(mCoolio, "Buffer", rb_cObject);
+    rb_define_alloc_func(cCoolio_Buffer, Coolio_Buffer_allocate);
+
+    rb_define_singleton_method(cCoolio_Buffer, "default_node_size",
+                   Coolio_Buffer_default_node_size, 0);
+    rb_define_singleton_method(cCoolio_Buffer, "default_node_size=",
+                   Coolio_Buffer_set_default_node_size, 1);
+
+    rb_define_method(cCoolio_Buffer, "initialize", Coolio_Buffer_initialize, -1);
+    rb_define_method(cCoolio_Buffer, "clear", Coolio_Buffer_clear, 0);
+    rb_define_method(cCoolio_Buffer, "size", Coolio_Buffer_size, 0);
+    rb_define_method(cCoolio_Buffer, "empty?", Coolio_Buffer_empty, 0);
+    rb_define_method(cCoolio_Buffer, "<<", Coolio_Buffer_append, 1);
+    rb_define_method(cCoolio_Buffer, "append", Coolio_Buffer_append, 1);
+    rb_define_method(cCoolio_Buffer, "write", Coolio_Buffer_append, 1);
+    rb_define_method(cCoolio_Buffer, "prepend", Coolio_Buffer_prepend, 1);
+    rb_define_method(cCoolio_Buffer, "read", Coolio_Buffer_read, -1);
+    rb_define_method(cCoolio_Buffer, "read_frame", Coolio_Buffer_read_frame, 2);
+    rb_define_method(cCoolio_Buffer, "to_str", Coolio_Buffer_to_str, 0);
+    rb_define_method(cCoolio_Buffer, "read_from", Coolio_Buffer_read_from, 1);
+    rb_define_method(cCoolio_Buffer, "write_to", Coolio_Buffer_write_to, 1);
+
+    rb_define_const(cCoolio_Buffer, "MAX_SIZE", INT2NUM(MAX_BUFFER_SIZE));
+}
+
+static VALUE
+Coolio_Buffer_allocate(VALUE klass)
+{
+    return Data_Wrap_Struct(klass, Coolio_Buffer_mark, Coolio_Buffer_free, buffer_new());
+}
+
+static void
+Coolio_Buffer_mark(struct buffer * buf)
+{
+    /* Naively discard the memory pool whenever Ruby garbage collects */
+    buffer_free_pool(buf);
+}
+
+static void
+Coolio_Buffer_free(struct buffer * buf)
+{
+    buffer_free(buf);
+}
+
+/**
+ * call-seq:
+ *   Coolio::Buffer.default_node_size -> 4096
+ *
+ * Retrieves the current value of the default node size.
+ */
+static VALUE
+Coolio_Buffer_default_node_size(VALUE klass)
+{
+    return UINT2NUM(default_node_size);
+}
+
+/*
+ * safely converts node sizes from Ruby numerics to C and raising
+ * ArgumentError or RangeError on invalid sizes
+ */
+static unsigned
+convert_node_size(VALUE size)
+{
+    if (
+        rb_funcall(size, rb_intern("<"), 1, INT2NUM(1)) == Qtrue ||
+        rb_funcall(size, rb_intern(">"), 1, INT2NUM(MAX_BUFFER_SIZE)) == Qtrue
+    )
+        rb_raise(rb_eArgError, "invalid buffer size");
+
+    return (unsigned) NUM2INT(size);
+}
+
+/**
+ * call-seq:
+ *   Coolio::Buffer.default_node_size = 16384
+ *
+ * Sets the default node size for calling Coolio::Buffer.new with no arguments.
+ */
+static VALUE
+Coolio_Buffer_set_default_node_size(VALUE klass, VALUE size)
+{
+    default_node_size = convert_node_size(size);
+
+    return size;
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer.new(size = Coolio::Buffer.default_node_size) -> Coolio::Buffer
+ *
+ * Create a new Coolio::Buffer with linked segments of the given size
+ */
+static VALUE
+Coolio_Buffer_initialize(int argc, VALUE * argv, VALUE self)
+{
+    VALUE           node_size_obj;
+    struct buffer *buf;
+
+    if (rb_scan_args(argc, argv, "01", &node_size_obj) == 1) {
+        Data_Get_Struct(self, struct buffer, buf);
+
+        /*
+         * Make sure we're not changing the buffer size after data
+         * has been allocated
+         */
+        assert(!buf->head);
+        assert(!buf->pool_head);
+
+        buf->node_size = convert_node_size(node_size_obj);
+    }
+    return Qnil;
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#clear -> nil
+ *
+ * Clear all data from the Coolio::Buffer
+ */
+static VALUE
+Coolio_Buffer_clear(VALUE self)
+{
+    struct buffer *buf;
+    Data_Get_Struct(self, struct buffer, buf);
+
+    buffer_clear(buf);
+
+    return Qnil;
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#size -> Integer
+ *
+ * Return the size of the buffer in bytes
+ */
+static VALUE
+Coolio_Buffer_size(VALUE self)
+{
+    struct buffer *buf;
+    Data_Get_Struct(self, struct buffer, buf);
+
+    return INT2NUM(buf->size);
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#empty? -> Boolean
+ *
+ * Is the buffer empty?
+ */
+static VALUE
+Coolio_Buffer_empty(VALUE self)
+{
+    struct buffer *buf;
+    Data_Get_Struct(self, struct buffer, buf);
+
+    return buf->size > 0 ? Qfalse : Qtrue;
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#append(data) -> String
+ *
+ * Append the given data to the end of the buffer
+ */
+static VALUE
+Coolio_Buffer_append(VALUE self, VALUE data)
+{
+    struct buffer *buf;
+    Data_Get_Struct(self, struct buffer, buf);
+
+    /* Is this needed?  Never seen anyone else do it... */
+    data = rb_convert_type(data, T_STRING, "String", "to_str");
+    buffer_append(buf, RSTRING_PTR(data), RSTRING_LEN(data));
+
+    return data;
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#prepend(data) -> String
+ *
+ * Prepend the given data to the beginning of the buffer
+ */
+static VALUE
+Coolio_Buffer_prepend(VALUE self, VALUE data)
+{
+    struct buffer *buf;
+    Data_Get_Struct(self, struct buffer, buf);
+
+    data = rb_convert_type(data, T_STRING, "String", "to_str");
+    buffer_prepend(buf, RSTRING_PTR(data), RSTRING_LEN(data));
+
+    return data;
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#read(length = nil) -> String
+ *
+ * Read the specified abount of data from the buffer.  If no value
+ * is given the entire contents of the buffer are returned.  Any data
+ * read from the buffer is cleared.
+ * The given length must be greater than 0 or an exception would raise.
+ * If the buffer size is zero then an empty string is returned (regardless
+ * the given length).
+ */
+static VALUE
+Coolio_Buffer_read(int argc, VALUE * argv, VALUE self)
+{
+    VALUE  length_obj, str;
+    int    length;
+    struct buffer *buf;
+
+    Data_Get_Struct(self, struct buffer, buf);
+
+    if (rb_scan_args(argc, argv, "01", &length_obj) == 1) {
+        length = NUM2INT(length_obj);
+        if(length < 1)
+          rb_raise(rb_eArgError, "length must be greater than zero");
+        if(length > buf->size)
+          length = buf->size;
+    } else
+        length = buf->size;
+
+    if(buf->size == 0)
+        return rb_str_new2("");
+
+    str = rb_str_new(0, length);
+    buffer_read(buf, RSTRING_PTR(str), length);
+
+    return str;
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#read_frame(str, mark) -> boolean
+ *
+ * Read up to and including the given frame marker (expressed a a
+ * Fixnum 0-255) byte, copying into the supplied string object. If the mark is
+ * not encountered before the end of the buffer, false is returned but data
+ * is still copied into str. True is returned if the end of a frame is reached.
+ *
+ */
+static VALUE
+Coolio_Buffer_read_frame(VALUE self, VALUE data, VALUE mark)
+{
+    char mark_c = (char) NUM2INT(mark);
+    struct buffer *buf;
+
+    Data_Get_Struct(self, struct buffer, buf);
+
+    if (buffer_read_frame(buf, data, mark_c)) {
+        return Qtrue;
+    } else {
+        return Qfalse;
+    }
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#to_str -> String
+ *
+ * Convert the Buffer to a String.  The original buffer is unmodified.
+ */
+static VALUE
+Coolio_Buffer_to_str(VALUE self)
+{
+    VALUE          str;
+    struct buffer *buf;
+
+    Data_Get_Struct(self, struct buffer, buf);
+
+    str = rb_str_new(0, buf->size);
+    buffer_copy(buf, RSTRING_PTR(str), buf->size);
+
+    return str;
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#read_from(io) -> Integer
+ *
+ * Perform a nonblocking read of the the given IO object and fill
+ * the buffer with any data received.  The call will read as much
+ * data as it can until the read would block.
+ */
+static VALUE
+Coolio_Buffer_read_from(VALUE self, VALUE io)
+{
+    struct buffer  *buf;
+    int             ret;
+#if defined(HAVE_RB_IO_T) || defined(HAVE_RB_IO_DESCRIPTOR)
+    rb_io_t        *fptr;
+#else
+    OpenFile       *fptr;
+#endif
+
+    Data_Get_Struct(self, struct buffer, buf);
+    io = rb_convert_type(io, T_FILE, "IO", "to_io");
+    GetOpenFile(io, fptr);
+    rb_io_set_nonblock(fptr);
+
+#ifdef HAVE_RB_IO_DESCRIPTOR
+    ret = buffer_read_from(buf, rb_io_descriptor(io));
+#else
+    ret = buffer_read_from(buf, FPTR_TO_FD(fptr));
+#endif
+    return ret == -1 ? Qnil : INT2NUM(ret);
+}
+
+/**
+ *  call-seq:
+ *    Coolio::Buffer#write_to(io) -> Integer
+ *
+ * Perform a nonblocking write of the buffer to the given IO object.
+ * As much data as possible is written until the call would block.
+ * Any data which is written is removed from the buffer.
+ */
+static VALUE
+Coolio_Buffer_write_to(VALUE self, VALUE io)
+{
+    struct buffer  *buf;
+#if defined(HAVE_RB_IO_T) || defined(HAVE_RB_IO_DESCRIPTOR)
+    rb_io_t        *fptr;
+#else
+    OpenFile       *fptr;
+#endif
+
+    Data_Get_Struct(self, struct buffer, buf);
+    io = rb_convert_type(io, T_FILE, "IO", "to_io");
+    GetOpenFile(io, fptr);
+    rb_io_set_nonblock(fptr);
+
+#ifdef HAVE_RB_IO_DESCRIPTOR
+    return INT2NUM(buffer_write_to(buf, rb_io_descriptor(io)));
+#else
+    return INT2NUM(buffer_write_to(buf, FPTR_TO_FD(fptr)));
+#endif
+}
+
+/*
+ * Ruby bindings end here.  Below is the actual implementation of
+ * the underlying byte queue ADT
+ */
+
+/* Create a new buffer */
+static struct buffer *
+buffer_new(void)
+{
+    struct buffer *buf;
+
+    buf = (struct buffer *) xmalloc(sizeof(struct buffer));
+    buf->head = buf->tail = buf->pool_head = buf->pool_tail = 0;
+    buf->size = 0;
+    buf->node_size = default_node_size;
+
+    return buf;
+}
+
+/* Clear all data from a buffer */
+static void
+buffer_clear(struct buffer * buf)
+{
+    /* Move everything into the buffer pool */
+    if (!buf->pool_tail) {
+        buf->pool_head = buf->pool_tail = buf->head;
+    } else {
+        buf->pool_tail->next = buf->head;
+    }
+
+    buf->head = buf->tail = 0;
+    buf->size = 0;
+}
+
+/* Free a buffer */
+static void
+buffer_free(struct buffer * buf)
+{
+    buffer_clear(buf);
+    buffer_free_pool(buf);
+
+    free(buf);
+}
+
+/* Free the memory pool */
+static void
+buffer_free_pool(struct buffer * buf)
+{
+    struct buffer_node *tmp;
+
+    while (buf->pool_head) {
+        tmp = buf->pool_head;
+        buf->pool_head = tmp->next;
+        free(tmp);
+    }
+
+    buf->pool_tail = 0;
+}
+
+/* Create a new buffer_node (or pull one from the memory pool) */
+static struct buffer_node *
+buffer_node_new(struct buffer * buf)
+{
+    struct buffer_node *node;
+
+    /* Pull from the memory pool if available */
+    if (buf->pool_head) {
+        node = buf->pool_head;
+        buf->pool_head = node->next;
+
+        if (node->next)
+            node->next = 0;
+        else
+            buf->pool_tail = 0;
+    } else {
+        node = (struct buffer_node *) xmalloc(sizeof(struct buffer_node) + buf->node_size);
+        node->next = 0;
+    }
+
+    node->start = node->end = 0;
+    return node;
+}
+
+/* Free a buffer node (i.e. return it to the memory pool) */
+static void
+buffer_node_free(struct buffer * buf, struct buffer_node * node)
+{
+    node->next = buf->pool_head;
+    buf->pool_head = node;
+
+    if (!buf->pool_tail) {
+        buf->pool_tail = node;
+    }
+}
+
+/* Prepend data to the front of the buffer */
+static void
+buffer_prepend(struct buffer * buf, char *str, unsigned len)
+{
+    struct buffer_node *node, *tmp;
+    buf->size += len;
+
+    /* If it fits in the beginning of the head */
+    if (buf->head && buf->head->start >= len) {
+        buf->head->start -= len;
+        memcpy(buf->head->data + buf->head->start, str, len);
+    } else {
+        node = buffer_node_new(buf);
+        node->next = buf->head;
+        buf->head = node;
+        if (!buf->tail)
+            buf->tail = node;
+
+        while (len > buf->node_size) {
+            memcpy(node->data, str, buf->node_size);
+            node->end = buf->node_size;
+
+            tmp = buffer_node_new(buf);
+            tmp->next = node->next;
+            node->next = tmp;
+
+            if (buf->tail == node)
+                buf->tail = tmp;
+            node = tmp;
+
+            str += buf->node_size;
+            len -= buf->node_size;
+        }
+
+        if (len > 0) {
+            memcpy(node->data, str, len);
+            node->end = len;
+        }
+    }
+}
+
+/* Append data to the front of the buffer */
+static void
+buffer_append(struct buffer * buf, char *str, unsigned len)
+{
+    unsigned nbytes;
+    buf->size += len;
+
+    /* If it fits in the remaining space in the tail */
+    if (buf->tail && len <= buf->node_size - buf->tail->end) {
+        memcpy(buf->tail->data + buf->tail->end, str, len);
+        buf->tail->end += len;
+        return;
+    }
+    /* Empty list needs initialized */
+    if (!buf->head) {
+        buf->head = buffer_node_new(buf);
+        buf->tail = buf->head;
+    }
+    /* Build links out of the data */
+    while (len > 0) {
+        nbytes = buf->node_size - buf->tail->end;
+        if (len < nbytes)
+            nbytes = len;
+
+        memcpy(buf->tail->data + buf->tail->end, str, nbytes);
+        str += nbytes;
+        len -= nbytes;
+
+        buf->tail->end += nbytes;
+
+        if (len > 0) {
+            buf->tail->next = buffer_node_new(buf);
+            buf->tail = buf->tail->next;
+        }
+    }
+}
+
+/* Read data from the buffer (and clear what we've read) */
+static void
+buffer_read(struct buffer * buf, char *str, unsigned len)
+{
+    unsigned nbytes;
+    struct buffer_node *tmp;
+
+    while (buf->size > 0 && len > 0) {
+        nbytes = buf->head->end - buf->head->start;
+        if (len < nbytes)
+            nbytes = len;
+
+        memcpy(str, buf->head->data + buf->head->start, nbytes);
+        str += nbytes;
+        len -= nbytes;
+
+        buf->head->start += nbytes;
+        buf->size -= nbytes;
+
+        if (buf->head->start == buf->head->end) {
+            tmp = buf->head;
+            buf->head = tmp->next;
+            buffer_node_free(buf, tmp);
+
+            if (!buf->head)
+                buf->tail = 0;
+        }
+    }
+}
+
+/*
+ * Read data from the buffer into str until byte frame_mark or empty. Bytes
+ * are copied into str and removed if a complete frame is read, a true value
+ * is returned
+ */
+static int
+buffer_read_frame(struct buffer * buf, VALUE str, char frame_mark)
+{
+    unsigned nbytes = 0;
+    struct buffer_node *tmp;
+
+    while (buf->size > 0) {
+        struct buffer_node *head = buf->head;
+        char           *loc, *s = head->data + head->start, *e = head->data + head->end;
+        nbytes = e - s;
+
+        loc = memchr(s, frame_mark, nbytes);
+
+        if (loc) {
+            nbytes = loc - s + 1;
+        }
+
+        /* Copy less than everything if we found a frame byte */
+        rb_str_cat(str, s, nbytes);
+
+        /* Fixup the buffer pointers to indicate the bytes were consumed */
+        head->start += nbytes;
+        buf->size -= nbytes;
+
+        if (head->start == head->end) {
+            buf->head = head->next;
+            buffer_node_free(buf, head);
+
+            if (!buf->head)
+                buf->tail = 0;
+        }
+
+        if (loc) {
+            return 1;
+        }
+    }
+
+    return 0;
+}
+
+/* Copy data from the buffer without clearing it */
+static void
+buffer_copy(struct buffer * buf, char *str, unsigned len)
+{
+    unsigned nbytes;
+    struct buffer_node *node;
+
+    node = buf->head;
+    while (node && len > 0) {
+        nbytes = node->end - node->start;
+        if (len < nbytes)
+            nbytes = len;
+
+        memcpy(str, node->data + node->start, nbytes);
+        str += nbytes;
+        len -= nbytes;
+
+        if (node->start + nbytes == node->end)
+            node = node->next;
+    }
+}
+
+/* Write data from the buffer to a file descriptor */
+static int
+buffer_write_to(struct buffer * buf, int fd)
+{
+    int bytes_written, total_bytes_written = 0;
+    struct buffer_node *tmp;
+
+    while (buf->head) {
+        bytes_written = write(fd, buf->head->data + buf->head->start, buf->head->end - buf->head->start);
+
+        /* If the write failed... */
+        if (bytes_written < 0) {
+            if (errno != EAGAIN)
+                rb_sys_fail("write");
+
+            return total_bytes_written;
+        }
+
+        total_bytes_written += bytes_written;
+        buf->size -= bytes_written;
+
+        /* If the write blocked... */
+        if (bytes_written < buf->head->end - buf->head->start) {
+            buf->head->start += bytes_written;
+            return total_bytes_written;
+        }
+        /* Otherwise we wrote the whole buffer */
+        tmp = buf->head;
+        buf->head = tmp->next;
+        buffer_node_free(buf, tmp);
+
+        if (!buf->head)
+            buf->tail = 0;
+    }
+
+    return total_bytes_written;
+}
+
+/* Read data from a file descriptor to a buffer */
+/* Append data to the front of the buffer */
+static int
+buffer_read_from(struct buffer * buf, int fd)
+{
+    int      bytes_read, total_bytes_read = 0;
+    unsigned nbytes;
+
+    /* Empty list needs initialized */
+    if (!buf->head) {
+        buf->head = buffer_node_new(buf);
+        buf->tail = buf->head;
+    }
+
+    do {
+        nbytes = buf->node_size - buf->tail->end;
+        bytes_read = read(fd, buf->tail->data + buf->tail->end, nbytes);
+
+        if (bytes_read == 0) {
+            return -1;
+            //When the file reaches EOF
+        } else if (bytes_read < 0) {
+            if (errno != EAGAIN)
+                rb_sys_fail("read");
+
+            return total_bytes_read;
+        }
+
+        total_bytes_read += bytes_read;
+        buf->tail->end += bytes_read;
+        buf->size += bytes_read;
+
+        if (buf->tail->end == buf->node_size) {
+            buf->tail->next = buffer_node_new(buf);
+            buf->tail = buf->tail->next;
+        }
+    } while (bytes_read == nbytes);
+
+    return total_bytes_read;
+}