klomp 0.0.8 → 1.0.0

data/lib/klomp.rb

@@ -1,5 +1,46 @@
-require 'klomp/client'
+class Klomp
+  VERSION = '1.0.0'
 
-module Klomp
-  VERSION = '0.0.8'
+  class Error < StandardError; end
+
+  attr_reader :connections
+
+  def initialize(servers, options = {})
+    servers = [servers].flatten
+    raise ArgumentError, "no servers given" if servers.empty?
+    @connections = servers.map {|s| Connection.new(s, options) }
+  end
+
+  def publish(queue, body, headers = {})
+    connections_remaining = connections.dup
+    begin
+      conn = connections_remaining.sample
+      conn.publish(queue, body, headers)
+    rescue
+      connections_remaining.delete conn
+      retry unless connections_remaining.empty?
+      raise
+    end
+  end
+
+  def subscribe(queue, subscriber = nil, &block)
+    connections.each {|conn| conn.subscribe(queue, subscriber, &block) }
+  end
+
+  def unsubscribe(queue)
+    connections.each {|conn| conn.unsubscribe(queue) rescue nil }
+  end
+
+  def connected?
+    connections.detect(&:connected?)
+  end
+
+  def disconnect
+    connections.each {|conn| conn.disconnect }
+    @connections = []
+  end
 end
+
+require 'klomp/connection'
+require 'klomp/sentinel'
+require 'klomp/frames'

data/lib/klomp/connection.rb

@@ -0,0 +1,163 @@
+require 'socket'
+require 'uri'
+
+class Klomp
+  FRAME_SEP = "\x00"          # null character is frame separator
+  class Connection
+
+    attr_reader :options, :subscriptions, :logger
+
+    def initialize(server, options={})
+      @options = options
+
+      if server =~ /^stomp:\/\//
+        uri                  = URI.parse server
+        host, port           = uri.host, uri.port
+        @options['login']    = uri.user if uri.user
+        @options['passcode'] = uri.password if uri.password
+        if uri.query && !uri.query.empty?
+          uri.query.split('&').each {|pair| k, v = pair.split('=', 2); @options[k] = v }
+        end
+      else
+        address            = server.split ':'
+        port, host         = address.pop.to_i, address.pop
+        @options['host'] ||= address.pop unless address.empty?
+      end
+
+      @options['server']   = [host, port]
+      @options['host']   ||= host
+      @subscriptions = {}
+      @logger = options['logger']
+      connect
+    end
+
+    def publish(queue, body, headers={})
+      write Frames::Send.new(queue, body, headers)
+    end
+
+    def subscribe(queue, subscriber = nil, &block)
+      raise Klomp::Error, "no subscriber provided" unless subscriber || block
+      raise Klomp::Error, "subscriber does not respond to #call" if subscriber && !subscriber.respond_to?(:call)
+      previous = subscriptions[queue]
+      subscriptions[queue] = subscriber || block
+      write Frames::Subscribe.new(queue) unless previous
+      start_subscriber_thread
+      previous
+    end
+
+    def unsubscribe(queue)
+      write Frames::Unsubscribe.new(queue) if subscriptions.delete queue
+    end
+
+    def connected?()    @socket end
+    def closed?()       @closing && @socket.nil? end
+
+    def disconnect
+      close!
+      stop_subscriber_thread
+      write Frames::Disconnect.new rescue nil
+      @socket.close rescue nil
+      @socket = nil
+    end
+
+    def reconnect
+      return if connected?
+      logger.warn "reconnect server=#{options['server'].join(':')}" if logger
+      connect
+      subs = subscriptions.dup
+      subscriptions.clear
+      subs.each {|queue, subscriber| subscribe(queue, subscriber) }
+    end
+
+    private
+    def connect
+      @socket  = TCPSocket.new *options['server']
+      @socket.set_encoding 'UTF-8'
+      write Frames::Connect.new(options)
+      frame = read Frames::Connected, 0.1
+      log_frame frame if logger
+      raise Error, frame.headers['message'] if frame.error?
+    end
+
+    def write(frame)
+      raise Error, "connection closed" if closed?
+      raise Error, "disconnected"      unless connected?
+
+      rs, ws, = IO.select(nil, [@socket], nil, 0.1)
+      raise Error, "connection unavailable for write" unless ws && !ws.empty?
+
+      @socket.write frame.to_s
+      log_frame frame if logger
+    rescue Error
+      raise
+    rescue
+      go_offline
+      raise
+    end
+
+    def read(type, timeout = nil)
+      rs, = IO.select([@socket], nil, nil, timeout)
+      raise Error, "connection unavailable for read" unless rs && !rs.empty?
+      type.new @socket.gets(FRAME_SEP)
+    rescue Error
+      raise
+    rescue
+      go_offline
+      raise
+    end
+
+    def log_frame(frame)
+      return unless logger.debug?
+      body = frame.body
+      body = body.lines.first.chomp + '...' if body =~ /\n/
+      logger.debug "frame=#{frame.name} #{frame.headers.map{|k,v| k + '=' + v }.join(' ')} body=#{body}"
+    end
+
+    def log_exception(ex, level = :error)
+      logger.send level, "exception=#{ex.class.name} message=#{ex.message.inspect} backtrace[0]=#{ex.backtrace[0]} backtrace[1]=#{ex.backtrace[1]}"
+      logger.debug "exception=#{ex.class.name} full_backtrace=" + ex.backtrace.join("\n")
+    end
+
+    def close!
+      @closing = true
+    end
+
+    def go_offline
+      if logger
+        msg = "offline server=#{options['server'].join(':')}"
+        msg << " exception=#{$!.class.name} message=#{$!.message.inspect}" if $!
+        logger.warn msg
+      end
+      @socket.close rescue nil
+      @socket = nil
+      Sentinel.new(self)
+      stop_subscriber_thread
+    end
+
+    INTERRUPT = Class.new(Error)
+
+    def start_subscriber_thread
+      @subscriber_thread ||= Thread.new do
+        loop do
+          begin
+            message = read Frames::Message
+            raise Error, message.headers['message'] if message.error?
+            if subscriber = subscriptions[message.headers['destination']]
+              subscriber.call message
+            end
+          rescue INTERRUPT
+            break
+          rescue => e
+            log_exception(e, :warn) if logger
+          end
+          break if @closing
+        end
+      end
+    end
+
+    def stop_subscriber_thread
+      thread, @subscriber_thread = @subscriber_thread, nil
+      thread.raise INTERRUPT, "disconnect" if thread
+    end
+  end
+end

