THRIFT-1189. rb: Ruby deserializer speed improvements

This patch gives the Ruby deserialization stack a healthy performance boost.

Patch: Ilya Maykov

git-svn-id: https://svn.apache.org/repos/asf/thrift/trunk@1140780 13f79535-47bb-0310-9956-ffa450edef68
diff --git a/lib/rb/Rakefile b/lib/rb/Rakefile
index 3315bca..7583c97 100644
--- a/lib/rb/Rakefile
+++ b/lib/rb/Rakefile
@@ -82,7 +82,7 @@
     p.summary = "Ruby bindings for the Apache Thrift RPC system"
     p.url = "http://thrift.apache.org"
     p.include_rakefile = true
-    p.version = "0.7.0-dev"
+    p.version = "0.7.0dev"
     p.rubygems_version = ">= 1.2.0"
   end
 
diff --git a/lib/rb/ext/binary_protocol_accelerated.c b/lib/rb/ext/binary_protocol_accelerated.c
index 3854887..bd1c2da 100644
--- a/lib/rb/ext/binary_protocol_accelerated.c
+++ b/lib/rb/ext/binary_protocol_accelerated.c
@@ -34,6 +34,7 @@
 static int VERSION_MASK;
 static int TYPE_MASK;
 static int BAD_VERSION;
+static ID rbuf_ivar_id;
 
 static void write_byte_direct(VALUE trans, int8_t b) {
   WRITE(trans, (char*)&b, 1);
@@ -226,26 +227,36 @@
 VALUE rb_thrift_binary_proto_read_i16(VALUE self);
 
 static char read_byte_direct(VALUE self) {
-  VALUE buf = READ(self, 1);
-  return RSTRING_PTR(buf)[0];
+  VALUE byte = rb_funcall(GET_TRANSPORT(self), read_byte_method_id, 0);
+  return (char)(FIX2INT(byte));
 }
 
 static int16_t read_i16_direct(VALUE self) {
-  VALUE buf = READ(self, 2);
-  return (int16_t)(((uint8_t)(RSTRING_PTR(buf)[1])) | ((uint16_t)((RSTRING_PTR(buf)[0]) << 8)));
+  VALUE rbuf = rb_ivar_get(self, rbuf_ivar_id);
+  rb_funcall(GET_TRANSPORT(self), read_into_buffer_method_id, 2, rbuf, INT2FIX(2));
+  return (int16_t)(((uint8_t)(RSTRING_PTR(rbuf)[1])) | ((uint16_t)((RSTRING_PTR(rbuf)[0]) << 8)));
 }
 
 static int32_t read_i32_direct(VALUE self) {
-  VALUE buf = READ(self, 4);
-  return ((uint8_t)(RSTRING_PTR(buf)[3])) | 
-    (((uint8_t)(RSTRING_PTR(buf)[2])) << 8) | 
-    (((uint8_t)(RSTRING_PTR(buf)[1])) << 16) | 
-    (((uint8_t)(RSTRING_PTR(buf)[0])) << 24);
+  VALUE rbuf = rb_ivar_get(self, rbuf_ivar_id);
+  rb_funcall(GET_TRANSPORT(self), read_into_buffer_method_id, 2, rbuf, INT2FIX(4));
+  return ((uint8_t)(RSTRING_PTR(rbuf)[3])) |
+    (((uint8_t)(RSTRING_PTR(rbuf)[2])) << 8) |
+    (((uint8_t)(RSTRING_PTR(rbuf)[1])) << 16) |
+    (((uint8_t)(RSTRING_PTR(rbuf)[0])) << 24);
 }
 
 static int64_t read_i64_direct(VALUE self) {
-  uint64_t hi = read_i32_direct(self);
-  uint32_t lo = read_i32_direct(self);
+  VALUE rbuf = rb_ivar_get(self, rbuf_ivar_id);
+  rb_funcall(GET_TRANSPORT(self), read_into_buffer_method_id, 2, rbuf, INT2FIX(8));
+  uint64_t hi = ((uint8_t)(RSTRING_PTR(rbuf)[3])) |
+    (((uint8_t)(RSTRING_PTR(rbuf)[2])) << 8) |
+    (((uint8_t)(RSTRING_PTR(rbuf)[1])) << 16) |
+    (((uint8_t)(RSTRING_PTR(rbuf)[0])) << 24);
+  uint32_t lo = ((uint8_t)(RSTRING_PTR(rbuf)[7])) |
+    (((uint8_t)(RSTRING_PTR(rbuf)[6])) << 8) |
+    (((uint8_t)(RSTRING_PTR(rbuf)[5])) << 16) |
+    (((uint8_t)(RSTRING_PTR(rbuf)[4])) << 24);
   return (hi << 32) | lo;
 }
 
@@ -425,4 +436,6 @@
   rb_define_method(bpa_class, "read_map_end", rb_thift_binary_proto_read_map_end, 0);
   rb_define_method(bpa_class, "read_list_end", rb_thift_binary_proto_read_list_end, 0);
   rb_define_method(bpa_class, "read_set_end", rb_thift_binary_proto_read_set_end, 0);
+
+  rbuf_ivar_id = rb_intern("@rbuf");
 }
