stompserver 0.9.7 → 0.9.8

data/lib/stomp_server/queue/activerecord_queue.rb

@@ -0,0 +1,104 @@
+## Queue implementation using ActiveRecord
+##
+## all messages are stored in a single table
+## they are indexed by 'stomp_id' which is the stomp 'message-id' header
+## which must be unique accross all queues
+##
+require 'stomp_server/queue/ar_message'
+require 'yaml'
+
+module StompServer
+class ActiveRecordQueue
+
+  def initialize(configdir, storagedir)
+    # Default configuration, use SQLite for simplicity
+    db_params = {
+      'adapter' => 'sqlite3',
+      'database' => "#{configdir}/stompserver_development"
+    }
+    # Load DB configuration
+    db_config = "#{configdir}/database.yml"
+    puts "reading from #{db_config}"
+    if File.exists? db_config
+      db_params.merge! YAML::load(File.open(db_config))
+    end
+
+    puts "using #{db_params['database']} DB"
+
+    # Setup activerecord
+    ActiveRecord::Base.establish_connection(db_params)
+    # Development <TODO> fix this
+    ActiveRecord::Base.logger = Logger.new(STDERR)
+    ActiveRecord::Base.logger.level = Logger::INFO
+    # we need the connection, it can't be done earlier
+    ArMessage.reset_column_information
+    reload_queues
+    @stompid = StompServer::StompId.new
+  end
+
+  # Add a frame to the queue
+  def enqueue(queue_name, frame)
+    unless @frames[queue_name]
+      @frames[queue_name] = {
+        :last_index => 0,
+        :frames => [],
+      }
+    end
+    affect_msgid_and_store(frame, queue_name)
+    @frames[queue_name][:frames] << frame
+  end
+
+  # Get and remove a frame from the queue
+  def dequeue(queue_name)
+    return nil unless @frames[queue_name] && !@frames[queue_name][:frames].empty?
+    frame = @frames[queue_name][:frames].shift
+    remove_from_store(frame.headers['message-id'])
+    return frame
+  end
+
+  # Requeue the frame previously pending
+  def requeue(queue_name, frame)
+    @frames[queue_name][:frames] << frame
+    ArMessage.create!(:stomp_id => frame.headers['message-id'],
+                      :frame => frame)
+  end
+
+  # remove a frame from the store
+  def remove_from_store(message_id)
+    ArMessage.find_by_stomp_id(message_id).destroy
+  end
+
+  # store a frame (assigning it a message-id)
+  def affect_msgid_and_store(frame, queue_name)
+    msgid = assign_id(frame, queue_name)
+    ArMessage.create!(:stomp_id => msgid, :frame => frame)
+  end
+
+  def message_for?(queue_name)
+    @frames[queue_name] && !@frames[queue_name][:frames].empty?
+  end
+
+  def assign_id(frame, queue_name)
+    msgid = @stompid[@frames[queue_name][:last_index] += 1]
+    frame.headers['message-id'] = msgid
+  end
+
+  private
+  def reload_queues
+    @frames = Hash.new
+    ArMessage.find(:all).each { |message|
+      frame = message.frame
+      destination = frame.dest
+      msgid = message.stomp_id
+      @frames[destination] ||= Hash.new
+      @frames[destination][:frames] ||= Array.new
+      @frames[destination][:frames] << frame
+    }
+    # compute base index for each destination
+    @frames.each_pair { |destination,hash|
+      hash[:last_index] = hash[:frames].map{|f|
+        f.headers['message-id'].match(/(\d+)\Z/)[0].to_i}.max
+    }
+  end
+end
+end

data/lib/stomp_server/queue/ar_message.rb

@@ -0,0 +1,5 @@
+require 'active_record'
+
+class ArMessage < ActiveRecord::Base
+  serialize :frame
+end

data/lib/stomp_server/queue/dbm_queue.rb

