kim-toms-starling 1.0.0

data/lib/starling/handler.rb

@@ -0,0 +1,239 @@
+module StarlingServer
+
+  ##
+  # This is an internal class that's used by Starling::Server to handle the
+  # MemCache protocol and act as an interface between the Server and the
+  # QueueCollection.
+
+  class Handler < EventMachine::Connection
+
+    DATA_PACK_FMT = "Ia*".freeze
+
+    # ERROR responses
+    ERR_UNKNOWN_COMMAND = "CLIENT_ERROR bad command line format\r\n".freeze
+
+    CMD_RE = /\A(g)et\s+([!-~]{1,250})\s{0,}\r\n\Z|\A(s)et\s+([!-~]{1,250})\s+(\d+)\s+(\d+)\s+(\d+)\s{0,}(noreply)?\s{0,}\r\n\Z|\A(d)elete\s+([!-~]{1,250})\s+(\d+)\s{0,}\r\n\Z|\As(t)ats\s{0,}\r\n\Z|\As(h)utdown\s{0,}\r\n\Z|\A(q)uit\s{0,}\r\n\Z/
+    # First character will be
+    # g for get,
+    # s for set,
+    # d for delete,
+    # t for stats
+    # h for shutdown
+    # q for quit
+
+    # GET Responses
+
+    GET_RESPONSE       = "VALUE %s %s %s\r\n%s\r\nEND\r\n".freeze
+    GET_RESPONSE_EMPTY = "END\r\n".freeze
+
+    # SET Responses
+    SET_RESPONSE_SUCCESS  = "STORED\r\n".freeze
+    SET_RESPONSE_FAILURE  = "NOT STORED\r\n".freeze
+    SET_CLIENT_DATA_ERROR = "CLIENT_ERROR bad data chunk\r\nERROR\r\n".freeze
+    
+    # DELETE Responses
+    DELETE_RESPONSE = "END\r\n".freeze
+
+    # STAT Response
+    STATS_RESPONSE = "STAT pid %d\r
+STAT uptime %d\r
+STAT time %d\r
+STAT version %s\r
+STAT rusage_user %0.6f\r
+STAT rusage_system %0.6f\r
+STAT curr_items %d\r
+STAT total_items %d\r
+STAT bytes %d\r
+STAT curr_connections %d\r
+STAT total_connections %d\r
+STAT cmd_get %d\r
+STAT cmd_set %d\r
+STAT get_hits %d\r
+STAT get_misses %d\r
+STAT bytes_read %d\r
+STAT bytes_written %d\r
+STAT limit_maxbytes %d\r
+%sEND\r\n".freeze
+    QUEUE_STATS_RESPONSE = "STAT queue_%s_items %d\r
+STAT queue_%s_total_items %d\r
+STAT queue_%s_logsize %d\r
+STAT queue_%s_expired_items %d\r
+STAT queue_%s_age %d\r\n".freeze
+
+    @@next_session_id = 1
+
+    ##
+    # Creates a new handler for the MemCache protocol that communicates with a
+    # given client.
+
+    def initialize(options = {})
+      @opts = options
+    end
+
+    ##
+    # Process incoming commands from the attached client.
+
+    def post_init
+      @stash = []
+      @data = ""
+      @data_buf = ""
+      @server = @opts[:server]
+      @logger = StarlingServer::Base.logger
+      @expiry_stats = Hash.new(0)
+      @expected_length = nil
+      @server.stats[:total_connections] += 1
+      set_comm_inactivity_timeout @opts[:timeout]
+      @queue_collection = @opts[:queue]
+
+      @session_id = @@next_session_id
+      @@next_session_id += 1
+
+      peer = Socket.unpack_sockaddr_in(get_peername)
+      #@logger.debug "(#{@session_id}) New session from #{peer[1]}:#{peer[0]}"
+    end
+
+    def receive_data(incoming)
+      @server.stats[:bytes_read] += incoming.size
+      @data << incoming
+
+      while ind = @data.index("\r\n")
+        to_proc = @data[0,ind+2]
+        @data = @data[ind+2,@data.size]
+        response = process(to_proc)
+        send_data response if response
+      end
+    end
+
+    def process(data)
+      data = @data_buf + data unless @data_buf.size == 0
+      if @expected_length && data.size == @expected_length
+        @data_buf = ""
+        return  set_data(data)
+      elsif @expected_length
+        @data_buf = data
+        return
+      end
+      m = CMD_RE.match(data)
+      if ! m.nil?
+        args = m.captures.compact
+        case args[0][0]
+        when ?s #set command
+          @server.stats[:set_requests] += 1
+          set(args[1], args[2], args[3], args[4].to_i)
+        when ?g #get command
+          @server.stats[:get_requests] += 1
+          get(args[1])
+        when ?t # stats command
+          stats
+        when ?h # shutdown command
+          Runner::shutdown
+        when ?d # delete command
+          delete args[1]
+        when ?q # quit command
+          # ignore the command, client is closing connection.
+          return nil
+        else
+          logger.warn "Unknown command (in case): #{data}."
+          respond ERR_UNKNOWN_COMMAND
+        end
+      else
+        logger.warn "Unknown command: #{data}."
+        respond ERR_UNKNOWN_COMMAND
+      end
+    end
+
+    def unbind
+      #@logger.debug "(#{@session_id}) connection ends"
+    end
+
+  private
+  
+    def delete(queue)
+      @queue_collection.delete(queue)
+      respond DELETE_RESPONSE
+    end
+  
+    def respond(str, *args)
+      response = sprintf(str, *args)
+      @server.stats[:bytes_written] += response.length
+      response
+    end
+
+    def set(key, flags, expiry, len)
+      @expected_length = len + 2
+      @stash = [ key, flags, expiry ]
+      nil
+    end
+
+    def set_data(incoming)
+      key, flags, expiry = @stash
+      data = incoming.slice(0...@expected_length-2)
+      @stash = []
+      @expected_length = nil
+
+      internal_data = [expiry.to_i, data].pack(DATA_PACK_FMT)
+      if @queue_collection.put(key, internal_data)
+        respond SET_RESPONSE_SUCCESS
+      else
+        respond SET_RESPONSE_FAILURE
+      end
+    end
+
+    def get(key)
+      now = Time.now.to_i
+
+      while response = @queue_collection.take(key)
+        expiry, data = response.unpack(DATA_PACK_FMT)
+
+        break if expiry == 0 || expiry >= now
+
+        @expiry_stats[key] += 1
+        expiry, data = nil
+      end
+
+      if data
+        respond GET_RESPONSE, key, 0, data.size, data
+      else
+        respond GET_RESPONSE_EMPTY
+      end
+    end
+
+    def stats
+      respond STATS_RESPONSE,
+        Process.pid, # pid
+        Time.now - @server.stats(:start_time), # uptime
+        Time.now.to_i, # time
+        StarlingServer::VERSION, # version
+        Process.times.utime, # rusage_user
+        Process.times.stime, # rusage_system
+        @queue_collection.stats(:current_size), # curr_items
+        @queue_collection.stats(:total_items), # total_items
+        @queue_collection.stats(:current_bytes), # bytes
+        @server.stats(:connections), # curr_connections
+        @server.stats(:total_connections), # total_connections
+        @server.stats(:get_requests), # get count
+        @server.stats(:set_requests), # set count
+        @queue_collection.stats(:get_hits),
+        @queue_collection.stats(:get_misses),
+        @server.stats(:bytes_read), # total bytes read
+        @server.stats(:bytes_written), # total bytes written
+        0, # limit_maxbytes
+        queue_stats
+    end
+
+    def queue_stats
+      @queue_collection.queues.inject("") do |m,(k,v)|
+        m + sprintf(QUEUE_STATS_RESPONSE,
+                      k, v.length,
+                      k, v.total_items,
+                      k, v.logsize,
+                      k, @expiry_stats[k],
+                      k, v.current_age)
+      end
+    end
+
+    def logger
+      @logger
+    end
+  end
+end

