object-stream 0.1

data/examples/udp.rb

@@ -0,0 +1,35 @@
+case ARGV[0]
+when "marshal", "yaml", "json", "msgpack"
+else
+  abort "Usage: #$0 marshal|yaml|json|msgpack"
+end
+
+type = ARGV.shift
+
+require 'object-stream'
+require 'socket'
+
+Socket.do_not_reverse_lookup = true
+s = UDPSocket.new; s.bind 'localhost', 0
+t = UDPSocket.new; t.bind 'localhost', 0
+s.connect *t.addr.values_at(2,1)
+t.connect *s.addr.values_at(2,1)
+
+th1 = Thread.new do
+  stream = ObjectStream.new(s, type: type)
+  10.times do |i|
+    stream << [i]
+  end
+  stream << "Bye."
+end
+
+th2 = Thread.new do
+  stream = ObjectStream.new(t, type: type)
+  stream.each do |obj|
+    p obj
+    break if /bye/i === obj
+  end
+end
+
+th1.join
+th2.join

data/lib/object-stream-wrapper.rb

@@ -0,0 +1,114 @@
+require 'object-stream'
+
+# Utility wrapper for basic ObjectStream class. Adds three groups of
+# functionality:
+#
+#  * peer_name
+#
+#  * expect
+#
+#  * consume
+#
+class ObjectStreamWrapper
+  include Enumerable
+
+  # Not set by this library, but available for users to keep track of
+  # the peer in a symbolic, application-specific manner. See funl for
+  # an example.
+  attr_accessor :peer_name
+
+  def initialize *args, **opts
+    @stream = ObjectStream.new(*args, **opts)
+    @peer_name = "unknown"
+    @expected_class = nil
+    @consumers = []
+    unexpect
+  end
+
+  def to_s
+    "#<Wrapped #{@stream.class} to #{peer_name}, io=#{@stream.inspect}>"
+  end
+  
+  # Set the stream state so that subsequent objects returned by read will be
+  # instances of a custom class +cl+. Does not affect #consume.
+  # Class +cl+ should define cl.from_serialized, plus #to_json, #to_msgpack,
+  # etc. as needed by the underlying serialization library.
+  def expect cl
+    @expected_class = cl
+  end
+  
+  # Turn off the custom class instantiation of #expect.
+  def unexpect; expect nil; end
+
+  # The block is appended to a queue of procs that are called for the
+  # subsequently read objects, instead of iterating over or returning them.
+  # Helps with handshake protocols. Not affected by #expect.
+  def consume &bl
+    @consumers << bl
+  end
+
+  def try_consume obj
+    if bl = @consumers.shift
+      bl[obj]
+      true
+    else
+      false
+    end
+  end
+  private :try_consume
+  
+  def convert_to_expected obj
+    if @expected_class and not obj.kind_of? @expected_class
+      @expected_class.from_serialized(obj)
+    else
+      obj
+    end
+  end
+  private :convert_to_expected
+
+  def read
+    if block_given?
+      @stream.read do |obj|
+        try_consume(obj) or yield convert_to_expected(obj)
+      end
+      return nil
+    else
+      begin
+        obj = @stream.read
+      end while try_consume(obj)
+      convert_to_expected(obj)
+    end
+  end
+  
+  def each
+    return to_enum unless block_given?
+    read {|obj| yield obj} until eof
+  rescue EOFError
+  end
+
+  def write *objects
+    @stream.write *objects
+  end
+  alias << write
+
+  def write_to_outbox *args, &bl
+    @stream.write_to_outbox *args, &bl
+  end
+
+  def eof?
+    @stream.eof?
+  end
+  alias eof eof?
+
+  def close
+    @stream.close
+  end
+
+  def closed?
+    @stream.closed?
+  end
+
+  def to_io
+    @stream.to_io
+  end
+end

data/lib/object-stream.rb

