Refactor spike into classes

lib/pipeline/rpc/router.rb

@@ -1,15 +1,160 @@
 module Pipeline::Rpc
+
+  class WorkChannel
+
+    def initialize(zmq_context, queue_address)
+      @socket = zmq_context.socket(ZMQ::PUSH)
+      @socket.setsockopt(ZMQ::SNDHWM, 1)
+      @socket.bind(queue_address)
+
+      @poller = ZMQ::Poller.new
+      @poller.register(@socket, ZMQ::POLLOUT)
+    end
+
+    def worker_available?
+      poll_result = @poller.poll(500)
+      poll_result != -1 && @poller.writables.size > 0
+    end
+
+    def forward_to_backend(req, context=nil)
+      m = req.parsed_msg.clone
+      m[:context] = context unless context.nil?
+      upstream_msg = [req.raw_address, "", m.to_json]
+      @socket.send_strings(upstream_msg, ZMQ::DONTWAIT)
+    end
+
+  end
+
+  class RequestRegister
+
+    def initialize
+      @in_flight = {}
+    end
+
+    def register(req)
+      timeout_at = Time.now.to_i + 5
+      @in_flight[req.raw_address] = {timeout: timeout_at, req: req}
+    end
+
+    def forward_response(msg)
+      addr = msg.binary_return_address
+      entry = @in_flight[addr]
+      if entry.nil?
+        puts "dropping response"
+      else
+        req = entry[:req]
+        req.send_result(msg.parsed_msg)
+        unregister(addr)
+      end
+    end
+
+    def flush_expired_requests
+      timed_out = []
+      now = Time.now.to_i
+      @in_flight.each do |addr, entry|
+        expiry = entry[:timeout]
+        timed_out << entry[:req] if expiry < now
+      end
+      timed_out.each do |req|
+        req.send_error({status: :timeout})
+        puts "Timing out #{req}"
+        unregister(req.raw_address)
+      end
+    end
+
+    def unregister(addr)
+      @in_flight.delete(addr)
+    end
+
+  end
+
+  class ServiceResponse
+
+    def self.recv(socket)
+      msg = ""
+      socket.recv_string(msg)
+      self.new(msg, socket)
+    end
+
+    attr_reader :parsed_msg
+
+    def initialize(raw_msg, socket)
+      @raw_msg = raw_msg
+      @socket = socket
+      @parsed_msg = JSON.parse(raw_msg)
+    end
+
+    def type
+      @parsed_msg["msg_type"]
+    end
+
+    def return_address
+      @parsed_msg["return_address"]
+    end
+
+    def binary_return_address
+      return_address.pack("c*")
+    end
+
+    def raw_msg
+      @raw_msg
+    end
+
+  end
+
+  class FrontEndRequest
+
+    def self.recv(socket)
+      msg = []
+      socket.recv_strings(msg)
+      self.new(msg, socket)
+    end
+
+    attr_reader :raw_address, :raw_msg, :parsed_msg
+
+    def initialize(msg_strings, socket)
+      @raw_address = msg_strings[0]
+      @raw_msg = msg_strings[2]
+      @socket = socket
+    end
+
+    def send_error(err)
+      reply = [raw_address, "", err.to_json]
+      @socket.send_strings(reply)
+    end
+
+    def send_result(result)
+      reply = [raw_address, "", result.to_json]
+      @socket.send_strings(reply)
+    end
+
+    def handle
+      begin
+        @parsed_msg = JSON.parse(raw_msg)
+      rescue JSON::ParserError => e
+        req.send_error({ status: :parse_error })
+        return
+      end
+      action = @parsed_msg["action"]
+      if action.nil?
+        req.send_error({ status: :no_action })
+      else
+        yield(action)
+      end
+    end
+
+  end
+
   class Router
-    attr_reader :context, :front_end_socket,
+    attr_reader :zmq_context, :front_end_socket, :response_socket, :poller
-                :response_socket, :poller, :workers_poller
 
-    def initialize(context)
+    def initialize(zmq_context)
-      @context = context
+      @zmq_context = zmq_context
 
