protobuffy 3.1.0

data/lib/protobuf/rpc/servers/zmq/server.rb

@@ -0,0 +1,313 @@
+require 'protobuf/rpc/servers/zmq/util'
+require 'protobuf/rpc/servers/zmq/worker'
+require 'protobuf/rpc/servers/zmq/broker'
+require 'protobuf/rpc/dynamic_discovery.pb'
+require 'securerandom'
+require 'thread'
+
+module Protobuf
+  module Rpc
+    module Zmq
+      class Server
+        include ::Protobuf::Rpc::Zmq::Util
+
+        DEFAULT_OPTIONS = {
+          :beacon_interval => 5,
+          :broadcast_beacons => false,
+          :broadcast_busy => false,
+          :zmq_inproc => true,
+        }
+
+        attr_accessor :options, :workers
+        attr_reader :zmq_context
+
+        def initialize(options)
+          @options = DEFAULT_OPTIONS.merge(options)
+          @workers = []
+
+          init_zmq_context
+          init_beacon_socket if broadcast_beacons?
+          init_shutdown_pipe
+        rescue
+          teardown
+          raise
+        end
+
+        def add_worker
+          @total_workers = total_workers + 1
+        end
+
+        def all_workers_busy?
+          workers.all? { |thread| !!thread[:busy] }
+        end
+
+        def backend_port
+          options[:worker_port] || frontend_port + 1
+        end
+
+        def backend_uri
+          if inproc?
+            "inproc://#{backend_ip}:#{backend_port}"
+          else
+            "tcp://#{backend_ip}:#{backend_port}"
+          end
+        end
+
+        def beacon_interval
+          [options[:beacon_interval].to_i, 1].max
+        end
+
+        def beacon_ip
+          "255.255.255.255"
+        end
+
+        def beacon_port
+          unless @beacon_port
+            unless port = options[:beacon_port]
+              port = ::Protobuf::Rpc::ServiceDirectory.port
+            end
+
+            @beacon_port = port.to_i
+          end
+
+          @beacon_port
+        end
+
+        def beacon_uri
+          "udp://#{beacon_ip}:#{beacon_port}"
+        end
+
+        def broadcast_beacons?
+          !brokerless? && options[:broadcast_beacons]
+        end
+
+        def broadcast_flatline
+          flatline = ::Protobuf::Rpc::DynamicDiscovery::Beacon.new(
+            :beacon_type => ::Protobuf::Rpc::DynamicDiscovery::BeaconType::FLATLINE,
+            :server => self.to_proto
+          )
+
+          @beacon_socket.send(flatline.encode, 0)
+        end
+
+        def broadcast_heartbeat
+          @last_beacon = Time.now.to_i
+
+          heartbeat = ::Protobuf::Rpc::DynamicDiscovery::Beacon.new(
+            :beacon_type => ::Protobuf::Rpc::DynamicDiscovery::BeaconType::HEARTBEAT,
+            :server => self.to_proto
+          )
+
+          @beacon_socket.send(heartbeat.encode, 0)
+
+          log_debug { sign_message("sent heartbeat to #{beacon_uri}") }
+        end
+
+        def broadcast_heartbeat?
+          Time.now.to_i >= next_beacon && broadcast_beacons?
+        end
+
+        def brokerless?
+          !!options[:workers_only]
+        end
+
+        def busy_worker_count
+          workers.count { |thread| !!thread[:busy] }
+        end
+
+        def frontend_ip
+          @frontend_ip ||= resolve_ip(options[:host])
+        end
+        alias_method :backend_ip, :frontend_ip
+
+        def frontend_port
+          options[:port]
+        end
+
+        def frontend_uri
+          "tcp://#{frontend_ip}:#{frontend_port}"
+        end
+
+        def inproc?
+          !!self.options[:zmq_inproc]
+        end
+
+        def maintenance_timeout
+          next_maintenance - Time.now.to_i
+        end
+
+        def next_maintenance
+          cycles = [next_reaping]
+          cycles << next_beacon if broadcast_beacons?
+
+          cycles.min
+        end
+
+        def minimum_timeout
+          0.1
+        end
+
+        def next_beacon
+          if @last_beacon.nil?
+            0
+          else
+            @last_beacon + beacon_interval
+          end
+        end
+
+        def next_reaping
+          if @last_reaping.nil?
+            0
+          else
+            @last_reaping + reaping_interval
+          end
+        end
+
+        def reap_dead_workers
+          @last_reaping = Time.now.to_i
+
+          @workers.keep_if do |worker|
+            worker.alive? or worker.join && false
+          end
+        end
+
+        def reap_dead_workers?
+          Time.now.to_i >= next_reaping
+        end
+
+        def reaping_interval
+          5
+        end
+
+        def run
+          @running = true
+
+          start_broker unless brokerless?
+          start_missing_workers
+
+          yield if block_given? # runs on startup
+          wait_for_shutdown_signal
+          broadcast_flatline if broadcast_beacons?
+          Thread.pass until reap_dead_workers.empty?
+          @broker.join unless brokerless?
+        ensure
+          @running = false
+          teardown
+        end
+
+        def running?
+          !!@running
+        end
+
+        def start_missing_workers
+          missing_workers = total_workers - @workers.size
+
+          if missing_workers > 0
+            missing_workers.times { start_worker }
+            log_debug { sign_message("#{total_workers} workers started") }
+          end
+        end
+
+        def stop
+          @running = false
+          @shutdown_w.write('.')
+        end
+
+        def teardown
+          @shutdown_r.try(:close)
+          @shutdown_w.try(:close)
+          @beacon_socket.try(:close)
+          @zmq_context.try(:terminate)
+          @last_reaping = @last_beacon = @timeout = nil
+        end
+
+        def total_workers
+          @total_workers ||= [@options[:threads].to_i, 1].max
+        end
+
+        def timeout
+          if @timeout.nil?
+            @timeout = 0
+          else
+            @timeout = [minimum_timeout, maintenance_timeout].max
+          end
+        end
+
+        def to_proto
+          @proto ||= ::Protobuf::Rpc::DynamicDiscovery::Server.new(
+            :uuid => uuid,
+            :address => frontend_ip,
+            :port => frontend_port.to_s,
+            :ttl => (beacon_interval * 1.5).ceil,
+            :services => ::Protobuf::Rpc::Service.implemented_services
+          )
+        end
+
+        def uuid
+          @uuid ||= SecureRandom.uuid
+        end
+
+        def wait_for_shutdown_signal
+          loop do
+            break if IO.select([@shutdown_r], nil, nil, timeout)
+
+            if reap_dead_workers?
+              reap_dead_workers
+              start_missing_workers
+            end
+
+            if broadcast_heartbeat?
+              if all_workers_busy? && options[:broadcast_busy]
+                broadcast_flatline
+              else
+                broadcast_heartbeat
+              end
+            end
+
+          end
+        end
+
+        private
+
+        def init_beacon_socket
+          @beacon_socket = UDPSocket.new
+          @beacon_socket.setsockopt(::Socket::SOL_SOCKET, ::Socket::SO_BROADCAST, true)
+          @beacon_socket.setsockopt(::Socket::SOL_SOCKET, ::Socket::SO_REUSEADDR, true)
+
+          if defined?(::Socket::SO_REUSEPORT)
+            @beacon_socket.setsockopt(::Socket::SOL_SOCKET, ::Socket::SO_REUSEPORT, true)
+          end
+
+          @beacon_socket.bind(frontend_ip, beacon_port)
+          @beacon_socket.connect(beacon_ip, beacon_port)
+        end
+
+        def init_shutdown_pipe
+          @shutdown_r, @shutdown_w = IO.pipe
+        end
+
+        def init_zmq_context
+          @zmq_context = ZMQ::Context.new
+        end
+
+        def start_broker
+          @broker = Thread.new(self) do |server|
+            ::Protobuf::Rpc::Zmq::Broker.new(server).run
+          end
+        end
+
+        def start_worker
+          @workers << Thread.new(self) do |server|
+            begin
+              ::Protobuf::Rpc::Zmq::Worker.new(server).run
+            rescue => e
+              message = "Worker failed: #{e.inspect}\n #{e.backtrace.join($/)}"
+              $stderr.puts(message)
+              log_error { message }
+            end
+          end
+        end
+      end
+    end
+  end
+end

