stompserver 0.9.7 → 0.9.8

data/lib/stomp_server/protocols/http.rb

@@ -0,0 +1,128 @@
+
+  class Mongrel::HttpRequest
+    attr_reader :body, :params
+
+    def initialize(params, initial_body)
+      @params = params
+      @body = StringIO.new
+      @body.write params.http_body
+    end
+  end
+
+module StompServer
+module StompServer::Protocols
+
+  class Http < EventMachine::Connection
+
+    def initialize *args
+      super
+      @buf = ''
+    end
+
+
+    def post_init
+      @parser = Mongrel::HttpParser.new
+      @params = Mongrel::HttpParams.new
+      @nparsed = 0
+      @request = nil
+      @request_method = nil
+      @request_length = 0
+      @state = :headers
+      @headers_out = {'Content-Length' => 0, 'Content-Type' => 'text/plain; charset=UTF-8'}
+    end
+
+    def receive_data data
+      parse_request(data)
+    end
+
+    def parse_request data
+      @buf << data
+      case @state
+      when :headers
+        @nparsed = @parser.execute(@params, @buf, @nparsed)
+        if @parser.finished?
+          @request = Mongrel::HttpRequest.new(@params,@buf)
+          @request_method = @request.params[Mongrel::Const::REQUEST_METHOD]
+          content_length = @request.params[Mongrel::Const::CONTENT_LENGTH].to_i
+          @request_length = @nparsed + content_length
+          @remain = content_length - @request.params.http_body.length
+          if @remain <= 0
+            @buf = @buf[@request_length+1..-1] || ''
+            process_request
+            post_init
+            return
+          end
+          @request.body.write @request.params.http_body
+          @state = :body
+        end
+      when :body
+        @remain -= @request.body.write data[0...@remain]
+        if @remain <= 0
+          @buf = @buf[@request_length+1..-1] || ''
+          process_request
+          post_init
+          return
+        end
+      end
+    end
+
+    def process_request
+      begin
+        @request.body.rewind
+        dest = @request.params[Mongrel::Const::REQUEST_PATH]
+        case @request_method
+        when 'PUT'
+          @frame = StompServer::StompFrame.new
+          @frame.command = 'SEND'
+          @frame.body = @request.body.read
+          @frame.headers['destination'] = dest
+          if @@queue_manager.enqueue(@frame)
+            create_response('200','Message Enqueued')
+          else
+            create_response('500','Error enqueueing message')
+          end
+        when 'GET'
+          if frame = @@queue_manager.dequeue(dest)
+            @headers_out['message-id'] = frame.headers['message-id']
+            create_response('200',frame.body)
+          else
+            create_response('404','No messages in queue')
+          end
+        else
+          create_response('500','Invalid Command')
+        end
+      rescue Exception => e
+        puts "err: #{e} #{e.backtrace.join("\n")}"
+        create_response('500',e)
+      end
+    end
+
+    def unbind
+      puts "Closing connection"
+      close_connection_after_writing
+    end
+
+    def create_response(code,response_text)
+      response = ''
+      @headers_out['Content-Length'] = response_text.size
+
+      case code
+      when '200'
+        response << "HTTP/1.1 200 OK\r\n"
+      when '500'
+        response << "HTTP/1.1 500 Server Error\r\n"
+      when '404'
+        response << "HTTP/1.1 404 Message Not Found\r\n"
+      end
+      @headers_out.each_pair do |key, value|
+        response << "#{key}:#{value}\r\n"
+      end
+      response << "\r\n"
+      response << response_text
+      send_data(response)
+      unbind if @request.params['HTTP_CONNECTION'] == 'close'
+    end
+  end
+
+end
+end

data/lib/stomp_server/protocols/stomp.rb

