famoseagle-carrot 0.1.0

data/lib/amqp/protocol.rb

@@ -0,0 +1,209 @@
+module AMQP
+  module Protocol
+    #:stopdoc:
+    class Class::Method
+      def initialize *args
+        opts = args.pop if args.last.is_a? Hash
+        opts ||= {}
+        
+        @debug = 1 # XXX hack, p(obj) == '' if no instance vars are set
+        
+        if args.size == 1 and args.first.is_a? Buffer
+          buf = args.shift
+        else
+          buf = nil
+        end
+
+        self.class.arguments.each do |type, name|
+          val = buf ? buf.read(type) :
+                      args.shift || opts[name] || opts[name.to_s]
+          instance_variable_set("@#{name}", val)
+        end
+      end
+
+      def arguments
+        self.class.arguments.inject({}) do |hash, (type, name)|
+          hash.update name => instance_variable_get("@#{name}")
+        end
+      end
+
+      def to_binary
+        buf = Buffer.new
+        buf.write :short, self.class.parent.id
+        buf.write :short, self.class.id
+
+        bits = []
+
+        self.class.arguments.each do |type, name|
+          val = instance_variable_get("@#{name}")
+          if type == :bit
+            bits << (val || false)
+          else
+            unless bits.empty?
+              buf.write :bit, bits
+              bits = []
+            end
+            buf.write type, val
+          end
+        end
+
+        buf.write :bit, bits unless bits.empty?
+        buf.rewind
+
+        buf
+      end
+      
+      def to_s
+        to_binary.to_s
+      end
+      
+      def to_frame channel = 0
+        Frame::Method.new(self, channel)
+      end
+    end
+
+    #:startdoc:
+    #
+    # Contains a properties hash that holds some potentially interesting 
+    # information.
+    # * :delivery_mode
+    # 1 equals transient.
+    # 2 equals persistent. Unconsumed persistent messages will survive
+    # a server restart when they are stored in a durable queue.
+    # * :redelivered
+    # True or False
+    # * :routing_key
+    # The routing string used for matching this message to this queue.
+    # * :priority
+    # An integer in the range of 0 to 9 inclusive.
+    # * :content_type
+    # Always "application/octet-stream" (byte stream)
+    # * :exchange
+    # The source exchange which published this message.
+    # * :message_count
+    # The number of unconsumed messages contained in the queue.
+    # * :delivery_tag
+    # A monotonically increasing integer. This number should not be trusted
+    # as a sequence number. There is no guarantee it won't get reset.
+    class Header
+      def initialize *args
+        opts = args.pop if args.last.is_a? Hash
+        opts ||= {}
+        
+        first = args.shift
+        
+        if first.is_a? ::Class and first.ancestors.include? Protocol::Class
+          @klass = first
+          @size = args.shift || 0
+          @weight = args.shift || 0
+          @properties = opts
+
+        elsif first.is_a? Buffer or first.is_a? String
+          buf = first
+          buf = Buffer.new(buf) unless buf.is_a? Buffer
+          
+          @klass = Protocol.classes[buf.read(:short)]
+          @weight = buf.read(:short)
+          @size = buf.read(:longlong)
+
+          props = buf.read(:properties, *klass.properties.map{|type,_| type })
+          @properties = Hash[*klass.properties.map{|_,name| name }.zip(props).reject{|k,v| v.nil? }.flatten]
+
+        else
+          raise ArgumentError, 'Invalid argument'
+        end
+        
+      end
+      attr_accessor :klass, :size, :weight, :properties
+      
+      def to_binary
+        buf = Buffer.new
+        buf.write :short, klass.id
+        buf.write :short, weight # XXX rabbitmq only supports weight == 0
+        buf.write :longlong, size
+        buf.write :properties, (klass.properties.map do |type, name|
+                                 [ type, properties[name] || properties[name.to_s] ]
+                               end)
+        buf.rewind
+        buf
+      end
+      
+      def to_s
+        to_binary.to_s
+      end
+
+      def to_frame channel = 0
+        Frame::Header.new(self, channel)
+      end
+
+      def == header
+        [ :klass, :size, :weight, :properties ].inject(true) do |eql, field|
+          eql and __send__(field) == header.__send__(field)
+        end
+      end
+
+      def method_missing meth, *args, &blk
+        @properties.has_key?(meth) || @klass.properties.find{|_,name| name == meth } ? @properties[meth] :
+                                                                                       super
+      end
+    end
+
+    def self.parse buf
+      buf = Buffer.new(buf) unless buf.is_a? Buffer
+      class_id, method_id = buf.read(:short, :short)
+      classes[class_id].methods[method_id].new(buf)
+    end
+    #:stopdoc:
+  end
+end
+
+if $0 =~ /bacon/ or $0 == __FILE__
+  require 'bacon'
+  include AMQP
+
+  describe Protocol do
+    should 'instantiate methods with arguments' do
+      meth = Protocol::Connection::StartOk.new nil, 'PLAIN', nil, 'en_US'
+      meth.locale.should == 'en_US'
+    end
+
+    should 'instantiate methods with named parameters' do
+      meth = Protocol::Connection::StartOk.new :locale => 'en_US',
+                                               :mechanism => 'PLAIN'
+      meth.locale.should == 'en_US'
+    end
+
+    should 'convert methods to binary' do
+      meth = Protocol::Connection::Secure.new :challenge => 'secret'
+      meth.to_binary.should.be.kind_of? Buffer
+
+      meth.to_s.should == [ 10, 20, 6, 'secret' ].pack('nnNa*')
+    end
+
+    should 'convert binary to method' do
+      orig = Protocol::Connection::Secure.new :challenge => 'secret'
+      copy = Protocol.parse orig.to_binary
+      orig.should == copy
+    end
+
+    should 'convert headers to binary' do
+      head = Protocol::Header.new Protocol::Basic,
+                                  size = 5,
+                                  weight = 0,
+                                  :content_type => 'text/json',
+                                  :delivery_mode => 1,
+                                  :priority => 1
+      head.to_s.should == [ 60, weight, 0, size, 0b1001_1000_0000_0000, 9, 'text/json', 1, 1 ].pack('nnNNnCa*CC')
+    end
+
+    should 'convert binary to header' do
+      orig = Protocol::Header.new Protocol::Basic,
+                                  size = 5,
+                                  weight = 0,
+                                  :content_type => 'text/json',
+                                  :delivery_mode => 1,
+                                  :priority => 1
+      Protocol::Header.new(orig.to_binary).should == orig
+    end
+  end
+end

