em-zmq-tp10 0.1.7

data/lib/em/protocols/zmq2/router.rb

@@ -0,0 +1,69 @@
+require 'em/protocols/zmq2/socket'
+require 'em/protocols/zmq2/queue_per_peer'
+module EventMachine
+  module Protocols
+    module Zmq2
+      # ZMQ socket which acts like Router but without outgoing message queueing.
+      # It counts first message string as peer identity when sending message and
+      # prepends socket identity to message on receiving.
+      class PreRouter < Socket
+        def receive_message_and_peer(message, peer_identity) # :nodoc:
+          message.unshift(peer_identity)
+          receive_message(message)
+        end
+
+        def receive_message(message)
+          raise NoMethodError
+        end
+
+        def send_message(message, even_if_busy = false)
+          if connect = choose_peer(message.first, even_if_busy)
+            connect.send_strings(message[1..-1])
+            true
+          end
+        end
+
+      private
+        # by default chooses peer by peer_identity, but you could wary it
+        def choose_peer(peer_identity, even_if_busy = false)
+          if (connect = @peers[peer_identity]) && !connect.error? &&
+             (even_if_busy || connect.not_too_busy?)
+            connect
+          end
+        end
+
+      end
+
+      # ZMQ socket which acts like Router.
+      # It counts first message string as peer identity when sending message and
+      # prepends socket identity to message on receiving.
+      class Router < PreRouter
+        include QueuePerPeer
+        def send_message(message)
+          peer_identity = message.first
+          unless (queue = @queues[peer_identity])
+            if generated_identity?(peer_identity)
+              return false
+            else
+              queue = @queues[peer_identity] = []
+            end
+          end
+          peer = choose_peer(peer_identity)
+          if peer && (queue.empty? || flush_queue(queue, peer)) &&
+             !peer.error? && peer.not_too_busy?
+            peer.send_strings(message[1..-1])
+            true
+          else
+            push_to_queue(queue, message)
+          end
+        end
+
+      private
+        def send_formed_message(peer, from_queue)
+          peer.send_strings(from_queue[1..-1])
+        end
+      end
+
+    end
+  end
+end

data/lib/em/protocols/zmq2/socket.rb

