示例#1
0
    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()
示例#2
0
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')