Refactor spike into classes
This commit is contained in:
@@ -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,
|
||||
:response_socket, :poller, :workers_poller
|
||||
attr_reader :zmq_context, :front_end_socket, :response_socket, :poller
|
||||
|
||||
def initialize(context)
|
||||
@context = context
|
||||
def initialize(zmq_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)
|
||||
channel.setsockopt(ZMQ::SNDHWM, 1)
|
||||
channel.bind(bind_address)
|
||||
|
||||
workers_poller = ZMQ::Poller.new
|
||||
workers_poller.register(channel, ZMQ::POLLOUT)
|
||||
|
||||
@backend_channels[type] = {
|
||||
socket: channel,
|
||||
poller: workers_poller
|
||||
}
|
||||
work_channel = WorkChannel.new(zmq_context, bind_address)
|
||||
@backend_channels[type] = work_channel
|
||||
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 = []
|
||||
front_end_socket.recv_strings(msg)
|
||||
raw_address = msg[0]
|
||||
raw_msg = msg[2]
|
||||
begin
|
||||
parsed_msg = JSON.parse(raw_msg)
|
||||
rescue JSON::ParserError => e
|
||||
reply = [msg.first, "", { status: :parse_error }.to_json]
|
||||
front_end_socket.send_strings(reply)
|
||||
return
|
||||
end
|
||||
action = parsed_msg["action"]
|
||||
if action.nil?
|
||||
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)
|
||||
req = FrontEndRequest.recv(front_end_socket)
|
||||
req.handle do |action|
|
||||
if action == "configure_worker"
|
||||
respond_with_worker_config(req)
|
||||
elsif action == "analyze_iteration"
|
||||
handle_with_worker(:static_analyzers, req)
|
||||
elsif action == "test_solution"
|
||||
handle_with_worker(:test_runners, req)
|
||||
elsif action == "represent"
|
||||
handle_with_worker(:representers, req)
|
||||
else
|
||||
req.send_error({ status: :unrecognised_action })
|
||||
end
|
||||
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]
|
||||
front_end_socket.send_strings(reply)
|
||||
elsif worker_available?(channel)
|
||||
forward_to_backend(channel, msg)
|
||||
req.send_error({ status: :worker_class_unknown })
|
||||
elsif channel.worker_available?
|
||||
context = { credentials: temp_credentials }
|
||||
@in_flight_requests.register(req)
|
||||
channel.forward_to_backend(req, context)
|
||||
else
|
||||
reply = [msg.first, "", { status: :worker_unavailable }.to_json]
|
||||
front_end_socket.send_strings(reply)
|
||||
req.send_error({ status: :worker_unavailable })
|
||||
end
|
||||
end
|
||||
|
||||
@@ -133,37 +254,18 @@ module Pipeline::Rpc
|
||||
end
|
||||
|
||||
def on_service_response
|
||||
msg = ""
|
||||
response_socket.recv_string(msg)
|
||||
status_message = JSON.parse(msg)
|
||||
type = status_message["msg_type"]
|
||||
if type == "response"
|
||||
return_address = status_message["return_address"]
|
||||
reply = [return_address.pack("c*"), "", msg]
|
||||
front_end_socket.send_strings(reply, ZMQ::DONTWAIT)
|
||||
elsif type == "heartbeat"
|
||||
flush_expired_requests
|
||||
msg = ServiceResponse.recv(response_socket)
|
||||
if msg.type == "response"
|
||||
@in_flight_requests.forward_response(msg)
|
||||
@in_flight_requests.unregister(msg.return_address)
|
||||
elsif msg.type == "heartbeat"
|
||||
@in_flight_requests.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]
|
||||
front_end_socket.send_strings(reply)
|
||||
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)
|
||||
analyzer_spec["credentials"] = temp_credentials
|
||||
req.send_result(analyzer_spec)
|
||||
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]
|
||||
msg
|
||||
session.to_h[:credentials]
|
||||
end
|
||||
end
|
||||
end
|
||||
|
||||
@@ -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"]
|
||||
|
||||
Reference in New Issue
Block a user