-      @front_end_socket = context.socket(ZMQ::ROUTER)
+      @front_end_socket = zmq_context.socket(ZMQ::ROUTER)
       @front_end_socket.bind('tcp://*:5566')
 
-      @in_flight = {}
+      @in_flight_requests = RequestRegister.new
 
       @public_hostname = "localhost"
       @response_port = 5555
@@ -20,7 +165,7 @@ module Pipeline::Rpc
         representers: 5579
       }
 
-      @response_socket = context.socket(ZMQ::SUB)
+      @response_socket = zmq_context.socket(ZMQ::SUB)
       @response_socket.setsockopt(ZMQ::SUBSCRIBE, "")
       @response_socket.bind("tcp://*:#{@response_port}")
 
@@ -31,24 +176,15 @@ module Pipeline::Rpc
       @backend_channels = {}
       @work_channel_ports.each do |type, port|
         bind_address = "tcp://*:#{@work_channel_ports[type]}"
-        channel = context.socket(ZMQ::PUSH)
+        work_channel = WorkChannel.new(zmq_context, bind_address)
-        channel.setsockopt(ZMQ::SNDHWM, 1)
+        @backend_channels[type] = work_channel
-        channel.bind(bind_address)
-
-        workers_poller = ZMQ::Poller.new
-        workers_poller.register(channel, ZMQ::POLLOUT)
-
-        @backend_channels[type] = {
-          socket: channel,
-          poller: workers_poller
-        }
       end
 
     end
 
     def run_heartbeater
       puts "STARTING heartbeat_socket"
-      heartbeat_socket = context.socket(ZMQ::PUB)
+      heartbeat_socket = zmq_context.socket(ZMQ::PUB)
       heartbeat_socket.connect("tcp://127.0.0.1:#{@response_port}")
       sleep 2
       loop do
@@ -80,47 +216,32 @@ module Pipeline::Rpc
     private
 
     def on_frontend_request
-      msg = []
+      req = FrontEndRequest.recv(front_end_socket)
-      front_end_socket.recv_strings(msg)
+      req.handle do |action|
-      raw_address = msg[0]
+        if action == "configure_worker"
-      raw_msg = msg[2]
+          respond_with_worker_config(req)
-      begin
+        elsif action == "analyze_iteration"
-        parsed_msg = JSON.parse(raw_msg)
+          handle_with_worker(:static_analyzers, req)
-      rescue JSON::ParserError => e
+        elsif action == "test_solution"
-        reply = [msg.first, "", { status: :parse_error }.to_json]
+          handle_with_worker(:test_runners, req)
-        front_end_socket.send_strings(reply)
+        elsif action == "represent"
-        return
+          handle_with_worker(:representers, req)
-      end
+        else
-      action = parsed_msg["action"]
+          req.send_error({ status: :unrecognised_action })
-      if action.nil?
+        end
-        reply = [msg.first, "", { status: :no_action }.to_json]
-        front_end_socket.send_strings(reply)
-        return
-      end
-      if action == "configure_worker"
-        respond_with_worker_config(raw_address, parsed_msg)
-      elsif action == "analyze_iteration"
-        handle_with_worker(:static_analyzers, parsed_msg, msg)
-      elsif action == "test_solution"
-        handle_with_worker(:test_runners, parsed_msg, msg)
-      elsif action == "represent"
-        handle_with_worker(:representers, parsed_msg, msg)
-      else
-        reply = [msg.first, "", { status: :unrecognised_action }.to_json]
-        front_end_socket.send_strings(reply)
       end
     end
 
-    def handle_with_worker(worker_class, parsed_msg, msg)
+    def handle_with_worker(worker_class, req)
       channel = @backend_channels[worker_class]
       if channel.nil?
-        reply = [msg.first, "", { status: :worker_class_unknown }.to_json]
+        req.send_error({ status: :worker_class_unknown })
-        front_end_socket.send_strings(reply)
+      elsif channel.worker_available?
-      elsif worker_available?(channel)
+        context = { credentials: temp_credentials }
-        forward_to_backend(channel, msg)
+        @in_flight_requests.register(req)
+        channel.forward_to_backend(req, context)
       else