diff --git a/lib/rb/ext/compact_protocol.c b/lib/rb/ext/compact_protocol.c
index 6c0123d..a47fe6c 100644
--- a/lib/rb/ext/compact_protocol.c
+++ b/lib/rb/ext/compact_protocol.c
@@ -34,6 +34,7 @@
 static ID last_field_id;
 static ID boolean_field_id;
 static ID bool_value_id;
+static ID rbuf_ivar_id;
 
 static int VERSION;
 static int VERSION_MASK;
@@ -354,8 +355,8 @@
 }
 
 static char read_byte_direct(VALUE self) {
-  VALUE buf = READ(self, 1);
-  return RSTRING_PTR(buf)[0];
+  VALUE byte = rb_funcall(GET_TRANSPORT(self), read_byte_method_id, 0);
+  return (char)(FIX2INT(byte));
 }
 
 static int64_t zig_zag_to_ll(int64_t n) {
@@ -528,15 +529,16 @@
     double f;
     int64_t l;
   } transfer;
-  VALUE bytes = READ(self, 8);
-  uint32_t lo = ((uint8_t)(RSTRING_PTR(bytes)[0]))
-    | (((uint8_t)(RSTRING_PTR(bytes)[1])) << 8)
-    | (((uint8_t)(RSTRING_PTR(bytes)[2])) << 16)
-    | (((uint8_t)(RSTRING_PTR(bytes)[3])) << 24);
-  uint64_t hi = (((uint8_t)(RSTRING_PTR(bytes)[4])))
-    | (((uint8_t)(RSTRING_PTR(bytes)[5])) << 8)
-    | (((uint8_t)(RSTRING_PTR(bytes)[6])) << 16)
-    | (((uint8_t)(RSTRING_PTR(bytes)[7])) << 24);
+  VALUE rbuf = rb_ivar_get(self, rbuf_ivar_id);
+  rb_funcall(GET_TRANSPORT(self), read_into_buffer_method_id, 2, rbuf, INT2FIX(8));
+  uint32_t lo = ((uint8_t)(RSTRING_PTR(rbuf)[0]))
+    | (((uint8_t)(RSTRING_PTR(rbuf)[1])) << 8)
+    | (((uint8_t)(RSTRING_PTR(rbuf)[2])) << 16)
+    | (((uint8_t)(RSTRING_PTR(rbuf)[3])) << 24);
+  uint64_t hi = (((uint8_t)(RSTRING_PTR(rbuf)[4])))
+    | (((uint8_t)(RSTRING_PTR(rbuf)[5])) << 8)
+    | (((uint8_t)(RSTRING_PTR(rbuf)[6])) << 16)
+    | (((uint8_t)(RSTRING_PTR(rbuf)[7])) << 24);
   transfer.l = (hi << 32) | lo;
 
   return rb_float_new(transfer.f);