data/lib/protobuf/rpc/servers/zmq/util.rb

@@ -0,0 +1,47 @@
+require 'resolv'
+
+module Protobuf
+  module Rpc
+    module Zmq
+
+      ADDRESS_MATCH = /\A\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\z/.freeze
+      WORKER_READY_MESSAGE = "\1"
+      CHECK_AVAILABLE_MESSAGE = "\3"
+      NO_WORKERS_AVAILABLE = "\4"
+      WORKERS_AVAILABLE = "\5"
+
+      module Util
+        include ::Protobuf::Logger::LogMethods
+
+        def self.included(base)
+          base.extend(::Protobuf::Rpc::Zmq::Util)
+        end
+
+        def zmq_error_check(return_code, source = nil)
+          unless ::ZMQ::Util.resultcode_ok?(return_code)
+            raise <<-ERROR
+            Last ZMQ API call #{source ? "to #{source}" : ""} failed with "#{::ZMQ::Util.error_string}".
+
+            #{caller(1).join($/)}
+            ERROR
+          end
+        end
+
+        def log_signature
+          unless @_log_signature
+            name = (self.class == Class ? self.name : self.class.name)
+            @_log_signature = "[server-#{name}-#{object_id}]"
+          end
+
+          @_log_signature
+        end
+
+        def resolve_ip(hostname)
+          ::Resolv.getaddresses(hostname).detect do |address|
+            address =~ ADDRESS_MATCH
+          end
+        end
+      end
+    end
+  end
+end