data/lib/starling/persistent_queue.rb

@@ -0,0 +1,156 @@
+module StarlingServer
+
+  ##
+  # PersistentQueue is a subclass of Ruby's thread-safe Queue class. It adds a
+  # transactional log to the in-memory Queue, which enables quickly rebuilding
+  # the Queue in the event of a sever outage.
+
+  class PersistentQueue < Queue
+
+    ##
+    # When a log reaches the SOFT_LOG_MAX_SIZE, the Queue will wait until
+    # it is empty, and will then rotate the log file.
+
+    SOFT_LOG_MAX_SIZE = 16 * (1024**2) # 16 MB
+
+    TRX_CMD_PUSH = "\000".freeze
+    TRX_CMD_POP = "\001".freeze
+
+    TRX_PUSH = "\000%s%s".freeze
+    TRX_POP = "\001".freeze
+
+    attr_reader :initial_bytes
+    attr_reader :total_items
+    attr_reader :logsize
+    attr_reader :current_age
+
+    ##
+    # Create a new PersistentQueue at +persistence_path+/+queue_name+.
+    # If a queue log exists at that path, the Queue will be loaded from
+    # disk before being available for use.
+
+    def initialize(persistence_path, queue_name, debug = false)
+      @persistence_path = persistence_path
+      @queue_name = queue_name
+      @total_items = 0
+      super()
+      @initial_bytes = replay_transaction_log(debug)
+      @current_age = 0
+    end
+
+    ##
+    # Pushes +value+ to the queue. By default, +push+ will write to the
+    # transactional log. Set +log_trx=false+ to override this behaviour.
+
+    def push(value, log_trx = true)
+      if log_trx
+        raise NoTransactionLog unless @trx
+        size = [value.size].pack("I")
+        transaction sprintf(TRX_PUSH, size, value)
+      end
+
+      @total_items += 1
+      super([now_usec, value])
+    end
+
+    ##
+    # Retrieves data from the queue.
+
+    def pop(log_trx = true)
+      raise NoTransactionLog if log_trx && !@trx
+
+      begin
+        rv = super(!log_trx)
+      rescue ThreadError
+        puts "WARNING: The queue was empty when trying to pop(). Technically this shouldn't ever happen. Probably a bug in the transactional underpinnings. Or maybe shutdown didn't happen cleanly at some point. Ignoring."
+        rv = [now_usec, '']
+      end
+      transaction "\001" if log_trx
+      @current_age = (now_usec - rv[0]) / 1000
+      rv[1]
+    end
+
+    ##
+    # Safely closes the transactional queue.
+
+    def close
+      # Ok, yeah, this is lame, and is *technically* a race condition. HOWEVER,
+      # the QueueCollection *should* have stopped processing requests, and I don't
+      # want to add yet another Mutex around all the push and pop methods. So we
+      # do the next simplest thing, and minimize the time we'll stick around before
+      # @trx is nil.
+      @not_trx = @trx
+      @trx = nil
+      @not_trx.close
+    end
+    
+    def purge
+      close
+      File.delete(log_path)
+    end
+
+    private
+
+    def log_path #:nodoc:
+      File.join(@persistence_path, @queue_name)
+    end
+
+    def reopen_log #:nodoc:
+      @trx = File.new(log_path, File::CREAT|File::RDWR)
+      @logsize = File.size(log_path)
+    end
+
+    def rotate_log #:nodoc:
+      @trx.close
+      backup_logfile = "#{log_path}.#{Time.now.to_i}"
+      File.rename(log_path, backup_logfile)
+      reopen_log
+      File.unlink(backup_logfile)
+    end
+
+    def replay_transaction_log(debug) #:nodoc:
+      reopen_log
+      bytes_read = 0
+
+      print "Reading back transaction log for #{@queue_name} " if debug
+
+      while !@trx.eof?
+        cmd = @trx.read(1)
+        case cmd
+        when TRX_CMD_PUSH
+          print ">" if debug
+          raw_size = @trx.read(4)
+          next unless raw_size
+          size = raw_size.unpack("I").first
+          data = @trx.read(size)
+          next unless data
+          push(data, false)
+          bytes_read += data.size
+        when TRX_CMD_POP
+          print "<" if debug
+          bytes_read -= pop(false).size
+        else
+          puts "Error reading transaction log: " +
+               "I don't understand '#{cmd}' (skipping)." if debug
+        end
+      end
+
+      print " done.\n" if debug
+
+      return bytes_read
+    end
+
+    def transaction(data) #:nodoc:
+      raise "no transaction log handle. that totally sucks." unless @trx
+
+      @trx.write_nonblock data
+      @logsize += data.size
+      rotate_log if @logsize > SOFT_LOG_MAX_SIZE && self.length == 0
+    end
+    
+    def now_usec
+      now = Time.now
+      now.to_i * 1000000 + now.usec
+    end
+  end
+end