@@ -0,0 +1,186 @@
+
+module StompServer
+module StompServer::Protocols
+VALID_COMMANDS = [:connect, :send, :subscribe, :unsubscribe, :begin, :commit, :abort, :ack, :disconnect]
+
+class Stomp < EventMachine::Connection
+
+  def initialize *args
+    super
+  end
+
+  def post_init
+    @sfr = StompServer::StompFrameRecognizer.new
+    @transactions = {}
+    @connected = false
+  end
+
+  def receive_data(data)
+    stomp_receive_data(data)
+  end
+ 
+  def stomp_receive_data(data)
+    begin
+      puts "receive_data: #{data.inspect}" if $DEBUG
+      @sfr << data
+      process_frames
+    rescue Exception => e
+      puts "err: #{e} #{e.backtrace.join("\n")}"
+      send_error(e.to_s)
+      close_connection_after_writing
+    end
+  end 
+ 
+  def stomp_receive_frame(frame)
+    begin
+      puts "receive_frame: #{frame.inspect}" if $DEBUG
+      process_frame(frame)
+    rescue Exception => e
+      puts "err: #{e} #{e.backtrace.join("\n")}"
+      send_error(e.to_s)
+      close_connection_after_writing
+    end
+  end
+ 
+  def process_frames
+    frame = nil
+    process_frame(frame) while frame = @sfr.frames.shift
+  end
+  
+  def process_frame(frame)
+    cmd = frame.command.downcase.to_sym
+    raise "Unhandled frame: #{cmd}" unless VALID_COMMANDS.include?(cmd)
+    raise "Not connected" if !@connected && cmd != :connect
+
+    # I really like this code, but my needs are a little trickier
+    # 
+
+    if trans = frame.headers['transaction']
+      handle_transaction(frame, trans, cmd)
+    else
+      cmd = :sendmsg if cmd == :send
+      send(cmd, frame) 
+    end
+    
+    send_receipt(frame.headers['receipt']) if frame.headers['receipt']
+  end
+  
+  def handle_transaction(frame, trans, cmd)
+    if [:begin, :commit, :abort].include?(cmd)
+      send(cmd, frame, trans)
+    else
+      raise "transaction does not exist" unless @transactions.has_key?(trans)
+      @transactions[trans] << frame
+    end    
+  end
+  
+  def connect(frame)
+    if @@auth_required
+      unless frame.headers['login'] and frame.headers['passcode'] and  @@stompauth.authorized[frame.headers['login']] == frame.headers['passcode']
+        raise "Invalid Login"
+      end
+    end
+    puts "Connecting" if $DEBUG
+    response = StompServer::StompFrame.new("CONNECTED", {'session' => 'wow'})
+    stomp_send_data(response)
+    @connected = true
+  end
+  
+  def sendmsg(frame)
+    # set message id
+    if frame.dest.match(%r|^/queue|)
+      @@queue_manager.sendmsg(frame)
+    else
+      frame.headers['message-id'] = "msg-#stompcma-#{@@topic_manager.next_index}"
+      @@topic_manager.sendmsg(frame)
+    end
+  end
+  
+  def subscribe(frame)
+    use_ack = false
+    use_ack = true  if frame.headers['ack'] == 'client'
+    if frame.dest =~ %r|^/queue|
+      @@queue_manager.subscribe(frame.dest, self,use_ack)
+    else
+      @@topic_manager.subscribe(frame.dest, self)
+    end
+  end
+  
+  def unsubscribe(frame)
+    if frame.dest =~ %r|^/queue|
+      @@queue_manager.unsubscribe(frame.dest,self)
+    else
+      @@topic_manager.unsubscribe(frame.dest,self)
+    end
+  end
+  
+  def begin(frame, trans=nil)
+    raise "Missing transaction" unless trans
+    raise "transaction exists" if @transactions.has_key?(trans)
+    @transactions[trans] = []
+  end
+  
+  def commit(frame, trans=nil)
+    raise "Missing transaction" unless trans
+    raise "transaction does not exist" unless @transactions.has_key?(trans)
+    
+    (@transactions[trans]).each do |frame|
+      frame.headers.delete('transaction')
+      process_frame(frame)
+    end
+    @transactions.delete(trans)
+  end
+  
+  def abort(frame, trans=nil)
+    raise "Missing transaction" unless trans
+    raise "transaction does not exist" unless @transactions.has_key?(trans)
+    @transactions.delete(trans)
+  end
+  
+  def ack(frame)
+    @@queue_manager.ack(self, frame)
+  end
+  
+  def disconnect(frame)
+    puts "Polite disconnect" if $DEBUG
+    close_connection_after_writing
+  end
+
+  def unbind
+    p "Unbind called" if $DEBUG
+    @connected = false
+    @@queue_manager.disconnect(self)
+    @@topic_manager.disconnect(self)
+  end
+ 
+  def connected?
+    @connected
+  end
+ 
+  def send_message(msg)
+    msg.command = "MESSAGE"
+    stomp_send_data(msg)
+  end
+    
+  def send_receipt(id)
+    send_frame("RECEIPT", { 'receipt-id' => id})
+  end
+  
+  def send_error(msg)
+    send_frame("ERROR",{'message' => 'See below'},msg)
+  end
+ 
+  def stomp_send_data(frame)
+    send_data(frame.to_s)
+    puts "Sending frame #{frame.to_s}" if $DEBUG
+  end
+
+  def send_frame(command, headers={}, body='')
+    headers['content-length'] = body.size.to_s
+    response = StompServer::StompFrame.new(command, headers, body)
+    stomp_send_data(response)
+  end
+end
+
+end
+end

data/lib/stomp_server/queue.rb