@@ -0,0 +1,324 @@
+# Stream of objects, with any underlying IO: File, Pipe, Socket, StringIO.
+# Stream is bidirectional if the IO is bidirectional.
+#
+# Serializes objects using any of several serializers: marshal, yaml, json,
+# msgpack. Works with select/readpartial if the serializer supports it (msgpack
+# and yajl do).
+#
+# ObjectStream supports three styles of iteration: Enumerable, blocking read,
+# and yielding (non-blocking) read.
+module ObjectStream
+  include Enumerable
+  
+  # The IO through which the stream reads and writes serialized object data.
+  attr_reader :io
+  
+  # Number of outgoing objects that can accumulate before the outbox is
+  # serialized to the byte buffer (and possibly to the io).
+  attr_reader :max_outbox
+  
+  MARSHAL_TYPE  = "marshal".freeze
+  YAML_TYPE     = "yaml".freeze
+  JSON_TYPE     = "json".freeze
+  MSGPACK_TYPE  = "msgpack".freeze
+  
+  TYPES = [
+    MARSHAL_TYPE, YAML_TYPE, JSON_TYPE, MSGPACK_TYPE
+  ]
+
+  DEFAULT_MAX_OUTBOX = 10
+
+  # Raised when maxbuf exceeded.
+  class OverflowError < StandardError; end
+
+  @stream_class_map =
+    Hash.new {|h,type| raise ArgumentError, "unknown type: #{type.inspect}"}
+  @mutex = Mutex.new
+
+  class << self
+    def new io, type: MARSHAL_TYPE, **opts
+      if io.kind_of? ObjectStream
+        raise ArgumentError,
+          "given io is already an ObjectStream: #{io.inspect}"
+      end
+      stream_class_for(type).new io, **opts
+    end
+
+    def stream_class_for type
+      cl = @stream_class_map[type]
+      return cl if cl.respond_to? :new
+
+      # Protect against race condition in msgpack and yajl extension
+      # initialization (bug #8374).
+      @mutex.synchronize do
+        return cl if cl.respond_to? :new
+        @stream_class_map[type] = cl.call
+      end
+    end
+    
+    def register_type type, &bl
+      @stream_class_map[type] = bl
+    end
+  end
+  
+  def initialize io, max_outbox: DEFAULT_MAX_OUTBOX, **opts
+    @io = io
+    @max_outbox = max_outbox
+    @inbox = nil
+    @outbox = []
+  end
+  
+  def to_s
+    "#<#{self.class} io=#{io.inspect}>"
+  end
+  
+  # If no block given, behaves just the same as #read_one. If block given,
+  # reads any available data and yields it to the block. This form is non-
+  # blocking, if supported by the underlying serializer (such as msgpack).
+  def read
+    if block_given?
+      read_from_inbox {|obj| yield obj}
+      read_from_stream {|obj| yield obj}
+      return nil
+    else
+      read_one
+    end
+  end
+
+  # Read one object from the stream, blocking if necessary. Returns the object.
+  # Raises EOFError at the end of the stream.
+  def read_one
+    if @inbox and not @inbox.empty?
+      return @inbox.shift
+    end
+    
+    have_result = false
+    result = nil
+    until have_result
+      read do |obj| # might not read enough bytes to yield an obj
+        if have_result
+          (@inbox||=[]) << obj
+        else
+          have_result = true
+          result = obj
+        end
+      end
+    end
+    result
+  end
+
+  def read_from_inbox
+    if @inbox and not @inbox.empty?
+      @inbox.each {|obj| yield obj}
+      @inbox.clear
+    end
+  end
+  private :read_from_inbox
+  
+  # Write the given objects to the stream, first flushing any objects in the
+  # outbox. Flushes the underlying byte buffer afterwards.
+  def write *objects
+    write_to_buffer *objects
+    flush_buffer
+  end
+  alias << write
+
+  # Push the given object into the outbox, to be written later when the outbox
+  # is flushed. If a block is given, it will be called when the outbox is
+  # flushed, and its value will be written instead.
+  def write_to_outbox object=nil, &bl
+    @outbox << (bl || object)
+    flush_outbox if @outbox.size > max_outbox
+    self
+  end
+
+  def flush_outbox
+    @outbox.each do |object|
+      object = object.call if object.kind_of? Proc
+      write_to_stream object
+    end
+    @outbox.clear
+    self
+  end
+
+  def write_to_buffer *objects
+    flush_outbox
+    objects.each do |object|
+      write_to_stream object
+    end
+    self
+  end
+
+  def flush_buffer
+    self
+  end
+
+  # Iterate through the (rest of) the stream of objects. Does not raise
+  # EOFError, but simply returns. All Enumerable and Enumerator methods are
+  # available.
+  def each
+    return to_enum unless block_given?
+    read {|obj| yield obj} until eof
+  rescue EOFError
+  end
+  
+  def eof?
+    (!@inbox || @inbox.empty?) && io.eof?
+  end
+  alias eof eof?
+
+  # Call this if the most recent write was a #write_to_buffer without
+  # a #flush_buffer. If you only use #write, there's no need to close
+  # the stream in any special way.
+  def close
+    flush_outbox
+    io.close
+  end
+  
+  def closed?
+    io.closed?
+  end
+  
+  # Makes it possible to use stream in a select.
+  def to_io
+    io
+  end
+  
+  class MarshalStream
+    include ObjectStream
+    
+    ObjectStream.register_type MARSHAL_TYPE do
+      self
+    end
+    
+    def read_from_stream
+      yield Marshal.load(io)
+    end
+    
+    def write_to_stream object
+      Marshal.dump(object, io)
+      self
+    end
+  end
+  
+  class YamlStream
+    include ObjectStream
+    
+    ObjectStream.register_type YAML_TYPE do
+      require 'yaml'
+      self
+    end
+    
+    def read_from_stream
+      YAML.load_stream(io) do |obj|
+        yield obj
+      end
+    end
+    
+    def write_to_stream object
+      YAML.dump(object, io)
+      self
+    end
+  end
+  
+  class JsonStream
+    include ObjectStream
+    
+    ObjectStream.register_type JSON_TYPE do
+      require 'yajl'
+      require 'yajl/json_gem'
+      self
+    end
+    
+    attr_accessor :chunk_size
+
+    DEFAULT_CHUNK_SIZE = 2000
+
+    def initialize io, chunk_size: DEFAULT_CHUNK_SIZE
+      super
+      @parser = Yajl::Parser.new
+      @encoder = Yajl::Encoder.new
+      @chunk_size = chunk_size
+    end
+
+    # Blocks only if no data available on io.
+    def read_from_stream(&bl)
+      @parser.on_parse_complete = bl
+      @parser << io.readpartial(chunk_size)
+    end
+    
+    def write_to_stream object
+      @encoder.encode object, io
+      self
+    end
+  end
+
+  class MsgpackStream
+    include ObjectStream
+    
+    ObjectStream.register_type MSGPACK_TYPE do
+      require 'msgpack'
+      self
+    end
+    
+    attr_accessor :chunk_size
+    attr_accessor :maxbuf
+
+    DEFAULT_CHUNK_SIZE = 2000
+    DEFAULT_MAXBUF = 4000
+    
+    def initialize io, chunk_size: DEFAULT_CHUNK_SIZE, maxbuf: DEFAULT_MAXBUF
+      super
+      @unpacker = MessagePack::Unpacker.new
+        # don't specify io, so don't have to read all of io in one loop
+      
+      @packer = MessagePack::Packer.new(io)
+      @chunk_size = chunk_size
+      @maxbuf = maxbuf
+    end
+
+    # Blocks only if no data available on io.
+    def read_from_stream
+      fill_buffer(chunk_size)
+      checkbuf if maxbuf
+      read_from_buffer do |obj|
+        yield obj
+      end
+    end
+    
+    def fill_buffer n
+      @unpacker.feed(io.readpartial(n))
+    end
+
+    def read_from_buffer
+      @unpacker.each do |obj|
+        yield obj
+      end
+    end
+
+    def checkbuf
+      if maxbuf and @unpacker.buffer.size > maxbuf
+        raise OverflowError,
+          "Exceeded buffer limit by #{@unpacker.buffer.size - maxbuf} bytes."
+      end
+    end
+    
+    def write_to_stream object
+      @packer.write(object).flush
+      self
+    end
+    
+    def write_to_buffer *objects
+      flush_outbox
+      objects.each do |object|
+        @packer.write(object)
+      end
+      self
+    end
+    
+    def flush_buffer
+      @packer.flush
+      self
+    end
+  end
+end