right_amqp 0.2.0

data/lib/right_amqp/amqp/buffer.rb

@@ -0,0 +1,395 @@
+if [].map.respond_to? :with_index
+  class Array #:nodoc:
+    def enum_with_index
+      each.with_index
+    end
+  end
+else
+  require 'enumerator'
+end
+
+module AMQP
+  class Buffer #:nodoc: all
+    class Overflow < StandardError; end
+    class InvalidType < StandardError; end
+    
+    def initialize data = ''
+      @data = data
+      @pos = 0
+    end
+
+    attr_reader :pos
+    
+    def data
+      @data.clone
+    end
+    alias :contents :data
+    alias :to_s :data
+
+    def << data
+      @data << data.to_s
+      self
+    end
+    
+    def length
+      @data.bytesize
+    end
+    
+    def empty?
+      pos == length
+    end
+    
+    def rewind
+      @pos = 0
+    end
+    
+    def read_properties *types
+      types.shift if types.first == :properties
+      
+      i = 0
+      values = []
+
+      while props = read(:short)
+        (0..14).each do |n|
+          # no more property types
+          break unless types[i]
+          
+          # if flag is set
+          if props & (1<<(15-n)) != 0
+            if types[i] == :bit
+              # bit values exist in flags only
+              values << true
+            else
+              # save type name for later reading
+              values << types[i]
+            end
+          else
+            # property not set or is false bit
+            values << (types[i] == :bit ? false : nil)
+          end
+
+          i+=1
+        end
+
+        # bit(0) == 0 means no more property flags
+        break unless props & 1 == 1
+      end
+
+      values.map do |value|
+        value.is_a?(Symbol) ? read(value) : value
+      end
+    end
+
+    def read *types
+      if types.first == :properties
+        return read_properties(*types)
+      end
+
+      values = types.map do |type|
+        case type
+        when :octet
+          _read(1, 'C')
+        when :short
+          _read(2, 'n')
+        when :long
+          _read(4, 'N')
+        when :longlong
+          upper, lower = _read(8, 'NN')
+          upper << 32 | lower
+        when :shortstr
+          _read read(:octet)
+        when :longstr
+          _read read(:long)
+        when :timestamp
+          Time.at read(:longlong)
+        when :table
+          t = Hash.new
+
+          table = Buffer.new(read(:longstr))
+          until table.empty?
+            key, type = table.read(:shortstr, :octet)
+            key = key.intern
+            t[key] ||= case type
+                       when 83 # 'S'
+                         table.read(:longstr)
+                       when 73 # 'I'
+                         table.read(:long)
+                       when 68 # 'D'
+                         exp = table.read(:octet)
+                         num = table.read(:long)
+                         num / 10.0**exp
+                       when 84 # 'T'
+                         table.read(:timestamp)
+                       when 70 # 'F'
+                         table.read(:table)
+                       end
+          end
+
+          t
+        when :bit
+          if (@bits ||= []).empty?
+            val = read(:octet)
+            @bits = (0..7).map{|i| (val & 1<<i) != 0 }
+          end
+
+          @bits.shift
+        else
+          raise InvalidType, "Cannot read data of type #{type}"
+        end
+      end
+      
+      types.size == 1 ? values.first : values
+    end
+    
+    def write type, data
+      case type
+      when :octet
+        _write(data, 'C')
+      when :short
+        _write(data, 'n')
+      when :long
+        _write(data, 'N')
+      when :longlong
+        lower =  data & 0xffffffff
+        upper = (data & ~0xffffffff) >> 32
+        _write([upper, lower], 'NN')
+      when :shortstr
+        data = (data || '').to_s
+        _write([data.bytesize, data], 'Ca*')
+      when :longstr
+        if data.is_a? Hash
+          write(:table, data)
+        else
+          data = (data || '').to_s
+          _write([data.bytesize, data], 'Na*')
+        end
+      when :timestamp
+        write(:longlong, data.to_i)
+      when :table
+        data ||= {}
+        write :longstr, (data.inject(Buffer.new) do |table, (key, value)|
+                          table.write(:shortstr, key.to_s)
+
+                          case value
+                          when String
+                            table.write(:octet, 83) # 'S'
+                            table.write(:longstr, value.to_s)
+                          when Fixnum
+                            table.write(:octet, 73) # 'I'
+                            table.write(:long, value)
+                          when Float
+                            table.write(:octet, 68) # 'D'
+                            # XXX there's gotta be a better way to do this..
+                            exp = value.to_s.split('.').last.bytesize
+                            num = value * 10**exp
+                            table.write(:octet, exp)
+                            table.write(:long, num)
+                          when Time
+                            table.write(:octet, 84) # 'T'
+                            table.write(:timestamp, value)
+                          when Hash
+                            table.write(:octet, 70) # 'F'
+                            table.write(:table, value)
+                          end
+
+                          table
+                        end)
+      when :bit
+        [*data].to_enum(:each_slice, 8).each{|bits|
+          write(:octet, bits.enum_with_index.inject(0){ |byte, (bit, i)|
+            byte |= 1<<i if bit
+            byte
+           })
+         }
+      when :properties
+        values = []
+        data.enum_with_index.inject(0) do |short, ((type, value), i)|
+          n = i % 15
+          last = i+1 == data.size
+
+          if (n == 0 and i != 0) or last
+            if data.size > i+1
+              short |= 1<<0
+            elsif last and value
+              values << [type,value]
+              short |= 1<<(15-n)
+            end
+
+            write(:short, short)
+            short = 0
+          end
+
+          if value and !last
+            values << [type,value] 
+            short |= 1<<(15-n)
+          end
+
+          short
+        end
+        
+        values.each do |type, value|
+          write(type, value) unless type == :bit
+        end
+      else
+        raise InvalidType, "Cannot write data of type #{type}"
+      end
+      
+      self
+    end
+
+    def extract
+      begin
+        cur_data, cur_pos = @data.clone, @pos
+        yield self
+      rescue Overflow
+        @data, @pos = cur_data, cur_pos
+        nil
+      end
+    end
+
+    def _read size, pack = nil
+      if @pos + size > length
+        raise Overflow
+      else
+        data = @data[@pos,size]
+        @data[@pos,size] = ''
+        if pack
+          data = data.unpack(pack)
+          data = data.pop if data.size == 1
+        end
+        data
+      end
+    end
+    
+    def _write data, pack = nil
+      data = [*data].pack(pack) if pack
+      @data[@pos,0] = data
+      @pos += data.bytesize
+    end
+  end
+end
+
+if $0 =~ /bacon/ or $0 == __FILE__
+  require 'bacon'
+  include AMQP
+
+  describe Buffer do
+    before do
+      @buf = Buffer.new
+    end
+
+    should 'have contents' do
+      @buf.contents.should == ''
+    end
+
+    should 'initialize with data' do
+      @buf = Buffer.new('abc')
+      @buf.contents.should == 'abc'
+    end
+
+    should 'append raw data' do
+      @buf << 'abc'
+      @buf << 'def'
+      @buf.contents.should == 'abcdef'
+    end
+
+    should 'append other buffers' do
+      @buf << Buffer.new('abc')
+      @buf.data.should == 'abc'
+    end
+
+    should 'have a position' do
+      @buf.pos.should == 0
+    end
+
+    should 'have a length' do
+      @buf.length.should == 0
+      @buf << 'abc'
+      @buf.length.should == 3
+    end
+
+    should 'know the end' do
+      @buf.empty?.should == true
+    end
+
+    should 'read and write data' do
+      @buf._write('abc')
+      @buf.rewind
+      @buf._read(2).should == 'ab'
+      @buf._read(1).should == 'c'
+    end
+
+    should 'raise on overflow' do
+      lambda{ @buf._read(1) }.should.raise Buffer::Overflow
+    end
+
+    should 'raise on invalid types' do
+      lambda{ @buf.read(:junk) }.should.raise Buffer::InvalidType
+      lambda{ @buf.write(:junk, 1) }.should.raise Buffer::InvalidType
+    end
+  
+    { :octet => 0b10101010,
+      :short => 100,
+      :long => 100_000_000,
+      :longlong => 666_555_444_333_222_111,
+      :shortstr => 'hello',
+      :longstr => 'bye'*500,
+      :timestamp => time = Time.at(Time.now.to_i),
+      :table => { :this => 'is', :a => 'hash', :with => {:nested => 123, :and => time, :also => 123.456} },
+      :bit => true
+    }.each do |type, value|
+
+      should "read and write a #{type}" do
+        @buf.write(type, value)
+        @buf.rewind
+        @buf.read(type).should == value
+        @buf.should.be.empty
+      end
+
+    end
+    
+    should 'read and write multiple bits' do
+      bits = [true, false, false, true, true, false, false, true, true, false]
+      @buf.write(:bit, bits)
+      @buf.write(:octet, 100)
+      
+      @buf.rewind
+      
+      bits.map do
+        @buf.read(:bit)
+      end.should == bits
+      @buf.read(:octet).should == 100
+    end
+
+    should 'read and write properties' do
+      properties = ([
+        [:octet, 1],
+        [:shortstr, 'abc'],
+        [:bit, true],
+        [:bit, false],
+        [:shortstr, nil],
+        [:timestamp, nil],
+        [:table, { :a => 'hash' }],
+      ]*5).sort_by{rand}
+      
+      @buf.write(:properties, properties)
+      @buf.rewind
+      @buf.read(:properties, *properties.map{|type,_| type }).should == properties.map{|_,value| value }
+      @buf.should.be.empty
+    end
+
+    should 'do transactional reads with #extract' do
+      @buf.write :octet, 8
+      orig = @buf.to_s
+
+      @buf.rewind
+      @buf.extract do |b|
+        b.read :octet
+        b.read :short
+      end
+
+      @buf.pos.should == 0
+      @buf.data.should == orig
+    end
+  end
+end