@@ -559,6 +561,7 @@
   last_field_id = rb_intern("@last_field");
   boolean_field_id = rb_intern("@boolean_field");
   bool_value_id = rb_intern("@bool_value");
+  rbuf_ivar_id = rb_intern("@rbuf");
 }
 
 static void Init_rb_methods() {
diff --git a/lib/rb/ext/constants.h b/lib/rb/ext/constants.h
index 57df544..9ea00d2 100644
--- a/lib/rb/ext/constants.h
+++ b/lib/rb/ext/constants.h
@@ -74,6 +74,7 @@
 extern ID skip_method_id;
 extern ID write_method_id;
 extern ID read_all_method_id;
+extern ID read_into_buffer_method_id;
 extern ID native_qmark_method_id;
 
 extern ID fields_const_id;
diff --git a/lib/rb/ext/memory_buffer.c b/lib/rb/ext/memory_buffer.c
index 74efa2c..23cd9eb 100644
--- a/lib/rb/ext/memory_buffer.c
+++ b/lib/rb/ext/memory_buffer.c
@@ -30,6 +30,11 @@
 
 #define GET_BUF(self) rb_ivar_get(self, buf_ivar_id)
 
+VALUE rb_thrift_memory_buffer_write(VALUE self, VALUE str);
+VALUE rb_thrift_memory_buffer_read(VALUE self, VALUE length_value);
+VALUE rb_thrift_memory_buffer_read_byte(VALUE self);
+VALUE rb_thrift_memory_buffer_read_into_buffer(VALUE self, VALUE buffer_value, VALUE size_value);
+
 VALUE rb_thrift_memory_buffer_write(VALUE self, VALUE str) {
   VALUE buf = GET_BUF(self);
   rb_str_buf_cat(buf, RSTRING_PTR(str), RSTRING_LEN(str));
@@ -62,10 +67,59 @@
   return data;
 }
 
+VALUE rb_thrift_memory_buffer_read_byte(VALUE self) {
+  VALUE index_value = rb_ivar_get(self, index_ivar_id);
+  int index = FIX2INT(index_value);
+
+  VALUE buf = GET_BUF(self);
+  if (index >= RSTRING_LEN(buf)) {
+    rb_raise(rb_eEOFError, "Not enough bytes remain in memory buffer");
+  }
+  char byte = RSTRING_PTR(buf)[index++];
+  rb_ivar_set(self, index_ivar_id, INT2FIX(index));
+
+  if (index >= GARBAGE_BUFFER_SIZE) {
+    rb_ivar_set(self, buf_ivar_id, rb_funcall(buf, slice_method_id, 2, INT2FIX(index), INT2FIX(RSTRING_LEN(buf) - 1)));
+    index = 0;
+  }
+  int result = (int) byte;
+  return INT2FIX(result);
+}
+
+VALUE rb_thrift_memory_buffer_read_into_buffer(VALUE self, VALUE buffer_value, VALUE size_value) {
+  int i = 0;
+  int size = FIX2INT(size_value);
+  int index;
+  VALUE buf = GET_BUF(self);
+
+  while (i < size) {
+    index = FIX2INT(rb_ivar_get(self, index_ivar_id));
+    if (index >= RSTRING_LEN(buf)) {
+      rb_raise(rb_eEOFError, "Not enough bytes remain in memory buffer");
+    }
+    char byte = RSTRING_PTR(buf)[index++];
+    rb_ivar_set(self, index_ivar_id, INT2FIX(index));
+
+    if (index >= GARBAGE_BUFFER_SIZE) {
+      rb_ivar_set(self, buf_ivar_id, rb_funcall(buf, slice_method_id, 2, INT2FIX(index), INT2FIX(RSTRING_LEN(buf) - 1)));
+      index = 0;
+    }
+
+    if (i >= RSTRING_LEN(buffer_value)) {
+      rb_raise(rb_eIndexError, "index %d out of string", i);
+    }
+    ((char*)RSTRING_PTR(buffer_value))[i] = byte;
+    i++;
+  }
+  return INT2FIX(i);
+}
+
 void Init_memory_buffer() {
   VALUE thrift_memory_buffer_class = rb_const_get(thrift_module, rb_intern("MemoryBufferTransport"));
   rb_define_method(thrift_memory_buffer_class, "write", rb_thrift_memory_buffer_write, 1);
   rb_define_method(thrift_memory_buffer_class, "read", rb_thrift_memory_buffer_read, 1);
+  rb_define_method(thrift_memory_buffer_class, "read_byte", rb_thrift_memory_buffer_read_byte, 0);
+  rb_define_method(thrift_memory_buffer_class, "read_into_buffer", rb_thrift_memory_buffer_read_into_buffer, 2);
   
   buf_ivar_id = rb_intern("@buf");
   index_ivar_id = rb_intern("@index");
diff --git a/lib/rb/ext/thrift_native.c b/lib/rb/ext/thrift_native.c
index 09b9fe4..2f6bb1a 100644
--- a/lib/rb/ext/thrift_native.c
+++ b/lib/rb/ext/thrift_native.c
@@ -88,6 +88,7 @@
 ID skip_method_id;
 ID write_method_id;
 ID read_all_method_id;
+ID read_into_buffer_method_id;
 ID native_qmark_method_id;
 
 // constant ids
@@ -170,6 +171,7 @@
   skip_method_id = rb_intern("skip");
   write_method_id = rb_intern("write");
   read_all_method_id = rb_intern("read_all");
+  read_into_buffer_method_id = rb_intern("read_into_buffer");
   native_qmark_method_id = rb_intern("native?");
 
   // constant ids
diff --git a/lib/rb/lib/thrift/protocol/binary_protocol.rb b/lib/rb/lib/thrift/protocol/binary_protocol.rb
index 04d149a..f9adb20 100644
--- a/lib/rb/lib/thrift/protocol/binary_protocol.rb
+++ b/lib/rb/lib/thrift/protocol/binary_protocol.rb
@@ -29,6 +29,11 @@
       super(trans)
       @strict_read = strict_read
       @strict_write = strict_write
+
+      # Pre-allocated read buffer for fixed-size read methods. Needs to be at least 8 bytes long for
+      # read_i64() and read_double().
+      @rbuf = "\0" * 8
+      @rbuf.force_encoding("BINARY") if @rbuf.respond_to?(:force_encoding)
     end
 
     def write_message_begin(name, type, seqid)
@@ -165,8 +170,7 @@
     end
 
     def read_byte
-      dat = trans.read_all(1)
-      val = dat[0].ord
+      val = trans.read_byte
       if (val > 0x7f)
         val = 0 - ((val - 1) ^ 0xff)
       end
@@ -174,8 +178,8 @@
     end
 
     def read_i16
-      dat = trans.read_all(2)
-      val, = dat.unpack('n')
+      trans.read_into_buffer(@rbuf, 2)
+      val, = @rbuf.unpack('n')
       if (val > 0x7fff)
         val = 0 - ((val - 1) ^ 0xffff)
       end
@@ -183,8 +187,8 @@
     end
 
     def read_i32
-      dat = trans.read_all(4)
-      val, = dat.unpack('N')
+      trans.read_into_buffer(@rbuf, 4)
+      val, = @rbuf.unpack('N')
       if (val > 0x7fffffff)
         val = 0 - ((val - 1) ^ 0xffffffff)
       end
@@ -192,8 +196,8 @@
     end
 
     def read_i64
-      dat = trans.read_all(8)
-      hi, lo = dat.unpack('N2')
+      trans.read_into_buffer(@rbuf, 8)
+      hi, lo = @rbuf.unpack('N2')
       if (hi > 0x7fffffff)
         hi ^= 0xffffffff
         lo ^= 0xffffffff
@@ -204,8 +208,8 @@
     end
 
     def read_double
-      dat = trans.read_all(8)
-      val = dat.unpack('G').first
+      trans.read_into_buffer(@rbuf, 8)
+      val = @rbuf.unpack('G').first
       val
     end
 
diff --git a/lib/rb/lib/thrift/protocol/compact_protocol.rb b/lib/rb/lib/thrift/protocol/compact_protocol.rb
index f5108ad..ede82f2 100644
--- a/lib/rb/lib/thrift/protocol/compact_protocol.rb
+++ b/lib/rb/lib/thrift/protocol/compact_protocol.rb
@@ -98,6 +98,10 @@
 
       @last_field = [0]
       @boolean_value = nil
+
+      # Pre-allocated read buffer for read_double().
+      @rbuf = "\0" * 8
+      @rbuf.force_encoding("BINARY") if @rbuf.respond_to?(:force_encoding)
     end
 
     def write_message_begin(name, type, seqid)
@@ -302,8 +306,7 @@
     end
 
     def read_byte
-      dat = trans.read_all(1)
-      val = dat[0].ord
+      val = trans.read_byte
       if (val > 0x7f)
         val = 0 - ((val - 1) ^ 0xff)
       end
@@ -323,8 +326,8 @@
     end
 
     def read_double
-      dat = trans.read_all(8)
-      val = dat.reverse.unpack('G').first
+      trans.read_into_buffer(@rbuf, 8)
+      val = @rbuf.reverse.unpack('G').first
       val
     end
 
diff --git a/lib/rb/lib/thrift/struct.rb b/lib/rb/lib/thrift/struct.rb
index 0525f53..3512463 100644
--- a/lib/rb/lib/thrift/struct.rb
+++ b/lib/rb/lib/thrift/struct.rb
@@ -55,7 +55,7 @@
     end
 
     def fields_with_default_values
-      fields_with_default_values = self.class.instance_variable_get("@fields_with_default_values")
+      fields_with_default_values = self.class.instance_variable_get(:@fields_with_default_values)
       unless fields_with_default_values
         fields_with_default_values = {}
         struct_fields.each do |fid, field_def|
@@ -63,7 +63,7 @@
             fields_with_default_values[field_def[:name]] = field_def[:default]
           end
         end
-        self.class.instance_variable_set("@fields_with_default_values", fields_with_default_values)
+        self.class.instance_variable_set(:@fields_with_default_values, fields_with_default_values)
       end
       fields_with_default_values
     end
@@ -114,9 +114,10 @@
     end
 
     def ==(other)
+      return false if other.nil?
       each_field do |fid, field_info|
         name = field_info[:name]
-        return false unless self.instance_variable_get("@#{name}") == other.instance_variable_get("@#{name}")
+        return false unless other.respond_to?(name) && self.send(name) == other.send(name)
       end
       true
     end
@@ -125,13 +126,15 @@
       self.class == other.class && self == other
     end
 
+    # This implementation of hash() is inspired by Apache's Java HashCodeBuilder class.
     def hash
-      field_values = []
+      total = 17
       each_field do |fid, field_info|
         name = field_info[:name]
-        field_values << self.instance_variable_get("@#{name}")
+        value = self.send(name)
+        total = (total * 37 + value.hash) & 0xffffffff
       end
-      field_values.hash
+      total
     end
 
     def differences(other)
diff --git a/lib/rb/lib/thrift/struct_union.rb b/lib/rb/lib/thrift/struct_union.rb
index 6be7ee7..0ff175f 100644
--- a/lib/rb/lib/thrift/struct_union.rb
+++ b/lib/rb/lib/thrift/struct_union.rb
@@ -21,13 +21,13 @@
 module Thrift
   module Struct_Union
     def name_to_id(name)
-      names_to_ids = self.class.instance_variable_get("@names_to_ids")
+      names_to_ids = self.class.instance_variable_get(:@names_to_ids)
       unless names_to_ids
         names_to_ids = {}
         struct_fields.each do |fid, field_def|
           names_to_ids[field_def[:name]] = fid
         end
-        self.class.instance_variable_set("@names_to_ids", names_to_ids)
+        self.class.instance_variable_set(:@names_to_ids, names_to_ids)
       end
       names_to_ids[name]
     end
diff --git a/lib/rb/lib/thrift/transport/base_transport.rb b/lib/rb/lib/thrift/transport/base_transport.rb
index 08a71da..0a12cea 100644
--- a/lib/rb/lib/thrift/transport/base_transport.rb
+++ b/lib/rb/lib/thrift/transport/base_transport.rb
@@ -34,6 +34,26 @@
     end
   end
 
+  module TransportUtils
+    if RUBY_VERSION >= '1.9'
+      def self.get_string_byte(string, index)
+        string.getbyte(index)
+      end
+
+      def self.set_string_byte(string, index, byte)
+        string.setbyte(index, byte)
+      end
+    else
+      def self.get_string_byte(string, index)
+        string[index]
+      end
+
+      def self.set_string_byte(string, index, byte)
+        string[index] = byte
+      end
+    end
+  end
+
   class BaseTransport
     def open?; end
     
@@ -45,9 +65,26 @@
       raise NotImplementedError
     end
 
+    # Returns an unsigned byte as a Fixnum in the range (0..255).
+    def read_byte
+      buf = read_all(1)
+      return ::Thrift::TransportUtils.get_string_byte(buf, 0)
+    end
+
+    # Reads size bytes and copies them into buffer[0..size].
+    def read_into_buffer(buffer, size)
+      tmp = read_all(size)
+      i = 0
+      tmp.each_byte do |byte|
+        ::Thrift::TransportUtils.set_string_byte(buffer, i, byte)
+        i += 1
+      end
+      i
+    end
+
     def read_all(size)
-      buf = ''
-    
+      return '' if size <= 0
+      buf = read(size)
       while (buf.length < size)
         chunk = read(size - buf.length)
         buf << chunk
diff --git a/lib/rb/lib/thrift/transport/buffered_transport.rb b/lib/rb/lib/thrift/transport/buffered_transport.rb
index 8dead4e..676a4d3 100644
--- a/lib/rb/lib/thrift/transport/buffered_transport.rb
+++ b/lib/rb/lib/thrift/transport/buffered_transport.rb
@@ -55,6 +55,37 @@
       ret
     end
 
+    def read_byte
+      # If the read buffer is exhausted, try to read up to DEFAULT_BUFFER more bytes into it.
+      if @index >= @rbuf.size
+        @rbuf = @transport.read(DEFAULT_BUFFER)
+        @index = 0
+      end
+
+      # The read buffer has some data now, read a single byte. Using get_string_byte() avoids
+      # allocating a temp string of size 1 unnecessarily.
+      @index += 1
+      return ::Thrift::TransportUtils.get_string_byte(@rbuf, @index - 1)
+    end
+
+    def read_into_buffer(buffer, size)
+      i = 0
+      while i < size
+        # If the read buffer is exhausted, try to read up to DEFAULT_BUFFER more bytes into it.
+        if @index >= @rbuf.size
+          @rbuf = @transport.read(DEFAULT_BUFFER)
+          @index = 0
+        end
+
+        # The read buffer has some data now, so copy bytes over to the output buffer.
+        byte = ::Thrift::TransportUtils.get_string_byte(@rbuf, @index)
+        ::Thrift::TransportUtils.set_string_byte(buffer, i, byte)
+        @index += 1
+        i += 1
+      end
+      i
+    end
+
     def write(buf)
       @wbuf << buf
     end
diff --git a/lib/rb/lib/thrift/transport/framed_transport.rb b/lib/rb/lib/thrift/transport/framed_transport.rb
index 558af74..e7630d0 100644
--- a/lib/rb/lib/thrift/transport/framed_transport.rb
+++ b/lib/rb/lib/thrift/transport/framed_transport.rb
@@ -52,6 +52,32 @@
       @rbuf.slice(@index - sz, sz) || ''
     end
 
+    def read_byte
+      return @transport.read_byte() unless @read
+
+      read_frame if @index >= @rbuf.length
+
+      # The read buffer has some data now, read a single byte. Using get_string_byte() avoids
+      # allocating a temp string of size 1 unnecessarily.
+      @index += 1
+      return ::Thrift::TransportUtils.get_string_byte(@rbuf, @index - 1)
+    end
+
+    def read_into_buffer(buffer, size)
+      i = 0
+      while i < size
+        read_frame if @index >= @rbuf.length
+
+        # The read buffer has some data now, so copy bytes over to the output buffer.
+        byte = ::Thrift::TransportUtils.get_string_byte(@rbuf, @index)
+        ::Thrift::TransportUtils.set_string_byte(buffer, i, byte)
+        @index += 1
+        i += 1
+      end
+      i
+    end
+
+
     def write(buf,sz=nil)
       return @transport.write(buf) unless @write
 
diff --git a/lib/rb/lib/thrift/transport/memory_buffer_transport.rb b/lib/rb/lib/thrift/transport/memory_buffer_transport.rb
index 3169b2b..5f740ca 100644
--- a/lib/rb/lib/thrift/transport/memory_buffer_transport.rb
+++ b/lib/rb/lib/thrift/transport/memory_buffer_transport.rb
@@ -70,6 +70,31 @@
       data
     end
 
+    def read_byte
+      raise EOFError.new("Not enough bytes remain in buffer") if @index >= @buf.size
+      val = ::Thrift::TransportUtils.get_string_byte(@buf, @index)
+      @index += 1
+      if @index >= GARBAGE_BUFFER_SIZE
+        @buf = @buf.slice(@index..-1)
+        @index = 0
+      end
+      val
+    end
+
+    def read_into_buffer(buffer, size)
+      i = 0
+      while i < size
+        raise EOFError.new("Not enough bytes remain in buffer") if @index >= @buf.size
+
+        # The read buffer has some data now, so copy bytes over to the output buffer.
+        byte = ::Thrift::TransportUtils.get_string_byte(@buf, @index)
+        ::Thrift::TransportUtils.set_string_byte(buffer, i, byte)
+        @index += 1
+        i += 1
+      end
+      i
+    end
+
     def write(wbuf)
       @buf << wbuf
     end
diff --git a/lib/rb/spec/binary_protocol_spec.rb b/lib/rb/spec/binary_protocol_spec.rb
index 0abccb8..6da467e 100644
--- a/lib/rb/spec/binary_protocol_spec.rb
+++ b/lib/rb/spec/binary_protocol_spec.rb
@@ -31,10 +31,8 @@
     end
 
     it "should read a message header" do
-      @trans.should_receive(:read_all).exactly(2).times.and_return(
-        [protocol_class.const_get(:VERSION_1) | Thrift::MessageTypes::REPLY].pack('N'),
-        [42].pack('N')
-      )      
+      @trans.write([protocol_class.const_get(:VERSION_1) | Thrift::MessageTypes::REPLY].pack('N'))
+      @trans.write([42].pack('N'))
       @prot.should_receive(:read_string).and_return('testMessage')
       @prot.read_message_begin.should == ['testMessage', Thrift::MessageTypes::REPLY, 42]
     end