data/lib/amqp/queue.rb

@@ -0,0 +1,70 @@
+module AMQP
+  class Queue
+    attr_reader :name, :server
+    attr_accessor :delivery_tag
+
+    def initialize(server, name, opts = {})
+      @server = server
+      @opts   = opts
+      @name   = name
+      server.send_frame(
+        Protocol::Queue::Declare.new({ :queue => name, :nowait => true }.merge(opts))
+      )
+    end
+
+    def delete(opts = {})
+      server.send_frame(
+        Protocol::Queue::Delete.new({ :queue => name, :nowait => true }.merge(opts))
+      )
+      pp server.next_method
+    end
+
+    def pop(opts = {})
+      self.delivery_tag = nil
+      server.send_frame(
+        Protocol::Basic::Get.new({ :queue => name, :consumer_tag => name, :no_ack => !opts.delete(:ack), :nowait => true }.merge(opts))
+      )
+      method = server.next_method
+      return if method.is_a?(Protocol::Basic::GetEmpty)
+
+      self.delivery_tag = method.delivery_tag
+
+      header = server.next_payload
+      msg    = server.next_payload
+      raise 'unexpected length' if msg.length < header.size
+
+      msg
+    end
+
+    def ack
+      server.send_frame(
+        Protocol::Basic::Ack.new(:delivery_tag => delivery_tag)
+      )
+    end
+
+    def publish(data, opts = {})
+      exchange.publish(data, opts)
+    end
+
+    def message_count
+      status.first
+    end
+
+    def consumer_count
+      status.last
+    end
+    
+    def status(opts = {}, &blk)
+      server.send_frame(
+        Protocol::Queue::Declare.new({ :queue => name, :passive => true }.merge(opts))
+      )
+      method = server.next_method
+      [method.message_count, method.consumer_count]
+    end
+
+  private
+    def exchange
+      @exchange ||= Exchange.new(server, :direct, '', :key => name)
+    end
+  end
+end

data/lib/amqp/server.rb