data/lib/klomp/frames.rb

@@ -0,0 +1,108 @@
+class Klomp
+  class FrameError < Error; end
+
+  module Frames
+    class Frame
+      def name
+        @name ||= self.class.name.split('::').last.upcase
+      end
+
+      def headers
+        @headers ||= {}
+      end
+
+      def body
+        @body ||= ""
+      end
+
+      def to_s
+        "#{name}\n#{dump_headers}\n#{@body}#{FRAME_SEP}"
+      end
+
+      def dump_headers
+        headers.map do |pair|
+          pair.map {|x| x.gsub("\n","\\n").gsub(":","\\c").gsub("\\", "\\\\") }.join(':')
+        end.join("\n").tap {|s| s << "\n" unless s.empty? }
+      end
+    end
+
+    class ServerFrame < Frame
+      def initialize(data)
+        @headers, @body = parse(data)
+      end
+
+      def error?
+        @error
+      end
+
+      private
+      def parse(data)
+        headers, body = data.split("\n\n")
+        [parse_headers(headers), body.chomp(FRAME_SEP)]
+      end
+
+      def parse_headers(data)
+        frame = nil
+        {}.tap do |headers|
+          data.lines.each do |line|
+            next if line == "\n"
+            unless frame
+              frame = line.chomp
+              @error = frame == "ERROR"
+              if !@error && frame != name
+                raise Klomp::FrameError,
+                  "unexpected frame #{frame} (expected #{name}):\n#{data}"
+              end
+              next
+            end
+            kv = line.chomp.split(':').map {|x| x.gsub("\\n","\n").gsub("\\c",":").gsub("\\\\", "\\") }
+            headers[kv.first] = kv.last
+          end
+        end
+      end
+    end
+
+    class Connect < Frame
+      def initialize(options)
+        headers['accept-version'] = '1.1'
+        headers['host'] = options['host'] if options['host']
+        headers['heart-beat'] = "0,0"
+        headers['login'] = options['login'] if options['login']
+        headers['passcode'] = options['passcode'] if options['passcode']
+      end
+    end
+
+    class Connected < ServerFrame
+    end
+
+    class Message < ServerFrame
+    end
+
+    class Send < Frame
+      def initialize(queue, body, hdrs)
+        headers['destination'] = queue
+        headers.update(hdrs.reject{|k,v| %w(destination content-length).include? k })
+        headers['content-type'] ||= 'text/plain'
+        headers['content-length'] = body.bytesize.to_s
+        @body = body
+      end
+    end
+
+    class Subscribe < Frame
+      def initialize(queue)
+        headers['id'] = queue
+        headers['destination'] = queue
+        headers['ack'] = 'auto'
+      end
+    end
+
+    class Unsubscribe < Frame
+      def initialize(queue)
+        headers['id'] = queue
+      end
+    end
+
+    class Disconnect < Frame
+    end
+  end
+end

data/lib/klomp/sentinel.rb

@@ -0,0 +1,21 @@
+class Klomp
+  class Sentinel
+    def initialize(connection)
+      @connection = connection
+      Thread.new { run } unless @connection.connected?
+    end
+
+    def run
+      fib_state = [0, 1]
+      loop do
+        begin
+          @connection.reconnect
+          break
+        rescue
+          sleep fib_state[1]
+          fib_state = [fib_state[1], fib_state[0]+fib_state[1]]
+        end
+      end
+    end
+  end
+end

