def __init__(self, host, slave_ip, screen, s3_bucket,
s3_folder, num_packup, max_download_sec, awake_frequency, port=5672):
self.slave_ip = slave_ip
self.port = port
self.screen = screen
self.s3_bucket = s3_bucket
self.s3_folder = s3_folder
self.num_packup = num_packup
self.host = host
self.max_download_sec = max_download_sec
self.frequency = awake_frequency
self.is_started = False
self.worker_lord = GapConveyWorkerLord(
PureHtmlWorker(),
ZipFolderConveyWorker(s3_dest_bucket=s3_bucket, s3_dest_folder=s3_folder),
self.s3_bucket,
self.s3_folder,
self.num_packup
)
self.task_connection = pika.BlockingConnection(
pika.ConnectionParameters(host=self.host, port=self.port))
self.task_channel = self.task_connection.channel()
self.task_channel.queue_declare(queue='task', durable=True)
self.task_channel.basic_qos(prefetch_count=1)
self.downloadertag = self.task_channel.basic_consume(self.callback, queue='task')
self.task_channel.exchange_declare(exchange='broadcast',
type='fanout')
self.task_channel.queue_bind(exchange='broadcast',
queue='task')
self.message_connection = pika.BlockingConnection(
pika.ConnectionParameters(host=self.host, port=self.port))
self.message_channel = self.message_connection.channel()
self.message_channel.queue_declare(queue='message')
self.beating = True
self.heart_beat_handler = threading.Thread(target=self.send_heart_beat)
self.heart_beat_handler.start()
logging.info(' [*] Waiting for messages. To exit press CTRL+C')
try:
self.is_started = True
self.task_channel.start_consuming()
except KeyboardInterrupt:
logging.info('KeyboardInterrupt. Stop.')
self.stop()
except RuntimeError:
logging.info('RuntimeError. Stop.')
self.stop()
class Slave:
def __init__(self, host, slave_ip, screen, s3_bucket,
s3_folder, num_packup, max_download_sec, awake_frequency, port=5672):
self.slave_ip = slave_ip
self.port = port
self.screen = screen
self.s3_bucket = s3_bucket
self.s3_folder = s3_folder
self.num_packup = num_packup
self.host = host
self.max_download_sec = max_download_sec
self.frequency = awake_frequency
self.is_started = False
self.worker_lord = GapConveyWorkerLord(
PureHtmlWorker(),
ZipFolderConveyWorker(s3_dest_bucket=s3_bucket, s3_dest_folder=s3_folder),
self.s3_bucket,
self.s3_folder,
self.num_packup
)
self.task_connection = pika.BlockingConnection(
pika.ConnectionParameters(host=self.host, port=self.port))
self.task_channel = self.task_connection.channel()
self.task_channel.queue_declare(queue='task', durable=True)
self.task_channel.basic_qos(prefetch_count=1)
self.downloadertag = self.task_channel.basic_consume(self.callback, queue='task')
self.task_channel.exchange_declare(exchange='broadcast',
type='fanout')
self.task_channel.queue_bind(exchange='broadcast',
queue='task')
self.message_connection = pika.BlockingConnection(
pika.ConnectionParameters(host=self.host, port=self.port))
self.message_channel = self.message_connection.channel()
self.message_channel.queue_declare(queue='message')
self.beating = True
self.heart_beat_handler = threading.Thread(target=self.send_heart_beat)
self.heart_beat_handler.start()
logging.info(' [*] Waiting for messages. To exit press CTRL+C')
try:
self.is_started = True
self.task_channel.start_consuming()
except KeyboardInterrupt:
logging.info('KeyboardInterrupt. Stop.')
self.stop()
except RuntimeError:
logging.info('RuntimeError. Stop.')
self.stop()
def send_heart_beat(self):
while self.beating:
reply = {'host': self.slave_ip, 'screen': self.screen}
self.message_channel.basic_publish(exchange='',
routing_key='message',
body='AWAKE ' + ujson.dumps(reply))
time.sleep(self.frequency)
def callback(self, ch, method, properties, body):
logging.info("Received %r" % (body,))
if type(body) is bytes:
body = body.decode('utf-8')
command = body[:body.find(' ')]
info = body[body.find(' ') + 1:]
if command == 'WORK':
task_string = ujson.loads(info)['task']
task = ujson.loads(task_string)
file_name, error = self.handle_work(task)
ch.basic_ack(delivery_tag=method.delivery_tag)
elif command == 'STOP':
# ch.basic_ack(delivery_tag=method.delivery_tag)
logging.info('Received stop')
self.stop()
else:
logging.warning("Unknown command %r" % (body,))
def handle_work(self, task):
error = ''
try:
signal.signal(signal.SIGALRM, timeout_handler)
signal.alarm(self.max_download_sec)
filename = self.worker_lord.run(task)
signal.alarm(0)
except RuntimeError as e:
if 'timeout' in str(e):
logging.error('Timeout')
else:
logging.info(e)
filename = None
error = 'timeout'
# except Exception as e:
# logging.info(e)
# filename = None
# error = e
return filename, error
def stop(self):
logging.info('stop')
if not self.is_started:
logging.info('I want to stop but I am already stopped.')
return
signal.alarm(0)
self.worker_lord.close()
self.beating = False
self.heart_beat_handler.join()
logging.info('Closing slave connection')
self.task_connection.close()
self.message_connection.close()
logging.info('Slave connection closed')
os.system('screen -S ' + self.screen + ' -X quit')