@@ -0,0 +1,213 @@
+require 'socket'
+require 'thread'
+require 'timeout'
+
+module AMQP
+  class Server
+    CONNECT_TIMEOUT = 1.0
+    RETRY_DELAY     = 10.0
+    DEFAULT_PORT    = 5672
+
+    attr_reader   :host, :port, :status
+    attr_accessor :retry_at, :channel, :ticket
+
+    class ConnectionError < StandardError; end
+    class ServerError     < StandardError; end
+    class ClientError     < StandardError; end
+    class ServerDown      < StandardError; end
+    class ProtocolError   < StandardError; end
+
+    def initialize(opts = {})
+      @host   = opts[:host]  || 'localhost'
+      @port   = opts[:port]  || DEFAULT_PORT
+      @user   = opts[:user]  || 'guest'
+      @pass   = opts[:pass]  || 'guest'
+      @vhost  = opts[:vhost] || '/'
+      @insist = opts[:insist]
+      @status = 'NOT CONNECTED'
+
+      @multithread = opts[:multithread]      
+      start_session
+    end
+
+    def start_session
+      @channel = 0
+      write(HEADER)
+      write([1, 1, VERSION_MAJOR, VERSION_MINOR].pack('C4'))
+      raise ProtocolError, 'bad start connection' unless next_method.is_a?(Protocol::Connection::Start)
+
+      send_frame(
+        Protocol::Connection::StartOk.new(
+          {:platform => 'Ruby', :product => 'Carrot', :information => 'http://github.com/famosagle/carrot', :version => VERSION},
+          'AMQPLAIN',
+          {:LOGIN => @user, :PASSWORD => @pass},
+          'en_US'
+        )
+      )
+
+      if next_method.is_a?(Protocol::Connection::Tune)
+        send_frame(
+          Protocol::Connection::TuneOk.new( :channel_max => 0, :frame_max => 131072, :heartbeat => 0)
+        )
+      end
+
+      send_frame(
+        Protocol::Connection::Open.new(:virtual_host => @vhost, :capabilities => '', :insist => @insist)
+      )
+      raise ProtocolError, 'bad open connection' unless next_method.is_a?(Protocol::Connection::OpenOk)
+
+      @channel = 1
+      send_frame(Protocol::Channel::Open.new)
+      raise ProtocolError, "cannot open channel #{channel}" unless next_method.is_a?(Protocol::Channel::OpenOk)
+
+      send_frame(
+        Protocol::Access::Request.new(:realm => '/data', :read => true, :write => true, :active => true, :passive => true)
+      )
+      method = next_method
+      raise ProtocolError, 'access denied' unless method.is_a?(Protocol::Access::RequestOk)
+      self.ticket = method.ticket
+    end
+
+    def send_frame(*args)
+      args.each do |data|
+        data.ticket  = ticket if ticket and data.respond_to?(:ticket=)
+        data         = data.to_frame(channel) unless data.is_a?(Frame)
+        data.channel = channel
+
+        log :send, data
+        write(data.to_s)
+      end
+    end
+
+    def next_frame
+      frame = Frame.get(self)
+      log :received, frame
+      frame
+    end
+
+    def next_method
+      next_payload
+    end
+
+    def next_payload
+      next_frame.payload
+    end
+
+    def close
+      send_frame(
+        Protocol::Channel::Close.new(:reply_code => 200, :reply_text => 'bye', :method_id => 0, :class_id => 0)
+      )
+      puts "Error closing channel #{channel}" unless next_method.is_a?(Protocol::Channel::CloseOk)
+
+      self.channel = 0
+      send_frame(
+        Protocol::Connection::Close.new(:reply_code => 200, :reply_text => 'Goodbye', :class_id => 0, :method_id => 0)
+      )
+      puts "Error closing connection" unless next_method.is_a?(Protocol::Connection::CloseOk)
+
+      close_socket
+    end
+
+    def read(*args)
+      with_socket_management do |socket|
+        socket.read(*args)
+      end
+    end
+
+    def write(*args)
+      with_socket_management do |socket|
+        socket.write(*args)
+      end
+    end
+
+  private
+
+    def with_socket_management(&block)
+      retried = false
+      begin
+        mutex.lock if multithread?
+        yield socket
+
+      rescue ClientError, ServerError, SocketError, SystemCallError, IOError => error
+        if not retried
+          # Close the socket and retry once.
+          close_socket
+          #start_session
+          retried = true
+          retry
+        else
+          # Mark the server dead and raise an error.
+          close(error.message)
+
+          # Reraise as a ConnectionError
+          new_error = ConnectionError.new("#{error.class}: #{error.message}")
+          new_error.set_backtrace(error.backtrace)
+          raise new_error
+        end
+      ensure
+        mutex.unlock if multithread?
+      end
+    end
+
+    def socket
+      return @socket if @socket and not @socket.closed?
+      raise ServerDown, "will retry at #{retry_at}" unless retry?
+
+      begin
+        # Attempt to connect.
+        mutex.lock if multithread?
+        @socket = timeout(CONNECT_TIMEOUT) do
+          TCPSocket.new(host, port)
+        end
+
+        if Socket.constants.include? 'TCP_NODELAY'
+          @socket.setsockopt Socket::IPPROTO_TCP, Socket::TCP_NODELAY, 1
+        end
+        @retry_at = nil
+        @status   = 'CONNECTED'
+      rescue SocketError, SystemCallError, IOError, Timeout::Error => e
+        close_socket
+        raise ServerDown, e.message
+      ensure
+        mutex.unlock if multithread?
+      end
+
+      @socket
+    end
+
+    def multithread?
+      @multithread
+    end
+
+    def retry?
+      @retry_at.nil? or @retry_at < Time.now
+    end
+
+    def unexpected_eof!
+      raise ConnectionError, 'unexpected end of file' 
+    end
+
+    def close_socket(reason=nil)
+      # Close the socket. The server is not considered dead.
+      mutex.lock if multithread?
+      @socket.close if @socket and not @socket.closed?
+      @socket   = nil
+      @retry_at = nil
+      @status   = "NOT CONNECTED"
+    ensure
+      mutex.unlock if multithread?
+    end
+
+    def mutex
+      @mutex ||= Mutex.new
+    end
+
+    def log(*args)
+      return unless Carrot.logging?
+      require 'pp'
+      pp args
+      puts
+    end
+
+  end
+end