data/spec/acceptance/acceptance_spec.rb

@@ -0,0 +1,152 @@
+require 'spec_helper'
+require 'json'
+require 'open-uri'
+
+describe "Klomp acceptance", :acceptance => true do
+
+  Given(:server) { "127.0.0.1:61613" }
+  Given(:credentials) { %w(admin password) }
+  Given(:options) { Hash[*%w(login passcode).zip(credentials).flatten] }
+  Given(:clients) { [] }
+  Given(:klomp) { Klomp.new(server, options).tap {|l| clients << l } }
+
+  context "connect" do
+
+    When { klomp }
+
+    Then { klomp.should be_connected }
+
+  end
+
+  context "publish" do
+
+    When { klomp.publish "/queue/greeting", "hello" }
+
+    Then do
+      vhosts = apollo_api_get_json "/broker/virtual-hosts.json"
+      vhost = vhosts['rows'].detect {|row| row['queues'].include?('greeting') }['id']
+      @queue_path = "/broker/virtual-hosts/#{vhost}/queues/greeting.json"
+      queue = apollo_api_get_json @queue_path
+      queue['metrics']['queue_items'].to_i.should > 0
+    end
+
+    after do
+      apollo_api_delete @queue_path
+    end
+  end
+
+  context "subscribe" do
+
+    Given(:subscriber) { double("subscriber") }
+    Given { klomp.publish "/queue/greeting", "hello subscriber!" }
+
+    When do
+      subscriber.stub!(:call).and_return {|msg| subscriber.stub!(:message => msg) }
+      klomp.subscribe "/queue/greeting", subscriber
+      sleep 1         # HAX: waiting for message to be pushed back and processed
+    end
+
+    Then do
+      subscriber.should have_received(:call).with(an_instance_of(Klomp::Frames::Message))
+      subscriber.message.body.should == "hello subscriber!"
+    end
+
+
+    context "and unsubscribe" do
+
+      When do
+        subscriber.reset
+        klomp.unsubscribe "/queue/greeting"
+        klomp.publish "/queue/greeting", "hello subscriber?"
+        sleep 1
+      end
+
+      Then do
+        subscriber.should_not have_received(:call)
+      end
+
+    end
+
+  end
+
+  context "throughput test", :performance => true do
+
+    require 'benchmark'
+
+    Given(:num_threads) { (ENV['THREADS'] || 4).to_i }
+    Given(:msgs_per_thread) { (ENV['MSGS'] || 10000).to_i }
+    Given(:total) { num_threads * msgs_per_thread }
+
+    Given do
+      trap("QUIT") do
+        Thread.list.each do |t|
+          $stderr.puts
+          $stderr.puts t.inspect
+          $stderr.puts t.backtrace.join("\n  ")
+        end
+      end
+    end
+
+    Given { klomp }
+
+    Then do
+      Thread.abort_on_exception = true
+
+      roundtrip_time = Benchmark.realtime do
+
+        Thread.new do
+          publish_time = Benchmark.realtime do
+            threads = []
+            1.upto(num_threads) do |i|
+              threads << Thread.new do
+                1.upto(msgs_per_thread) do |j|
+                  id = i * j
+                  print "." if id % 100 == 0
+                  klomp.publish "/queue/greeting", "hello #{id}!", "id" => "greeting-#{id}"
+                end
+              end
+            end
+            threads.each(&:join)
+          end
+
+          puts "\n--------------------------------------------------------------------------------\n" \
+          "Sending   #{total} messages took #{publish_time} using #{num_threads} threads\n" \
+          "--------------------------------------------------------------------------------\n"
+        end
+
+        ids = []
+        subscribe_time = Benchmark.realtime do
+          klomp.subscribe "/queue/greeting" do |msg|
+            id = msg.headers['id'][/(\d+)/, 1].to_i
+            print "," if id % 100 == 0
+            ids << id
+          end
+
+          Thread.pass until ids.length == total
+        end
+
+        puts "\n--------------------------------------------------------------------------------\n" \
+        "Receiving #{total} messages took #{subscribe_time}\n" \
+        "--------------------------------------------------------------------------------\n"
+      end
+      puts "\n--------------------------------------------------------------------------------\n" \
+      "Roundtrip to process #{total} messages: #{roundtrip_time} (#{total/roundtrip_time} msgs/sec)\n" \
+      "--------------------------------------------------------------------------------\n"
+    end
+  end
+
+  after { clients.each(&:disconnect) }
+
+  def apollo_mgmt_url(path)
+    "http://localhost:61680#{path}"
+  end
+
+  def apollo_api_get_json(path)
+    open(apollo_mgmt_url(path), :http_basic_authentication => credentials) {|f| JSON::parse(f.read) }
+  end
+
+  def apollo_api_delete(path)
+    `curl -s -f -X DELETE -u #{credentials.join(':').inspect} #{apollo_mgmt_url path}`
+    $?.should be_success
+  end
+end