data/lib/starling/queue_collection.rb

@@ -0,0 +1,147 @@
+require 'thread'
+require File.join(File.dirname(__FILE__), 'persistent_queue')
+
+module StarlingServer
+  class InaccessibleQueuePath < Exception #:nodoc:
+  end
+
+  ##
+  # QueueCollection is a proxy to a collection of PersistentQueue instances.
+
+  class QueueCollection
+
+    ##
+    # Create a new QueueCollection at +path+
+
+    def initialize(path)
+      unless File.directory?(path) && File.writable?(path)
+        raise InaccessibleQueuePath.new("'#{path}' must exist and be read-writable by #{Etc.getpwuid(Process.uid).name}.")
+      end
+
+      @shutdown_mutex = Mutex.new
+
+      @path = path
+      @logger = StarlingServer::Base.logger
+
+      @queues = {}
+      @queue_init_mutexes = {}
+
+      @stats = Hash.new(0)
+    end
+
+    ##
+    # Puts +data+ onto the queue named +key+
+
+    def put(key, data)
+      queue = queues(key)
+      return nil unless queue
+
+      @stats[:current_bytes] += data.size
+      @stats[:total_items] += 1
+
+      queue.push(data)
+
+      return true
+    end
+
+    ##
+    # Retrieves data from the queue named +key+
+
+    def take(key)
+      queue = queues(key)
+      if queue.nil? || queue.length == 0
+        @stats[:get_misses] += 1
+        return nil
+      else
+        @stats[:get_hits] += 1
+      end
+      result = queue.pop
+      @stats[:current_bytes] -= result.size
+      result
+    end
+    
+    def delete(key)
+      queue = @queues.delete(key)
+      return if queue.nil?
+      queue.purge
+    end
+
+    ##
+    # Returns all active queues.
+
+    def queues(key=nil)
+      return nil if @shutdown_mutex.locked?
+
+      return @queues if key.nil?
+
+      # First try to return the queue named 'key' if it's available.
+      return @queues[key] if @queues[key]
+
+      # If the queue wasn't available, create or get the mutex that will
+      # wrap creation of the Queue.
+      @queue_init_mutexes[key] ||= Mutex.new
+
+      # Otherwise, check to see if another process is already loading
+      # the queue named 'key'.
+      if @queue_init_mutexes[key].locked?
+        # return an empty/false result if we're waiting for the queue
+        # to be loaded and we're not the first process to request the queue
+        return nil
+      else
+        begin
+          @queue_init_mutexes[key].lock
+          # we've locked the mutex, but only go ahead if the queue hasn't
+          # been loaded. There's a race condition otherwise, and we could
+          # end up loading the queue multiple times.
+          if @queues[key].nil?
+            @queues[key] = PersistentQueue.new(@path, key)
+            @stats[:current_bytes] += @queues[key].initial_bytes
+          end
+        rescue Object => exc
+          puts "ZOMG There was an exception reading back the queue. That totally sucks."
+          puts "The exception was: #{exc}. Backtrace: #{exc.backtrace.join("\n")}"
+        ensure
+          @queue_init_mutexes[key].unlock
+        end
+      end
+
+      return @queues[key]
+    end
+
+    ##
+    # Returns statistic +stat_name+ for the QueueCollection.
+    #
+    # Valid statistics are:
+    #
+    #   [:get_misses]    Total number of get requests with empty responses
+    #   [:get_hits]      Total number of get requests that returned data
+    #   [:current_bytes] Current size in bytes of items in the queues
+    #   [:current_size]  Current number of items across all queues
+    #   [:total_items]   Total number of items stored in queues.
+
+    def stats(stat_name)
+      case stat_name
+      when nil; @stats
+      when :current_size; current_size
+      else; @stats[stat_name]
+      end
+    end
+
+    ##
+    # Safely close all queues.
+
+    def close
+      @shutdown_mutex.lock
+      @queues.each_pair do |name,queue|
+        queue.close
+        @queues.delete(name)
+      end
+    end
+
+    private
+
+    def current_size #:nodoc:
+      @queues.inject(0) { |m, (k,v)| m + v.length }
+    end
+  end
+end