@@ -0,0 +1,236 @@
+require 'em/protocols/zmq2/socket_connection'
+require 'em/protocols/zmq2/inproc'
+module EM
+  module Protocols
+    module Zmq2
+      # Base class for all ZMQ sockets.
+      # It implements address parsing, binding, connecting and reconnecting
+      # For implementing your own kind of socket you should override at least
+      # #receive_message, and define method which will choose peer from @peers,
+      # and call its peer#send_strings
+      class Socket
+        attr :identity, :hwm
+        attr_accessor :hwm_strategy
+        GENERATED = '%GN%'.freeze
+
+        # Accept options, which are dependend on socket type
+        # Common options are:
+        #  [+:identity+] socket identity
+        #  [+:hwm+]      highwater mark
+        #  [+:hwm_strategy+]  what to do on send_message when :hwm reached, hwm_strategy could be:
+        #                     +:drop_last+ - do not accept message for sending - what Zmq does
+        #                     +:drop_first+ - remove message from queue head and add this message to queue tail
+        #
+        # this class provides convinient method +#push_to_queue+ for default strategy, but
+        # it is up to subclass how to use it.
+        #
+        # Another note concerning highwatermark: EventMachine does not allow precise control on
+        # outgoing data buffer, so that there is a bit more message will be lost, when outgoing
+        # peer disconnected.
+        def initialize(opts = {})
+          @hwm = opts[:hwm] || HWM_INFINITY
+          @hwm_strategy = opts[:hwm_strategy] || :drop_last
+          @identity = opts[:identity] || EMPTY
+          @peers = {}
+          @free_peers = {}
+          @connections = {}
+          @conn_addresses = {}
+          @reconnect_timers = {}
+          @bindings = []
+          @after_writting = nil
+          @uniq_identity = '%GN%aaaaaaaaaaaa' # ~ 100 years to overflow
+        end
+
+        # binding to port
+        # :call-seq:
+        #   bind('tcp://host:port') - bind to tcp port
+        #   bind('ipc://filename') - bind to unix port
+        def bind(addr)
+          kind, *socket = parse_address(addr)
+          EM.schedule {
+            @bindings << case kind
+                when :tcp, :ipc
+                  EM.start_server(*socket, SocketConnection, self)
+                when :inproc
+                  InProc.bind(addr, self)
+                end
+          }
+        end
+
+        # connect to port
+        # :call-seq:
+        #   connect('tcp://host:port') - connect to tcp port
+        #   connect('ipc://filename') - connect to unix port
+        def connect(addr)
+          kind, *socket = parse_address(addr)
+          EM.schedule lambda{
+            @reconnect_timers.delete addr
+            unless @conn_addresses[ addr ]
+              connection = case kind
+                  when :tcp
+                    EM.connect(*socket, SocketConnection, self)
+                  when :ipc
+                    begin
+                      EM.connect_unix_domain(*socket, SocketConnection, self)
+                    rescue RuntimeError
+                      timer = EM.add_timer(SMALL_TIMEOUT) do
+                        connect(addr)
+                      end
+                      @reconnect_timers[addr] = timer
+                      break
+                    end
+                  when :inproc
+                    InProc.connect(addr, self)
+                  end
+              @connections[ connection ] = addr
+              @conn_addresses[ addr ] = connection
+            end
+          }
+        end
+
+        # :stopdoc:
+        def not_connected(connection)
+          if addr = @connections.delete(connection)
+            @conn_addresses.delete addr
+            timer = EM.add_timer(SMALL_TIMEOUT) do
+              connect(addr)
+            end
+            @reconnect_timers[addr] = timer
+          end
+        end
+
+        def register_peer(peer_identity, connection)
+          peer_identity = next_uniq_identity  if peer_identity.empty?
+          @peers[peer_identity] = connection
+          EM.next_tick{ peer_free(peer_identity, connection) }
+          peer_identity
+        end
+
+        def unregister_peer(peer_identity)
+          @peers.delete peer_identity
+          @free_peers.delete peer_identity
+          if @peers.empty? && @after_writting.respond_to?(:call)
+            @after_writting.call
+          end
+        end
+
+        # :startdoc:
+
+        # close all connections
+        # if callback is passed, then it will be called after all messages written to sockets
+        def close(cb = nil, &block)
+          @connections.clear
+          @conn_addresses.clear
+          @reconnect_timers.each{|_, timer| EM.cancel_timer(timer)}
+          @after_writting = cb || block
+          flush_all_queue  if @after_writting
+          @peers.values.each{|c| c.close_connection(!!@after_writting)}
+          @bindings.each{|c|
+            unless String === c
+              EM.stop_server c
+            else
+              InProc.unbind(c, self)
+            end
+          }
+        end
+
+        # override to make sure all messages are sent before socket closed
+        def flush_all_queue
+          true
+        end
+
+        # stub method for sending message to a socket
+        # note that every socket type should define proper behaviour here
+        # or/and define another useful, semantic clear methods
+        def send_message(message)
+          raise NoMethodError
+        end
+
+        # callback method called with underlied connection when
+        # some message arrives
+        def receive_message_and_peer(message, peer_identity)
+          raise NoMethodError
+        end
+
+        # Change hwm
+        def hwm=(new_hwm)
+          old_hwm, @hwm = @hwm, new_hwm
+          react_on_hwm_decrease if old_hwm > @hwm
+          @hwm
+        end
+
+        # callback method, called when underlying peer is free for writing in
+        # should be used in subclasses for proper reaction on network instability
+        def peer_free(peer_identity, peer) # :doc:
+          @free_peers[peer_identity] = peer
+        end
+
+      private
+        # splits message into envelope and message as defined by ZMQ 2.x
+        def split_message(message) # :doc:
+          i = message.index(EMPTY)
+          [message.slice(0, i), message.slice(i+1, message.size)]
+        end
+
+        # helper method for managing queue concerning @hwm setting
+        def push_to_queue(queue, message = nil) # :doc
+          if queue.size >= @hwm
+            case @hwm_strategy
+            when :drop_last
+              if queue.size > @hwm
+                queue.pop(queue.size - @hwm).each{|message|
+                  cancel_message(message)
+                }
+              end
+              false
+            when :drop_first
+              hwm = @hwm - (message ? 1 : 0)
+              queue.shift(queue.size - hwm).each{|earlier|
+                cancel_message(earlier)
+              }
+              queue.push(message.dup)  if message
+              true
+            end
+          else
+            queue.push(message.dup) if message
+            true
+          end
+        end
+
+        # override to correctly react on hwm decrease
+        def react_on_hwm_decrease # :doc:
+          true
+        end
+
+        # overried if you should react on dropped requests
+        def cancel_message(message) # :doc:
+          true
+        end
+
+        def next_uniq_identity
+          res = @uniq_identity
+          @uniq_identity = res.next
+          res
+        end
+
+        def generated_identity?(id)
+          id.start_with?(GENERATED)
+        end
+
+        def parse_address(addr)
+          case addr
+          when %r{tcp://([^:]+):(\d+)}
+            [:tcp, $1, $2.to_i]
+          when %r{ipc://(.+)}
+            [:ipc, $1]
+          when %r{inproc://(.+)}
+            [:inproc, $1]
+          else
+            raise 'Not supported ZMQ socket kind'
+          end
+        end
+
+      end
+    end
+  end
+end

data/lib/em/protocols/zmq2/socket_connection.rb

@@ -0,0 +1,151 @@
+require 'eventmachine'
+require 'em/protocols/zmq2'
+require 'em/protocols/zmq2/connection'
+module EventMachine
+  module Protocols
+    module Zmq2
+      # Main implementation peace, heart of protocol
+      module PackString
+        BIG_PACK = 'CNNCa*'.freeze
+        SMALL_PACK = 'CCa*'.freeze
+        def pack_string(string, more = false)
+          bytesize = string.bytesize + 1
+          if bytesize <= 254
+            [bytesize, more ? 1 : 0, string].pack(SMALL_PACK)
+          else
+            [255, 0, bytesize, more ? 1 : 0, string].pack(BIG_PACK)
+          end
+        end
+
+        def prepare_message(message)
+          if String === message
+            pack_string(message, false)
+          else
+            message = Array(message)
+            buffer = ''
+            i = 0
+            last = message.size - 1
+            while i < last
+              buffer << pack_string(message[i].to_s, true)
+              i += 1
+            end
+            buffer << pack_string(message[last].to_s, false)
+          end
+        end
+      end
+
+      # Heavy duty worker class
+      # Implements ZMTP1.0 - ZMQ2.x transport protocol
+      #
+      # It calls following callbacks
+      #
+      # It is not for end user usage
+      class SocketConnection < EM::Connection
+        include ConnectionMixin
+        # :stopdoc:
+        def initialize(socket)
+          @socket = socket
+          @recv_buffer = ''
+          @recv_frames = [[]]
+        end
+
+        def unbind(err)
+          if @peer_identity
+            @socket.unregister_peer(@peer_identity)
+          end
+          @socket.not_connected(self)
+        end
+
+        # use watching on outbound queue when possible
+        # rely on https://github.com/eventmachine/eventmachine/pull/317 if were accepted
+        # or on https://github.com/funny-falcon/eventmachine/tree/sent_data
+        # use timers otherwise
+        def sent_data
+          @socket.peer_free(@peer_identity, self) if not_too_busy?
+        end
+
+        if method_defined?(:outbound_data_count)
+          def _not_too_busy?
+            outbound_data_count < 32
+          end
+        else
+          def _not_too_busy?
+            get_outbound_data_size < 2048
+          end
+        end
+
+        if method_defined?(:notify_sent_data=)
+          alias notify_when_free= notify_sent_data=
+        else
+          def notify_when_free=(v)
+            if v
+              @when_free_timer ||= EM.add_timer(SMALL_TIMEOUT) do
+                @when_free_timer = nil
+                sent_data
+              end
+            elsif @when_free_timer
+              EM.cancel_timer @when_free_timer
+              @when_free_timer = nil
+            end
+          end
+        end
+
+        def not_too_busy?
+          free = _not_too_busy?
+          self.notify_when_free = !free
+          free
+        end
+
+        def receive_data(data)
+          parse_frames(data)
+          while message = pop_message
+            receive_strings(message)
+          end
+        end
+
+        def pop_message
+          if @recv_frames.size > 1
+            @recv_frames.shift
+          end
+        end
+
+        FF = "\xff".freeze
+        BIG_UNPACK = 'CNNC'.freeze
+        SMALL_UNPACK = 'CC'.freeze
+        def parse_frames(data)
+          unless @recv_buffer.empty?
+            data = @recv_buffer << data
+          end
+          while data.bytesize >= 2
+            if data.start_with?(FF)
+              break  if data.bytesize < 10
+              _, _, length, more = data.unpack(BIG_UNPACK)
+              start_at = 10
+            else
+              length = data.getbyte(0)
+              more = data.getbyte(1)
+              start_at = 2
+            end
+            length -= 1
+            break  if data.bytesize < start_at + length
+            str = data.byteslice(start_at, length)
+            data[0, start_at + length] = EMPTY
+            @recv_frames.last << str
+            @recv_frames << [] if more & 1 == 0
+          end
+          @recv_buffer = data
+        end
+
+        include PackString
+
+        def send_strings(strings)
+          send_data prepare_message(strings)
+        end
+
+        def send_strings_or_prepared(strings, prepared)
+          send_data prepared
+        end
+      end
+    end
+  end
+end