class Slave(Actionable):
"""
This class manages connections to the master and handles client-server logistics.
Actions on it are run by master.
"""
id = None
last_bw_out = None
bw_out = 0
total_bw_out = 0
last_bw_in = None
bw_in = 0
total_bw_in = 0
running_tests = []
test_runner = None
sink_obj = None
sink = None
client_tests = {}
message_lock = BoundedSemaphore(1)
def __init__(self, context):
self.context = context
# This socket sends messages to master
self.socket_out = context.socket(zmq.PUSH)
self.socket_out.connect("tcp://localhost:{}".format(SLAVE_MASTER_PORT))
# Syntax allowing self.master.<action>(data)
self.master = Sender(self.send_to_master)
# This socket receives broadcasts from master
self.socket_in = context.socket(zmq.SUB)
self.socket_in.setsockopt(zmq.SUBSCRIBE, "")
self.socket_in.connect("tcp://localhost:{}".format(SLAVE_PUB_PORT))
self.sink_host = '{}:{}'.format(self.hostname(), SLAVE_SINK_PORT)
self.uuid = uuid4().hex
def send_to_master(self, data):
"""
Sends a bson serialized message to master, inserting slave specific information such as id
:param data: Message data to send
action: (required) RPC to call on serverside slave instance
"""
if self.id:
data['slave_id'] = self.id
if data['action'] != 'heartbeat':
log.debug('MASTER->{}'.format(data['action']))
self.socket_out.send(bson.dumps(data))
if data['action'] != 'heartbeat':
log.debug('MASTER->{} end'.format(data['action']))
def send_to_sink(self, data):
"""
Send bson encoded data to the sink for collation
:param data: Message data to send
"""
if self.id:
data['slave_id'] = self.id
if self.sink_obj:
self.sink_obj.send_to_sink(data)
else:
log.warn('{}->SINK failed, no sink'.format(data['action']))
def listen_to_master(self):
"""
Wait for messages from the server and delegate them
"""
log.info("Listening to master")
while True:
self.handle_master_message(self.socket_in.recv())
def handle_master_message(self, msg):
"""
Decode a message from the server, if it matches our id or has no id decode and run it.
:param msg: BSON encoded message from the server
action: Local function to run
slave_id: If present and not our id ignore message
"""
data = bson.loads(msg)
if 'action' in data:
if 'slave_id' not in data or data['slave_id'] == self.id:
log.debug('{}<-MASTER'.format(data['action']))
self.run_action(data)
log.debug('{}<-MASTER end'.format(data['action']))
else:
log.warn('Server sent message with no action')
def heartbeat(self):
"""
Generate a status report and send it to master
"""
# physical_memory = psutil.phymem_usage()
virtual_memory = psutil.virtual_memory()
swap_memory = psutil.swap_memory()
load = psutil.cpu_percent(interval=1)
network = psutil.net_io_counters()
hostname = self.hostname()
if not self.last_bw_out:
self.last_bw_out = network.bytes_sent
self.bw_out = (network.bytes_sent - self.last_bw_out)/HEARTBEAT_PERIOD
self.total_bw_out += self.bw_out * HEARTBEAT_PERIOD
self.last_bw_out = network.bytes_sent
if not self.last_bw_in:
self.last_bw_in = network.bytes_recv
self.bw_in = (network.bytes_recv - self.last_bw_in)/HEARTBEAT_PERIOD
self.total_bw_in += self.bw_in * HEARTBEAT_PERIOD
self.last_bw_in = network.bytes_recv
self.master.heartbeat({
'hostname': hostname,
'memory': {
# 'physical': (physical_memory.used, physical_memory.total),
'physical': (0, 0),
'virtual': (virtual_memory.used, virtual_memory.total),
'swap': (swap_memory.used, swap_memory.total)
},
'load': load,
'bandwidth': {
'in': (self.bw_in, self.total_bw_in),
'out': (self.bw_out, self.total_bw_out)
},
'generated': int(time()),
'sink_id': self.sink_obj.id if self.sink_obj else None,
'slave_uuid': self.uuid
})
def heartbeat_forever(self):
log.info('Beating heart forever')
while True:
self.heartbeat()
gevent.sleep(HEARTBEAT_PERIOD)
@action
def set_id(self, data):
"""
Set id if it matches our uuid
:param data: Message specific data:
slave_uuid: UUID to check against ours
new_slave_id: Id to set if match
"""
if 'slave_uuid' not in data or 'new_slave_id' not in data:
log.warn('receiving id without slave_uuid/new_slave_id')
return
if data['slave_uuid'] == self.uuid:
self.id = data['new_slave_id']
log.info('Received id: {}'.format(self.id))
@action
def set_sink(self, data):
"""
Set sink
:param data: Message specific data:
sink_id: Sink id to connect to
sink_host: Host to connect to
:return:
"""
if 'sink_id' in data and 'sink_host' in data and data['sink_host'] == self.sink_host:
self.connect_to_sink(data['sink_id'], data['sink_host'])
else:
log.warn('Incomplete sink call')
@staticmethod
def hostname():
return socket.gethostname()
def connect_to_sink(self, sink_id, host):
"""
Attempt to connect to the given sink, and create a local sink instance
:param sink_id: ID to assign to the sink
:param host: host to connect to
"""
log.info('Attempting to connect to sink({}) - {}'.format(sink_id, host))
if self.sink_obj:
log.info('Closing previous sink')
self.sink_obj.close()
# Create our sink
self.sink_obj = Sink(id=sink_id,
host=host,
slave=self)
# Initialize it, and tell master if successful
if self.sink_obj.setup():
log.info('Ok!')
# RPC for sink
self.sink = Sender(self.send_to_sink)
self.master.connected_to_sink({
'sink_id': self.sink_obj.id
})
@action
def quit(self, data):
"""
If this RPC is run master told this slave to quit
"""
log.info('Master told us to quit, quitting.')
sys.exit(0)
@action
def run_test(self, data):
"""
Receive a test, instantiate it and run
:param data: Message specific data:
client_id: Client that started the test
test: Test data to hand off to the test
:return:
"""
log.info('slave({}).run_test client({})'.format(self.id, data['client_id']))
if not self.sink_obj:
log.warn('Error, no sink')
return
if 'test' in data:
client_id = data['client_id']
test = Test(self.sink, data)
# Only allow one test at a time
if client_id in self.client_tests:
log.info('Stopping previous test')
self.client_tests[client_id].stop()
self.client_tests[data['client_id']] = test
test.run()
else:
log.warn('No test given')
@action
def stop_test(self, data):
"""
Stop any tests from that client
:param data: Message specific data:
client_id: Client to stop tests for
"""
client_id = data['client_id']
if client_id in self.client_tests:
log.info('slave({}).stop_test client({})'.format(self.id, client_id))
self.client_tests[client_id].stop()
del self.client_tests[client_id]
else:
self.master.test_stopped()
class Client(Actionable):
"""
This class is the local representation of a remote client (browser)
When a websocket message is received it is routed into the proper local instance
When the master sends a message to a client it goes through this class
"""
def __init__(self, master, websocket, id):
self.master = master
self.websocket = websocket
self.id = id
# This is syntax for sending message through a websocket.
# client.send.<action>(data)
self.send = Sender(self.send_data)
def send_data(self, data):
self.websocket.send(json.dumps(data))
def handle_websocket_message(self, msg):
"""
This handles any message that comes in by looking at the action key and seeing if it is a local method,
then calling it
:param msg: JSON encoded data from the websocket
"""
if msg:
data = json.loads(msg)
if "action" in data:
self.run_action(data)
@action
def get_id(self, data):
self.send.set_id({"client_id": self.id})
@action
def quit(self, data):
"""
Websocket telling a slave to quit
:param data: Message specific data:
slave_id: slave to tell to quit
"""
if "slave_id" in data:
self.master.remove_slave(data["slave_id"])
else:
self.send.error({"error": "Id not specified in data"})
@action
def request_slaves(self, data):
"""
Websocket requesting slave data
"""
data = {"slaves": dict((slave.id, slave.last_beat) for slave in self.master.slave_registry.values())}
# Send a message with action=receive_slaves to the browser
self.send.receive_slaves(data)
@action
def set_sink(self, data):
"""
Set the sink on a particular slave, which sets the sink for all slaves on that host
:param data: message specific data:
slave_id: Slave to set the sink for (unused)
sink_id: Slave to set the sink to
"""
if "sink_id" not in data or "slave_id" not in data:
raise Exception("Missing sink_id or slave_id")
# Let the master handle actually doing this
self.master.set_sink(data["slave_id"], data["sink_id"], self)
@action
def request_available_tests(self, data):
"""
Client asking to load tests
:param data:
"""
tests_glob = os.path.join(os.path.dirname(__file__), TEST_DIR, "*.json")
files = glob.glob(tests_glob)
files_stripped = [os.path.basename(filename)[:-5] for filename in files]
# Send list of files to the client
self.send.receive_available_tests({"tests": files_stripped})
@action
def request_test(self, data):
"""
Client asking to load a test
:param data: Message specific data:
name: filename to load
"""
filename = os.path.join(os.path.dirname(__file__), TEST_DIR, data["name"] + ".json")
# TODO: Prevent path traversal
if os.path.exists(filename):
with open(filename) as file:
return self.send.receive_test({"test": json.loads(file.read())})
return self.send.error({"error": "Test not found"})
@action
def save_test(self, data):
"""
Save a test to a file
:param data: Message specific data:
name: Filename to save to
test: data representing a test
"""
if "test" in data:
test = data["test"]
if "name" in test:
# TODO: Prevent path traversal
filename = os.path.join(os.path.dirname(__file__), TEST_DIR, test["name"] + ".json")
with open(filename, "w") as file:
file.write(json.dumps(test))
self.send.save_successful()
@action
def delete_test(self, data):
"""
Delete a test
:param data: Message specific data:
test_name: Name of test to delete
"""
if "test_name" in data:
# TODO: Prevent path traversal
filename = os.path.join(os.path.dirname(__file__), TEST_DIR, data["test_name"] + ".json")
if os.path.exists(filename):
os.remove(filename)
else:
return self.send.error({"error": "Test not found"})
else:
return self.send.error({"error": "Cannot delete test, not found"})
@action
def run_test(self, data):
"""
Send a test to all slaves to run
:param data: Message specific data:
test: Test data to run
"""
if "test" in data:
data["client_id"] = self.id
if "runs" in data["test"]:
try:
# Each slave gets a portion of the total runs
data["runs"] = int(data["test"]["runs"]) / len(self.master.slave_registry)
except ValueError:
data["runs"] = 1
else:
data["runs"] = 1
self.master.slaves.run_test(data)
@action
def stop_test(self, data):
"""
Halt all tests run by this client
:param data:
"""
data["client_id"] = self.id
if len(self.master.slave_registry) != 0:
self.master.slaves.stop_test(data)
else:
self.send.test_stopped()