@@ -0,0 +1,68 @@
+
+module StompServer
+class DBMQueue < Queue
+
+  def initialize *args
+    super
+    # Please don't use dbm files for storing large frames, it's problematic at best and uses large amounts of memory.
+    # sdbm croaks on marshalled data that contains certain characters, so we don't use it at all
+    @dbm = false
+    if RUBY_PLATFORM =~/linux|bsd/
+      types = ['bdb','dbm','gdbm']
+    else
+      types = ['bdb','gdbm']
+    end
+    types.each do |dbtype|
+      begin
+        require dbtype
+        @dbm = dbtype
+        puts "#{@dbm} loaded"
+        break
+      rescue LoadError => e
+      end
+    end
+    raise "No DBM library found. Tried bdb,dbm,gdbm" unless @dbm
+    @db = Hash.new
+    @queues.keys.each {|q| _open_queue(q)}
+  end
+
+  def dbmopen(dbname)
+    if @dbm == 'bdb'
+      BDB::Hash.new(dbname, nil, "a")
+    elsif @dbm == 'dbm'
+      DBM.open(dbname)
+    elsif @dbm == 'gdbm'
+      GDBM.open(dbname)
+    end
+  end
+
+
+  def _open_queue(dest)
+    queue_name = dest.gsub('/', '_')
+    dbname = @directory + '/' + queue_name
+    @db[dest] = Hash.new
+    @db[dest][:dbh] = dbmopen(dbname)
+    @db[dest][:dbname] = dbname
+  end
+
+
+  def _close_queue(dest)
+    @db[dest][:dbh].close
+    dbname = @db[dest][:dbname]
+    File.delete(dbname) if File.exists?(dbname)
+    File.delete("#{dbname}.db") if File.exists?("#{dbname}.db")
+    File.delete("#{dbname}.pag") if File.exists?("#{dbname}.pag")
+    File.delete("#{dbname}.dir") if File.exists?("#{dbname}.dir")
+  end
+
+  def _writeframe(dest,frame,msgid)
+    @db[dest][:dbh][msgid] = Marshal::dump(frame)
+  end
+
+  def _readframe(dest,msgid)
+    frame_image = @db[dest][:dbh][msgid]
+    Marshal::load(frame_image)
+  end
+
+end
+end

data/lib/stomp_server/queue/file_queue.rb

@@ -0,0 +1,47 @@
+
+module StompServer
+class FileQueue < Queue
+
+  def _close_queue(dest)
+    Dir.delete(@queues[dest][:queue_dir]) if File.directory?(@queues[dest][:queue_dir])
+  end
+
+  def _open_queue(dest)
+    # handle clashes between _ and /
+    queue_name = dest.gsub('_','__')
+    queue_name = dest.gsub('/','_')
+    queue_dir = @directory + '/' + queue_name
+    @queues[dest][:queue_dir] = queue_dir
+    Dir.mkdir(queue_dir) unless File.directory?(queue_dir)
+  end
+
+  def _writeframe(dest,frame_todump,msgid)
+    filename = "#{@queues[dest][:queue_dir]}/#{msgid}"
+    frame = frame_todump.dup
+    frame_body = frame.body
+    frame.body = ''
+    frame_image = Marshal.dump(frame)
+    framelen = sprintf("%08x", frame_image.length)
+    bodylen = sprintf("%08x", frame_body.length)
+    File.open(filename,'wb') {|f| f.syswrite("#{framelen}#{bodylen}#{frame_image}#{frame_body}")}
+    return true
+  end
+
+  def _readframe(dest,msgid)
+    filename = "#{@queues[dest][:queue_dir]}/#{msgid}"
+    file = nil
+    File.open(filename,'rb') {|f| file = f.read}
+    frame_len = file[0,8].hex
+    body_len = file[8,8].hex
+    frame = Marshal::load(file[16,frame_len])
+    frame.body = file[(frame_len + 16),body_len]
+    if File.delete(filename)
+      result = frame
+    else
+      result = false
+    end
+    return result
+  end
+end
+end
+

data/lib/stomp_server/queue/memory_queue.rb

