ryana-starling 0.9.8

data/lib/starling/handler.rb

@@ -0,0 +1,233 @@
+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
+
+    # GET Responses
+    GET_COMMAND = /\Aget (.{1,250})\s*\r\n/m
+    GET_RESPONSE       = "VALUE %s %s %s\r\n%s\r\nEND\r\n".freeze
+    GET_RESPONSE_EMPTY = "END\r\n".freeze
+
+    # SET Responses
+    SET_COMMAND = /\Aset (.{1,250}) ([0-9]+) ([0-9]+) ([0-9]+)\r\n/m
+    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_COMMAND = /\Adelete (.{1,250}) ([0-9]+)\r\n/m
+    DELETE_RESPONSE = "END\r\n".freeze
+
+    # STAT Response
+    STATS_COMMAND = /\Astats\r\n/m
+    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
+
+    SHUTDOWN_COMMAND = /\Ashutdown\r\n/m
+
+
+    @@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 data = @data.slice!(/.*?\r\n/m)
+        response = process(data)
+      end
+
+      send_data response if response
+    end
+
+    def process(data)
+      data = @data_buf + data if @data_buf.size > 0
+      # our only non-normal state is consuming an object's data
+      # when @expected_length is present
+      if @expected_length && data.size == @expected_length
+        response = set_data(data)
+        @data_buf = ""
+        return response
+      elsif @expected_length
+        @data_buf = data
+        return
+      end
+      case data
+      when SET_COMMAND
+        @server.stats[:set_requests] += 1
+        set($1, $2, $3, $4.to_i)
+      when GET_COMMAND
+        @server.stats[:get_requests] += 1
+        get($1)
+      when STATS_COMMAND
+        stats
+      when SHUTDOWN_COMMAND
+        # no point in responding, they'll never get it.
+        Runner::shutdown
+      when DELETE_COMMAND
+        delete $1
+      else
+        logger.warn "Unknown command: #{data}."
+        respond ERR_UNKNOWN_COMMAND
+      end
+    rescue => e
+      logger.error "Error handling request: #{e}."
+      logger.debug e.backtrace.join("\n")
+      respond GET_RESPONSE_EMPTY
+    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 'starling/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