data/lib/protobuf/rpc/servers/zmq/worker.rb

@@ -0,0 +1,105 @@
+require 'protobuf/rpc/server'
+require 'protobuf/rpc/servers/zmq/util'
+require 'thread'
+
+module Protobuf
+  module Rpc
+    module Zmq
+      class Worker
+        include ::Protobuf::Rpc::Server
+        include ::Protobuf::Rpc::Zmq::Util
+
+        ##
+        # Constructor
+        #
+        def initialize(server)
+          @server = server
+
+          init_zmq_context
+          init_backend_socket
+        rescue
+          teardown
+          raise
+        end
+
+        ##
+        # Instance Methods
+        #
+        def process_request
+          client_address, _, data = read_from_backend
+          return unless data
+
+          gc_pause do
+            encoded_response = handle_request(data)
+            write_to_backend([client_address, "", encoded_response])
+          end
+        end
+
+        def run
+          poller = ::ZMQ::Poller.new
+          poller.register_readable(@backend_socket)
+          poller.register_readable(@shutdown_socket)
+
+          # Send request to broker telling it we are ready
+          write_to_backend([::Protobuf::Rpc::Zmq::WORKER_READY_MESSAGE])
+
+          loop do
+            rc = poller.poll(500)
+
+            # The server was shutdown and no requests are pending
+            break if rc == 0 && !running?
+
+            # Something went wrong
+            break if rc == -1
+
+            if rc > 0
+              ::Thread.current[:busy] = true
+              process_request
+              ::Thread.current[:busy] = false
+            end
+          end
+        ensure
+          teardown
+        end
+
+        def running?
+          @server.running?
+        end
+
+        private
+
+        def init_zmq_context
+          if inproc?
+            @zmq_context = @server.zmq_context
+          else
+            @zmq_context = ZMQ::Context.new
+          end
+        end
+
+        def init_backend_socket
+          @backend_socket = @zmq_context.socket(ZMQ::REQ)
+          zmq_error_check(@backend_socket.connect(@server.backend_uri))
+        end
+
+        def inproc?
+          !!@server.try(:inproc?)
+        end
+
+        def read_from_backend
+          frames = []
+          zmq_error_check(@backend_socket.recv_strings(frames))
+          frames
+        end
+
+        def teardown
+          @backend_socket.try(:close)
+          @zmq_context.try(:terminate)
+        end
+
+        def write_to_backend(frames)
+          zmq_error_check(@backend_socket.send_strings(frames))
+        end
+      end
+    end
+  end
+end