@@ -0,0 +1,58 @@
+
+module StompServer
+class MemoryQueue
+
+  def initialize
+    @frame_index =0
+    @stompid = StompServer::StompId.new
+    @stats = Hash.new
+    @messages = Hash.new { Array.new }
+    puts "MemoryQueue initialized"
+  end
+
+  def stop
+  end
+
+  def monitor
+    stats = Hash.new
+    @messages.keys.each do |dest|
+     stats[dest] = {'size' => @messages[dest].size, 'enqueued' => @stats[dest][:enqueued], 'dequeued' => @stats[dest][:dequeued]}
+    end
+    stats
+  end
+
+  def dequeue(dest)
+    if frame = @messages[dest].shift
+      @stats[dest][:dequeued] += 1
+      return frame
+    else
+      return false
+    end
+  end
+
+  def enqueue(dest,frame)
+    @frame_index += 1
+    if @stats[dest]
+      @stats[dest][:enqueued] += 1
+    else
+      @stats[dest] = Hash.new
+      @stats[dest][:enqueued] = 1
+      @stats[dest][:dequeued] = 0
+    end
+    assign_id(frame, dest)
+    requeue(dest, frame)
+  end
+
+  def requeue(dest,frame)
+    @messages[dest] += [frame]
+  end
+
+  def message_for?(dest)
+    !@messages[dest].empty?
+  end
+
+  def assign_id(frame, dest)
+    frame.headers['message-id'] = @stompid[@frame_index]
+  end
+end
+end

data/lib/stomp_server/queue_manager.rb