data/lib/right_amqp/amqp/client.rb

@@ -0,0 +1,282 @@
+require 'right_support'
+require File.expand_path('../frame', __FILE__)
+
+module AMQP
+  class Error < StandardError; end
+
+  module BasicClient
+    def process_frame frame
+      if mq = channels[frame.channel]
+        mq.process_frame(frame)
+        return
+      end
+
+      case frame
+      when Frame::Method
+        case method = frame.payload
+        when Protocol::Connection::Start
+          send Protocol::Connection::StartOk.new({:platform => 'Ruby/EventMachine',
+                                                  :product => 'AMQP',
+                                                  :information => 'http://github.com/tmm1/amqp',
+                                                  :version => VERSION},
+                                                 'AMQPLAIN',
+                                                 {:LOGIN => @settings[:user],
+                                                  :PASSWORD => @settings[:pass]},
+                                                 'en_US')
+
+        when Protocol::Connection::Tune
+          send Protocol::Connection::TuneOk.new(:channel_max => 0,
+                                                :frame_max => 131072,
+                                                :heartbeat => @settings[:heartbeat] || 0)
+
+          send Protocol::Connection::Open.new(:virtual_host => @settings[:vhost],
+                                              :capabilities => '',
+                                              :insist => @settings[:insist])
+
+        when Protocol::Connection::OpenOk
+          logger.debug("[amqp] Received open completion from broker #{@settings[:identity]}")
+          succeed(self)
+
+        when Protocol::Connection::Close
+          # raise Error, "#{method.reply_text} in #{Protocol.classes[method.class_id].methods[method.method_id]}"
+          STDERR.puts "#{method.reply_text} in #{Protocol.classes[method.class_id].methods[method.method_id]}"
+
+        when Protocol::Connection::CloseOk
+          logger.debug("[amqp] Received close completion from broker #{@settings[:identity]}")
+          @on_disconnect.call if @on_disconnect
+        end
+
+      when Frame::Heartbeat
+        logger.debug("[amqp] Received heartbeat from broker #{@settings[:identity]}")
+        @last_server_heartbeat = Time.now
+
+      end
+
+      # Make callback now that handshake with the broker has completed
+      # The 'connected' status callback happens before the handshake is done and if it results in
+      # a lot of activity it might prevent EM from being able to call the code handling the
+      # incoming handshake packet in a timely fashion causing the broker to close the connection
+      @connection_status.call(:ready) if @connection_status && frame.payload.is_a?(AMQP::Protocol::Connection::Start)
+    end
+  end
+
+  def self.client
+    @client ||= BasicClient
+  end
+
+  def self.client= mod
+    mod.__send__ :include, AMQP
+    @client = mod
+  end
+
+  module Client
+    include EM::Deferrable
+    include RightSupport::Log::Mixin
+
+    def self.included(base)
+      base.extend(RightSupport::Log::Mixin::ClassMethods)
+    end
+
+    def initialize opts = {}
+      @settings = opts
+      extend AMQP.client
+
+      @on_disconnect ||= proc{ @connection_status.call(:failed) if @connection_status }
+
+      timeout @settings[:timeout] if @settings[:timeout]
+      errback{ @on_disconnect.call } unless @reconnecting
+
+      @connected = false
+    end
+
+    def connection_completed
+      start_tls if @settings[:ssl]
+      log 'connected'
+      # @on_disconnect = proc{ raise Error, 'Disconnected from server' }
+      unless @closing
+        @on_disconnect = method(:disconnected)
+        @reconnecting = false
+      end
+
+      @connected = true
+      @connection_status.call(:connected) if @connection_status
+
+      @buf = Buffer.new
+      send_data HEADER
+      send_data [1, 1, VERSION_MAJOR, VERSION_MINOR].pack('C4')
+
+      if heartbeat = @settings[:heartbeat]
+        init_heartbeat if (@settings[:heartbeat] = heartbeat.to_i) > 0
+      end
+    end
+
+    def init_heartbeat
+      logger.debug("[amqp] Initializing heartbeat for broker #{@settings[:identity]} to #{@settings[:heartbeat]}")
+      @last_server_heartbeat = Time.now
+
+      @timer.cancel if @timer
+      @timer = EM::PeriodicTimer.new(@settings[:heartbeat]) do
+        if connected?
+          if @last_server_heartbeat < (Time.now - (@settings[:heartbeat] * 2))
+            log "Reconnecting due to missing server heartbeats"
+            logger.info("[amqp] Reconnecting to broker #{@settings[:identity]} due to missing server heartbeats")
+            reconnect(true)
+          else
+            logger.debug("[amqp] Sending heartbeat to broker #{@settings[:identity]}")
+            @last_server_heartbeat = Time.now
+            send AMQP::Frame::Heartbeat.new, :channel => 0
+          end
+        end
+      end
+    end
+
+    def connected?
+      @connected
+    end
+
+    def unbind
+      log 'disconnected'
+      @connected = false
+      EM.next_tick{ @on_disconnect.call }
+    end
+
+    def add_channel mq
+      (@_channel_mutex ||= Mutex.new).synchronize do
+        channels[ key = (channels.keys.max || 0) + 1 ] = mq
+        key
+      end
+    end
+
+    def channels
+      @channels ||= {}
+    end
+
+    # Catch exceptions that would otherwise cause EM to stop or be in a bad
+    # state if a top level EM error handler was setup. Instead close the connection and leave EM
+    # alone.
+    # Don't log an error if the environment variable IGNORE_AMQP_FAILURES is set (used in the
+    # enroll script)
+    def receive_data data
+      begin
+        # log 'receive_data', data
+        @buf << data
+
+        while frame = Frame.parse(@buf)
+          log 'receive', frame
+          process_frame frame
+        end
+      rescue Exception => e
+        logger.exception("[amqp] Failed processing frame, closing connection", e, :trace) unless ENV['IGNORE_AMQP_FAILURES']
+        failed
+      end
+    end
+
+    def process_frame frame
+      # this is a stub meant to be
+      # replaced by the module passed into initialize
+    end
+
+    def send data, opts = {}
+      channel = opts[:channel] ||= 0
+      data = data.to_frame(channel) unless data.is_a? Frame
+      data.channel = channel
+
+      log 'send', data
+      send_data data.to_s
+    end
+
+    #:stopdoc:
+    # def send_data data
+    #   log 'send_data', data
+    #   super
+    # end
+    #:startdoc:
+
+    def close &on_disconnect
+      if on_disconnect
+        @closing = true
+        @on_disconnect = proc{
+          on_disconnect.call
+          @closing = false
+        }
+      end
+
+      callback{ |c|
+        if c.channels.any?
+          c.channels.each do |ch, mq|
+            mq.close
+          end
+        else
+          send Protocol::Connection::Close.new(:reply_code => 200,
+                                               :reply_text => 'Goodbye',
+                                               :class_id => 0,
+                                               :method_id => 0)
+        end
+      }
+    end
+
+    def reconnect force = false
+      if @reconnecting and not force
+        # Wait after first reconnect attempt and in between each subsequent attempt
+        EM.add_timer(@settings[:reconnect_interval] || 5) { reconnect(true) }
+        return
+      end
+
+      unless @reconnecting
+        @deferred_status = nil
+        initialize(@settings)
+
+        mqs = @channels
+        @channels = {}
+        mqs.each{ |_,mq| mq.reset } if mqs
+
+        @reconnecting = true
+
+        again = @settings[:reconnect_delay]
+        again = again.call if again.is_a?(Proc)
+        if again.is_a?(Numeric)
+          # Wait before making initial reconnect attempt
+          EM.add_timer(again) { reconnect(true) }
+          return
+        elsif ![nil, true].include?(again)
+          raise ::AMQP::Error, "Could not interpret :reconnect_delay => #{again.inspect}; expected nil, true, or Numeric"
+        end
+      end
+
+      log 'reconnecting'
+      logger.info("[amqp] Attempting to reconnect to #{@settings[:identity]}")
+      EM.reconnect(@settings[:host], @settings[:port], self)
+    end
+
+    def self.connect opts = {}
+      opts = AMQP.settings.merge(opts)
+      EM.connect opts[:host], opts[:port], self, opts
+    end
+
+    def connection_status &blk
+      @connection_status = blk
+    end
+
+    def failed
+      @connection_status.call(:failed) if @connection_status
+      @failed = true
+      close_connection
+    end
+
+    private
+
+    def disconnected
+      unless @failed
+        @connection_status.call(:disconnected) if @connection_status
+        reconnect
+      end
+    end
+
+    def log *args
+      return unless @settings[:logging] or AMQP.logging
+      require 'pp'
+      pp args
+      puts
+    end
+  end
+end