homeq 1.1.4

data/lib/homeq/sobs/connection.rb

@@ -0,0 +1,635 @@
+#############################################################################
+#
+# $Id: connection.rb 48 2008-11-19 05:11:59Z colin $
+#
+# Author:: Colin Steele (colin@colinsteele.org)
+# Homepage::  
+# 
+# TODO: info 
+#
+#----------------------------------------------------------------------------
+#
+#      Copyright (C) 2008 by Colin Steele. All Rights Reserved.
+#                        colin@colinsteele.org
+# 
+# This program is free software; you can redistribute it and/or modify
+# it under the terms of either: 1) the GNU General Public License
+# as published by the Free Software Foundation; either version 2 of the
+# License, or (at your option) any later version; or 2) Ruby's License.
+# 
+# See the file COPYING for complete licensing information.
+#
+#---------------------------------------------------------------------------
+#
+# 
+#############################################################################
+
+require 'socket'
+
+module HomeQ
+
+  module SOBS
+
+    class Connection < EventMachine::Connection
+
+      HEADER_REGEX                      = /([^\r\n]+?)\r\n/m
+      PUSH_TIMER_LENGTH                	= 1
+      HELLO_TIMER_LENGTH                = 5
+      OUTBOUND_MAX                      = 1024 * 64   # 64 KB
+      REFUSE_SEND_THRESHOLD             = 1024 * 1024 # 1 MB
+
+      # Shouldn't make any difference to throughput, but a low(ish)
+      # number should smooth out latency.
+      MAX_INPUT_MSGS_PER_READ_CYCLE     = 2
+
+      include Base::Logging
+      include Sender
+      include HomeQ
+
+      attr :server, true     # our server/listener
+      attr :peer, false
+      attr :state, false
+      attr :outbound_max, true
+      attr :refuse_send_threshold, true
+      
+      ######################################################################
+      # EventMachine::Connection entry points
+      ######################################################################
+
+      def post_init
+        @state = Statemachine.build do
+          state :start do
+            event :connect, :waiting_for_hello, :hello  # send hello
+            event :close, :closing
+            event :unbind, :closed
+          end
+          state :waiting_for_hello do
+            on_entry :wait_for_hello
+            on_exit :cancel_hello_timer
+            event :close, :closing
+            event :good_hello, :open
+            event :bad_hello, :closing
+            event :hello_timeout, :closing
+            event :unbind, :closed
+          end
+          state :open do
+            on_entry :opened
+            event :close, :closing
+            event :unbind, :closed
+          end
+          state :closing do
+            on_entry :start_closing
+            event :unbind, :closed
+            event :good_hello, :closing
+            event :bad_hello, :closing
+            event :hello_timeout, :closing
+            event :close, :closing
+          end
+          state :closed do
+            on_entry :closed
+            event :close, :closed
+          end
+        end
+        @state.context = self
+
+        @hello_timer    = nil
+        @data           = ""
+        @scanner        = StringScanner.new(@data)
+        @peer           = nil
+        @outgoing       = []
+
+        # Stats
+
+        @rcvs           = 0
+        @bytes_in       = 0
+        @bytes_out      = 0
+        @msgs_in        = 0
+        @msgs_out       = 0
+        @rcv_stats      = {}
+        @snd_stats      = {}
+        @opened_at      = nil
+        @closed_at      = nil
+        @outbound_bytes_highwater = 0
+
+        # Congestion
+        
+        @congested      = false
+        @refused_send   = 0
+        @outbound_max   ||= OUTBOUND_MAX
+        @refuse_send_threshold ||= REFUSE_SEND_THRESHOLD
+        
+        setup_peer
+        @state.connect if peer
+
+      end
+
+
+      # NB: EM does NOT call this for passive (server) connections
+      def connection_completed
+        @opened_at = Time.now
+      end
+      
+      # Called by EM when data is recieved from the peer
+      def receive_data data
+        return if (@state.state == :closed || @state.state == :closing)
+        logger.debug3 {
+          "Receive data callback."
+        }
+        @rcvs += 1
+        @bytes_in += data.length
+        @data << data
+        process_data
+      end
+
+      # Called by EM when the connection goes away.
+      def unbind
+        @closed_at = Time.now
+        if peer
+          logger.info {
+            "Connection has terminated to sobs peer #{remote_endpoint}."
+          }
+        end
+        if @opened_at
+          logger.info {
+            "Open for #{'%6.4fs' % (@closed_at - @opened_at)}. "+
+              "In: #{@msgs_in}. Out: #{@msgs_out}. "+
+              "IO/sec: #{io_per_sec}"
+          }
+        end
+        @data = ''  # drop anything we were doing for the peer.
+        @state.unbind
+        server.connections.delete(self) if server
+      end
+      
+      ######################################################################
+      # Message parsing routines
+      ######################################################################
+
+      # Do some parsing, and then, if needed, schedule a block to
+      # parse some more the next spin of the reactor.
+      def process_data
+        parse_incoming_data
+        logger.debug4 {
+          "check if callback needed? " +
+            "needed: #{@need_next_tick_parse.inspect} && " +
+            "ticker: #{@process_ticker.inspect}"
+        }
+        if @need_next_tick_parse && !@ticker
+          EventMachine::next_tick {
+            logger.debug3 {
+              "Next tick parse callback."
+            }
+            @process_ticker = false
+            process_data
+          }
+          @process_ticker = true
+          logger.debug3 {
+            "Scheduled tick parse callback."
+          }
+        end
+        @need_next_tick_parse = false
+      end
+
+      def parse_incoming_data
+        if outbound_maxed_out?
+          # No more.  Please.  Wait to gulp more down until we've spit
+          # some out.
+          logger.debug3 {
+            "Short circuit deferring parsing incoming data until next tick. " +
+              "'#{@data[0..20]}' #{@data.length}"
+          }
+          @need_next_tick_parse = true
+          return
+        end
+        msgs_this_spin = 0
+        logger.debug4 {
+          '<-- ' + @data.inspect
+        }
+        @scanner.pos = 0
+        while (header = @scanner.scan(HEADER_REGEX)) do
+          len = @data.length
+          parse_header(header)
+          if len == @data.length
+            break # it didn't take anything out, so there's not enough
+          end
+          
+          logger.debug3 {
+            "<<<< #{header.inspect}"
+          }
+
+          @msgs_in              += 1
+          msgs_this_spin        += 1
+
+          logger.debug3 {
+            "Msgs in: #{@msgs_in}"
+          }
+
+          # We've done enough work this turn of the reactor.  If we're
+          # closed/closing and we still have data we haven't handled,
+          # keep going.  Otherwise, go do some IO.
+          
+          if (msgs_this_spin == MAX_INPUT_MSGS_PER_READ_CYCLE &&
+              @data.length > 0 &&
+              (@state.state != :closed || @state.state != :closing))
+            logger.debug3 {
+              "Max input msgs per read cycle.  " +
+                "Deferring parsing incoming data until next tick. " +
+                "'#{@data[0..20]}' #{@data.class} #{@data.length}"
+            }
+            @scanner = StringScanner.new(@data)
+            @need_next_tick_parse = true
+            return
+          end
+          logger.debug4 {
+            "BUF NOW: '#{@data}'"
+          }
+          @scanner.pos = 0
+        end
+      end
+      
+      def parse_header(header_data)
+        elements = header_data.split(/\s/)
+        command_name = elements[0].downcase
+        if self.respond_to?("parse_#{command_name}")
+          self.send("parse_#{command_name}", elements[1..-1])
+        else
+          unknown_message(command_name)
+          @data.slice!(0, @scanner.pos)
+        end
+      end
+
+      NO_ARG_MESSAGES = [
+        'not_ignored',
+        'unknown_command',
+        'internal_error',
+        'out_of_memory',
+        'draining',
+        'bad_format',
+        'job_too_big',
+        'timed_out',
+        'reserve',
+        'subscribe',
+        'unsubscribe'
+      ]
+      NO_ARG_MESSAGES.each { |method_name|
+        class_eval do
+          define_method "parse_#{method_name}" do |args|
+            m = Message.new(method_name)
+            receive_message(m)
+            @data.slice!(0, @scanner.pos)
+          end
+        end
+      }
+
+
+      ARG_MESSAGES = [
+        'kicked',
+        'inserted',
+        'buried',
+        'hello',
+        'release',
+        'delete',
+        'bury',
+        'kick',
+        'deadline_soon',
+        'not_found',
+        'deleted',
+        'released',
+      ]
+      ARG_MESSAGES.each { |method_name|
+        class_eval do
+          define_method "parse_#{method_name}" do |args|
+            m = Message.new(method_name, nil, nil, args)
+            receive_message(m)
+            @data.slice!(0, @scanner.pos)
+          end
+        end
+      }
+
+      BODY_MESSAGES = ['ok']
+      BODY_MESSAGES.each { |method_name|
+        class_eval do
+          define_method "parse_#{method_name}" do |args|
+            body_length = args[0].to_i
+            # Add two bytes for delimiting "\r\n"
+            if @scanner.rest_size >= (body_length + 2)
+              m = Message.new(method_name,
+                              nil,
+                              @scanner.rest[0..(body_length - 1)])
+              receive_message(m)
+              @scanner.pos = @scanner.pos + body_length + 2
+              @data.slice!(0, @scanner.pos)
+            end
+          end
+        end
+      }
+
+      ARG_AND_BODY_MESSAGES = ['reserved']
+      ARG_AND_BODY_MESSAGES.each { |method_name|
+        class_eval do
+          define_method "parse_#{method_name}" do |args|
+            job_id = args[0]
+            body_length = args[1].to_i
+            # Add two bytes for delimiting "\r\n"
+            if @scanner.rest_size >= (body_length + 2)
+              m = Message.new(method_name,
+                              job_id,
+                              @scanner.rest[0..(body_length - 1)])
+              receive_message(m)
+              @scanner.pos = @scanner.pos + body_length + 2
+              @data.slice!(0, @scanner.pos)
+            end
+          end
+        end
+      }
+
+      ARGS_AND_BODY_MESSAGES = ['put']
+      ARGS_AND_BODY_MESSAGES.each { |method_name|
+        class_eval do
+          define_method "parse_#{method_name}" do |args|
+            body_length = args.last.to_i
+            # Add two bytes for delimiting "\r\n"
+            if @scanner.rest_size >= (body_length + 2)
+              pos = @scanner.pos
+              endpos = pos + (body_length - 1)
+              m = Message.new(method_name,
+                              nil,
+                              @data[pos..endpos], 
+                              args)
+              receive_message(m)
+              @data.slice!(0, endpos + 3) # endpos is an index; plus "\r\n"
+            end
+          end
+        end
+      }
+
+      ALL_MESSAGES = 
+        NO_ARG_MESSAGES + ARG_MESSAGES +
+        BODY_MESSAGES + ARG_AND_BODY_MESSAGES + ARGS_AND_BODY_MESSAGES
+
+      def receive_message(message)
+        @rcv_stats[message.type] ||= 0
+        @rcv_stats[message.type] += 1
+        if self.respond_to?("receive_#{message.type}")
+          self.send("receive_#{message.type}", message)
+        else
+          logger.error {
+            "Ignoring message: #{message.type}"
+          }
+        end
+      end
+
+      ######################################################################
+      # Message interface - receiving
+      ######################################################################
+
+      ALL_MESSAGES.each{ |method_name|
+        class_eval do
+          define_method "receive_#{method_name}" do |message|
+            if !ok_to_receive?
+              return
+            end
+            logger.debug4 {
+              "Got message #{message.type}"
+            }
+          end
+        end
+      }
+
+      def receive_hello(message)
+        if !protocol_version_ok?(message.args[0].to_i)
+          logger.info {
+            "SOBS Peer #{remote_endpoint} sent incompatible " +
+              "protocol version #{message.args[0].to_i}"
+          }
+          @state.bad_hello
+        else
+          @state.good_hello
+        end
+      end
+
+      ######################################################################
+      # Other
+      ######################################################################
+
+      def to_s
+        str = ''
+        str << "EM signature: #{signature} peer: #{remote_endpoint} \n"
+        str << "state: #{state.state} "
+        str << "opened: #{@opened_at} (#{Time.now - @opened_at}s)\n"
+        str << "rcvs: #{@rcvs} in: #{@msgs_in} out: #{@msgs_out}\n"
+        str << "bytes in: #{@bytes_in} bytes_out: #{@bytes_out}\n"
+        str << "outgoing buffers: #{@outgoing.length}, "
+        str << "#{queued_outgoing_data_size} total bytes\n"
+        str << "#{self.get_outbound_data_size} total outbound bytes\n"
+        str << "send ticker is #{@send_ticker ? '' : 'not '}running\n"
+        str << "outbound_max: #{@outbound_max || OUTBOUND_MAX} "
+        str << "refuse_send_threshold: "
+        str << "#{@refuse_send_threshold || REFUSE_SEND_THRESHOLD}\n"
+        str << "congested?: #{@congested ? 'Y' : 'N'}, "
+        str << "refused sends: #{@refused_send}\n"
+        str << "outbound_bytes_highwater: #{@outbound_bytes_highwater}\n"
+        @rcv_stats.each {|msg_name,count|
+          str << "#{msg_name} in: #{count}\n"
+        }
+        @snd_stats.each {|msg_name,count|
+          str << "#{msg_name} out: #{count}\n"
+        }
+        str
+      end
+
+      def remote_endpoint
+        peer ? "#{peer[1]}:#{peer[0]}" : '?:?'
+      end
+
+      def should_refuse_send?
+        outbound_size = queued_outgoing_data_size
+        @outbound_bytes_highwater = [
+          @outbound_bytes_highwater,
+          outbound_size
+        ].max
+        outbound_size > (@refuse_send_threshold || REFUSE_SEND_THRESHOLD)
+      end
+
+      def enter_congestion
+        unless @congested
+          logger.warn {
+            "Connection #{remote_endpoint} congested.  Peer having problems?"
+          }
+          @congested = true
+          @congested_at = Time.now
+        end
+      end
+
+      def exit_congestion
+        @congested = false
+        logger.warn {
+          "Congested #{remote_endpoint} cleared.  Was congested for " +
+            "#{Time.now - @congested_at} seconds."
+        }
+      end
+      
+      def send_to_peer(data)
+        if should_refuse_send?
+          @refused_send += 1
+          enter_congestion unless @congested
+          return false
+        end
+        @outgoing.push(data)
+        push_to_peer
+        true # could return queued_outgoing_data_size?
+      end
+
+      def unknown_message(command_name)
+        msg = "Unknown/unsupported command: #{command_name}"
+        sys.die(msg)
+      end
+      
+      def setup_peer
+        if (peername = get_peername)
+          @peer = Socket.unpack_sockaddr_in(peername)
+          @opened_at = Time.now
+          logger.info {
+            "SOBS connection to #{remote_endpoint} complete."
+          }
+        end
+      end
+
+      def protocol_version_ok?(their_version)
+        their_version == HomeQ::SOBS::PROTOCOL_VERSION       
+      end
+
+      def ok_to_receive?
+        if (@state.state == :start ||
+            @state.state == :waiting_for_hello ||
+            @state.state == :closed)
+          logger.error {
+            "SOBS peer #{remote_endpoint} sent isn't open."
+          }
+          close
+          false
+        else
+          true
+        end
+      end
+
+      def close
+        @state.close
+      end
+
+      def wait_for_hello
+        @hello_timer = EventMachine::Timer.new(HELLO_TIMER_LENGTH) do
+          logger.info {
+            "No hello received from #{remote_endpoint}.  Bye."
+          }
+          @hello_timer = nil
+          @state.hello_timeout
+        end
+      end
+
+      def cancel_hello_timer
+        @hello_timer.cancel if @hello_timer
+        @hello_timer = nil
+      end
+
+      def start_closing
+        close_connection_after_writing
+        closing
+      end
+
+      ######################################################################
+      # Higher-level interface for subclasses
+      ######################################################################
+
+      def opened
+      end
+      def closing
+      end
+      def closed
+      end
+
+      ######################################################################
+      # Workhorse methods
+      ######################################################################
+
+      protected
+
+      def queued_outgoing_data_size
+        @outgoing.inject(0) { |sum,data| sum += data.length }
+      end
+
+      def outbound_maxed_out?
+        self.get_outbound_data_size >= @outbound_max
+      end
+
+      def push_to_peer
+        logger.debug2 {
+          "push #{self.get_outbound_data_size}"
+        }
+
+        while !outbound_maxed_out? do
+          data = @outgoing.shift
+          unless data
+            logger.debug2 {
+              "No more to send."
+            }
+            exit_congestion if @congested
+            return
+          end
+          logger.debug4 {
+            '--> ' + data.inspect
+          }
+          logger.debug3 {
+            '>>>> ' + data[0..20].chomp
+          }
+          @msgs_out += 1
+          @bytes_out += data.length
+          send_data(data)
+        end
+
+        exit_congestion if (!should_refuse_send? && @congested)
+
+        logger.debug4 {
+          "push maxed #{outbound_maxed_out?}; " +
+          "outbound #{self.get_outbound_data_size}; " +
+          "queued #{self.queued_outgoing_data_size}; " +
+          "ticker #{@send_ticker}"
+        }
+
+        raise 'WTF' unless outbound_maxed_out?
+        
+        if @outgoing.any? && !@send_ticker
+          logger.debug2 {
+            "Send buffer size exceeded (now " +
+            "#{self.get_outbound_data_size}).  Deferring sends."
+          }
+
+          @send_ticker = true
+          @push_timer = EventMachine::Timer.new(PUSH_TIMER_LENGTH) do
+            logger.debug3 {
+              "Push to peer callback."
+            }
+            @send_ticker = false
+            @push_timer  = nil
+            push_to_peer
+          end
+
+          logger.debug3 {
+            "Scheduled push callback."
+          }
+        end
+      end
+
+      def io_per_sec
+        '%6.4f.' % ((@msgs_in + @msgs_out) / how_long_open).to_f
+      end
+
+      def how_long_open
+        (@closed_at || Time.now) - (@opened_at || Time.now)
+      end
+
+    end # Connection
+
+  end # module SOBS
+
+end # module HomeQ