Refactor spike into classes

This commit is contained in:
Charles Care
2019-10-10 19:36:51 +01:00
parent c692f92d11
commit e0507d894a
2 changed files with 191 additions and 103 deletions

View File

@@ -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

View File

@@ -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"]