data/lib/starling/server.rb

@@ -0,0 +1,125 @@
+require 'socket'
+require 'logger'
+require 'rubygems'
+require 'eventmachine'
+
+here = File.dirname(__FILE__)
+
+require File.join(here, 'queue_collection')
+require File.join(here, 'handler')
+
+module StarlingServer
+
+  VERSION = "1.0.0"
+  
+  class Base
+    attr_reader :logger
+
+    DEFAULT_HOST    = '127.0.0.1'
+    DEFAULT_PORT    = 22122
+    DEFAULT_PATH    = "/tmp/starling/"
+    DEFAULT_TIMEOUT = 60
+
+    ##
+    # Initialize a new Starling server and immediately start processing
+    # requests.
+    #
+    # +opts+ is an optional hash, whose valid options are:
+    #
+    #   [:host]     Host on which to listen (default is 127.0.0.1).
+    #   [:port]     Port on which to listen (default is 22122).
+    #   [:path]     Path to Starling queue logs. Default is /tmp/starling/
+    #   [:timeout]  Time in seconds to wait before closing connections.
+    #   [:logger]   A Logger object, an IO handle, or a path to the log.
+    #   [:loglevel] Logger verbosity. Default is Logger::ERROR.
+    #
+    # Other options are ignored.
+
+    def self.start(opts = {})
+      server = self.new(opts)
+      server.run
+    end
+
+    ##
+    # Initialize a new Starling server, but do not accept connections or
+    # process requests.
+    #
+    # +opts+ is as for +start+
+
+    def initialize(opts = {})
+      @opts = {
+        :host    => DEFAULT_HOST,
+        :port    => DEFAULT_PORT,
+        :path    => DEFAULT_PATH,
+        :timeout => DEFAULT_TIMEOUT,
+        :server  => self
+      }.merge(opts)
+
+      @stats = Hash.new(0)
+
+      FileUtils.mkdir_p(@opts[:path])
+
+    end
+
+    ##
+    # Start listening and processing requests.
+
+    def run
+      @stats[:start_time] = Time.now
+
+      if @opts[:syslog_channel]
+        begin
+          require 'syslog_logger'
+          @@logger = SyslogLogger.new(@opts[:syslog_channel])
+        rescue LoadError
+          # SyslogLogger isn't available, so we're just going to use Logger
+        end
+      end
+
+      @@logger ||= case @opts[:logger]
+                   when IO, String; Logger.new(@opts[:logger])
+                   when Logger; @opts[:logger]
+                   else; Logger.new(STDERR)
+                   end
+
+      begin
+        @opts[:queue] = QueueCollection.new(@opts[:path])
+      rescue InaccessibleQueuePath => e
+        puts "Error: #{e.message}"
+        exit 1
+      end
+      @@logger.level = @opts[:log_level] || Logger::ERROR
+
+      @@logger.info "Starling STARTUP on #{@opts[:host]}:#{@opts[:port]}"
+
+      EventMachine.epoll
+      EventMachine.set_descriptor_table_size(4096)
+      EventMachine.run do
+        EventMachine.start_server(@opts[:host], @opts[:port], Handler, @opts)
+      end
+
+      # code here will get executed on shutdown:
+      @opts[:queue].close
+    end
+
+    def self.logger
+      @@logger
+    end
+
+
+    ##
+    # Stop accepting new connections and shutdown gracefully.
+
+    def stop
+      EventMachine.stop_event_loop
+    end
+
+    def stats(stat = nil) #:nodoc:
+      case stat
+      when nil; @stats
+      when :connections; 1
+      else; @stats[stat]
+      end
+    end
+  end
+end