-        reply = [msg.first, "", { status: :worker_unavailable }.to_json]
+        req.send_error({ status: :worker_unavailable })
-        front_end_socket.send_strings(reply)
       end
     end
 
@@ -133,37 +254,18 @@ module Pipeline::Rpc
     end
 
     def on_service_response
-      msg = ""
+      msg = ServiceResponse.recv(response_socket)
-      response_socket.recv_string(msg)
+      if msg.type == "response"
-      status_message = JSON.parse(msg)
+        @in_flight_requests.forward_response(msg)
-      type = status_message["msg_type"]
+        @in_flight_requests.unregister(msg.return_address)
-      if type == "response"
+      elsif msg.type == "heartbeat"
-        return_address = status_message["return_address"]
+        @in_flight_requests.flush_expired_requests
-        reply = [return_address.pack("c*"), "", msg]
-        front_end_socket.send_strings(reply, ZMQ::DONTWAIT)
-      elsif type == "heartbeat"
-        flush_expired_requests
         emit_current_spec
       else
         puts "Unrecognised message"
       end
     end
 
-    def flush_expired_requests
-      timed_out = []
-      now = Time.now.to_i
-      @in_flight.each do |addr, v|
-        expiry = v[:timeout]
-        timed_out << addr if expiry < now
-      end
-      timed_out.each do |addr|
-        reply = [addr, "", { status: :timeout }.to_json]
-        front_end_socket.send_strings(reply)
-        puts "Timing out #{@in_flight[addr]}"
-        @in_flight.delete(addr)
-      end
-    end
-
     def emit_current_spec
       analyzer_spec = analyzer_versions
       m = {
@@ -174,41 +276,27 @@ module Pipeline::Rpc
       message = ["_", "", m.to_json]
       puts "TODO"
       puts message
-      # back_end_socket.send_strings(message, ZMQ::DONTWAIT)
     end
 
-    def respond_with_worker_config(address, message)
+    def respond_with_worker_config(req)
       analyzer_spec = analyzer_versions
-      set_temp_credentials(analyzer_spec)
       analyzer_spec[:channels] = {
         workqueue_address: "tcp://#{@public_hostname}:#{@work_channel_ports[:static_analyzers]}",
         response_address: "tcp://#{@public_hostname}:#{@response_port}"
       }
-      reply = [address, "", analyzer_spec.to_json]
+      analyzer_spec["credentials"] = temp_credentials
-      front_end_socket.send_strings(reply)
+      req.send_result(analyzer_spec)
-    end
-
-    def worker_available?(channel)
-      poller = channel[:poller]
-      poll_result = poller.poll(500)
-      poll_result != -1 && poller.writables.size > 0
-    end
-
-    def forward_to_backend(channel, msg)
-      @in_flight[msg.first] = {msg: msg, timeout: Time.now.to_i + 5}
-      raw_msg = msg[2]
-      m = JSON.parse(raw_msg)
-      set_temp_credentials(m)
-      upstream_msg = [msg.first, "", m.to_json]
-      socket = channel[:socket]
-      socket.send_strings(upstream_msg, ZMQ::DONTWAIT)
     end
 
     def set_temp_credentials(msg)
+      msg["credentials"] = temp_credentials
+      msg
+    end
+
+    def temp_credentials
       sts =  Aws::STS::Client.new(region: "eu-west-1")
       session = sts.get_session_token(duration_seconds: 900)
-      msg["credentials"] = session.to_h[:credentials]
+      session.to_h[:credentials]
-      msg
     end
   end
 end

lib/pipeline/rpc/worker.rb

@@ -105,7 +105,7 @@ class Pipeline::Rpc::Worker
 
   def analyze(request)
     s3 = Aws::S3::Client.new(
-      credentials: parse_credentials(request),
+      credentials: parse_credentials(request["context"]),
       region: "eu-west-1")
 
     language_slug = request["track_slug"]