@@ -0,0 +1,209 @@
+# QueueManager is used in conjunction with a storage class.  The storage class MUST implement the following two methods:
+#
+# - enqueue(queue name, frame)
+# enqueue pushes a frame to the top of the queue in FIFO order. It's return value is ignored. enqueue must also set the 
+# message-id and add it to the frame header before inserting the frame into the queue.
+#
+# - dequeue(queue name)
+# dequeue removes a frame from the bottom of the queue and returns it.
+#
+# - requeue(queue name,frame)
+# does the same as enqueue, except it puts the from at the bottom of the queue
+#
+# The storage class MAY implement the stop() method which can be used to do any housekeeping that needs to be done before 
+# stompserver shuts down. stop() will be called when stompserver is shut down.
+#
+# The storage class MAY implement the monitor() method.  monitor() should return a hash of hashes containing the queue statistics.
+# See the file queue for an example. Statistics are available to clients in /queue/monitor.
+#
+
+module StompServer
+class QueueMonitor
+
+  def initialize(qstore,queues)
+    @qstore = qstore
+    @queues = queues
+    @stompid = StompServer::StompId.new
+    puts "QueueManager initialized"
+  end
+
+  def start
+    count =0
+    EventMachine::add_periodic_timer 5, proc {count+=1; monitor(count) }
+  end
+
+  def monitor(count)
+    return unless @qstore.methods.include?('monitor')
+    users = @queues['/queue/monitor']
+    return if users.size == 0
+    stats = @qstore.monitor
+    return if stats.size == 0
+    body = ''
+
+    stats.each do |queue,qstats|
+      body << "Queue: #{queue}\n"
+      qstats.each {|stat,value| body << "#{stat}: #{value}\n"}
+      body << "\n"
+    end
+
+    headers = {
+      'message-id' => @stompid[count],
+      'destination' => '/queue/monitor',
+      'content-length' => body.size.to_s
+    }
+
+    frame = StompServer::StompFrame.new('MESSAGE', headers, body)
+    users.each {|user| user.user.stomp_send_data(frame)}
+  end
+end
+
+class QueueManager
+  Struct::new('QueueUser', :connection, :ack)
+
+  def initialize(qstore)
+    @qstore = qstore
+    @queues = Hash.new { Array.new }
+    @pending = Hash.new
+    if $STOMP_SERVER
+      monitor = StompServer::QueueMonitor.new(@qstore,@queues)
+      monitor.start
+      puts "Queue monitor started" if $DEBUG
+    end
+  end
+
+  def stop
+    @qstore.stop if @qstore.methods.include?('stop')
+  end
+
+  def subscribe(dest, connection, use_ack=false)
+    puts "Subscribing to #{dest}"
+    user = Struct::QueueUser.new(connection, use_ack)
+    @queues[dest] += [user]
+    send_destination_backlog(dest,user) unless dest == '/queue/monitor'
+  end
+
+  # Send at most one frame to a connection
+  # used when use_ack == true
+  def send_a_backlog(connection)
+    puts "Sending a backlog" if $DEBUG
+    # lookup queues with data for this connection
+    possible_queues = @queues.select{ |destination,users|
+      @qstore.message_for?(destination) &&
+        users.detect{|u| u.connection == connection}
+    }
+    if possible_queues.empty?
+      puts "Nothing left" if $DEBUG
+      return
+    end
+    # Get a random one (avoid artificial priority between queues
+    # without coding a whole scheduler, which might be desirable later)
+    dest,users = possible_queues[rand(possible_queues.length)]
+    user = users.find{|u| u.connection == connection}
+    frame = @qstore.dequeue(dest)
+    puts "Chose #{dest}" if $DEBUG
+    send_to_user(frame, user)
+  end
+
+  def send_destination_backlog(dest,user)
+    puts "Sending destination backlog for #{dest}" if $DEBUG
+    if user.ack
+      # only send one message (waiting for ack)
+      frame = @qstore.dequeue(dest)
+      send_to_user(frame, user) if frame
+    else
+      while frame = @qstore.dequeue(dest)
+        send_to_user(frame, user)
+      end
+    end
+  end
+
+  def unsubscribe(dest, connection)
+    puts "Unsubscribing from #{dest}"
+    @queues.each do |d, queue|
+      queue.delete_if { |qu| qu.connection == connection and d == dest}
+    end
+    @queues.delete(dest) if @queues[dest].empty?
+  end
+
+  def ack(connection, frame)
+    puts "Acking #{frame.headers['message-id']}" if $DEBUG
+    unless @pending[connection]
+      puts "No message pending for connection!"
+      return
+    end
+    msgid = frame.headers['message-id']
+    p_msgid = @pending[connection].headers['message-id']
+    if p_msgid != msgid
+      # We don't know what happened, we requeue
+      # (probably a client connecting to a restarted server)
+      frame = @pending[connection]
+      @qstore.requeue(frame.headers['destination'],frame)
+      puts "Invalid message-id (received #{msgid} != #{p_msgid})"
+    end
+    @pending.delete connection
+    # We are free to work now, look if there's something for us
+    send_a_backlog(connection)
+  end
+
+  def disconnect(connection)
+    puts "Disconnecting"
+    frame = @pending[connection]
+    if frame
+      @qstore.requeue(frame.headers['destination'],frame)
+      @pending.delete connection
+    end
+
+    @queues.each do |dest, queue|
+      queue.delete_if { |qu| qu.connection == connection }
+      @queues.delete(dest) if queue.empty?
+    end
+  end
+
+  def send_to_user(frame, user)
+    connection = user.connection
+    if user.ack
+      raise "other connection's end already busy" if @pending[connection]
+      @pending[connection] = frame
+    end
+    connection.stomp_send_data(frame)
+  end
+
+  def sendmsg(frame)
+    frame.command = "MESSAGE"
+    dest = frame.headers['destination']
+    puts "Sending a message to #{dest}: #{frame}"
+    # Lookup a user willing to handle this destination
+    available_users = @queues[dest].reject{|user| @pending[user.connection]}
+    if available_users.empty?
+      @qstore.enqueue(dest,frame)
+      return
+    end
+
+    # Look for a user with ack (we favor reliability)
+    reliable_user = available_users.find{|u| u.ack}
+
+    if reliable_user
+      # give it a message-id
+      @qstore.assign_id(frame, dest)
+      send_to_user(frame, reliable_user)
+    else
+      random_user = available_users[rand(available_users.length)]
+      # Note message-id header isn't set but we won't need it anyway
+      # <TODO> could break some clients: fix this
+      send_to_user(frame, random_user)
+    end
+  end
+
+  # For protocol handlers that want direct access to the queue
+  def dequeue(dest)
+    @qstore.dequeue(dest)
+  end
+
+  def enqueue(frame)
+    frame.command = "MESSAGE"
+    dest = frame.headers['destination']
+    @qstore.enqueue(dest,frame)
+  end
+
+end
+end