@@ -0,0 +1,140 @@
+
+module StompServer
+class Queue
+
+  def initialize(directory='.stompserver', delete_empty=true)
+    @stompid = StompServer::StompId.new
+    @delete_empty = delete_empty
+    @directory = directory
+    Dir.mkdir(@directory) unless File.directory?(@directory)
+    if File.exists?("#{@directory}/qinfo")
+      qinfo = Hash.new
+      File.open("#{@directory}/qinfo", "rb") { |f| qinfo = Marshal.load(f.read)}
+      @queues = qinfo[:queues]
+      @frames = qinfo[:frames]
+    else
+      @queues = Hash.new
+      @frames = Hash.new
+    end
+
+    @queues.keys.each do |dest|
+      puts "Queue #{dest} size=#{@queues[dest][:size]} enqueued=#{@queues[dest][:enqueued]} dequeued=#{@queues[dest][:dequeued]}" if $DEBUG
+    end
+
+    puts "Queue initialized in #{@directory}"
+
+    # Cleanup dead queues and save the state of the queues every so often.  Alternatively we could save the queue state every X number
+    # of frames that are put in the queue.  Should probably also read it after saving it to confirm integrity.
+    # Removed, this badly corrupt the queue when stopping with messages
+    #EventMachine::add_periodic_timer 1800, proc {@queues.keys.each {|dest| close_queue(dest)};save_queue_state }
+  end
+
+  def stop
+    puts "Shutting down Queue"
+
+    @queues.keys.each {|dest| close_queue(dest)}
+    @queues.keys.each do |dest|
+      puts "Queue #{dest} size=#{@queues[dest][:size]} enqueued=#{@queues[dest][:enqueued]} dequeued=#{@queues[dest][:dequeued]}" if $DEBUG
+    end
+    save_queue_state
+  end
+
+  def save_queue_state
+    puts "Saving Queue State" if $DEBUG
+    qinfo = {:queues => @queues, :frames => @frames}
+    File.open("#{@directory}/qinfo", "wb") { |f| f.write Marshal.dump(qinfo)}
+  end
+
+  def monitor
+    stats = Hash.new
+    @queues.keys.each do |dest|
+      stats[dest] = {'size' => @queues[dest][:size], 'enqueued' => @queues[dest][:enqueued], 'dequeued' => @queues[dest][:dequeued]}
+    end
+    stats
+  end
+
+  def close_queue(dest)
+    if @queues[dest][:size] == 0 and @queues[dest][:frames].size == 0 and @delete_empty
+      _close_queue(dest)
+      @queues.delete(dest)
+      @frames.delete(dest)
+      puts "Queue #{dest} removed." if $DEBUG
+    end
+  end
+
+  def open_queue(dest)
+    @queues[dest] = Hash.new
+    @frames[dest] = Hash.new
+    @queues[dest][:size] = 0
+    @queues[dest][:frames] = Array.new
+    @queues[dest][:msgid] = 1
+    @queues[dest][:enqueued] = 0
+    @queues[dest][:dequeued] = 0
+    @queues[dest][:exceptions] = 0
+    _open_queue(dest)
+    puts "Created queue #{dest}" if $DEBUG
+  end
+
+  def requeue(dest,frame)
+    open_queue(dest) unless @queues.has_key?(dest)
+    msgid = frame.headers['message-id']
+    if frame.headers['max-exceptions'] and @frames[dest][msgid][:exceptions] >= frame.headers['max-exceptions'].to_i
+      enqueue("/queue/deadletter",frame)
+      return
+    end
+    writeframe(dest,frame,msgid)
+    @queues[dest][:frames].unshift(msgid)
+    @frames[dest][msgid][:exceptions] += 1
+    @queues[dest][:dequeued] -= 1
+    @queues[dest][:exceptions] += 1
+    @queues[dest][:size] += 1
+    save_queue_state
+    return true
+  end
+
+  def enqueue(dest,frame)
+    open_queue(dest) unless @queues.has_key?(dest)
+    msgid = assign_id(frame, dest)
+    writeframe(dest,frame,msgid)
+    @queues[dest][:frames].push(msgid)
+    @frames[dest][msgid] = Hash.new
+    @frames[dest][msgid][:exceptions] =0
+    @frames[dest][msgid][:client_id] = frame.headers['client-id'] if frame.headers['client-id']
+    @frames[dest][msgid][:expires] = frame.headers['expires'] if frame.headers['expires']
+    @queues[dest][:msgid] += 1
+    @queues[dest][:enqueued] += 1
+    @queues[dest][:size] += 1
+    save_queue_state
+    return true
+  end
+
+  def dequeue(dest)
+    return false unless message_for?(dest)
+    msgid = @queues[dest][:frames].shift
+    frame = readframe(dest,msgid)
+    @queues[dest][:size] -= 1
+    @queues[dest][:dequeued] += 1
+    @queues[dest].delete(msgid)
+    close_queue(dest)
+    save_queue_state
+    return frame
+  end
+
+  def message_for?(dest)
+    return (@queues.has_key?(dest) and (!@queues[dest][:frames].empty?))
+  end
+
+  def writeframe(dest,frame,msgid)
+    _writeframe(dest,frame,msgid)
+  end
+
+  def readframe(dest,msgid)
+    _readframe(dest,msgid)
+  end
+
+  def assign_id(frame, dest)
+    frame.headers['message-id'] = @stompid[@queues[dest][:msgid]]
+  end
+end
+end
+