def main(): parser = gather_parsers(get_argument_parser(), (res.scheduling,)) args = parser.parse_args() Logger.setup_logging() initialize_configuration() logger = logging.getLogger("main") loop = asyncio.get_event_loop() session_id = str(uuid.uuid4()) try: yield from Logger.duplicate_logs_to_mongo(session_id, "main") except: logger.exception("Failed to setup logging to MongoDB") install_watchdog(r.watchdog.interval) db_manager = DBManager(**r.db) yield from db_manager.initialize() heap = Heap() heap.initialize((yield from db_manager.fetch_all())) worker = Worker(db_manager, heap, r.amqp, r.scheduler.poll_interval, r.scheduler.timeout, (yield from db_manager.fetch_pending())) yield from worker.initialize() yield from worker.work() for signame in "SIGINT", "SIGTERM": @asyncio.coroutine def terminate(): if not args.disable_logging_to_mongo: yield from Logger.discard_logs_to_mongo() yield from worker.stop() yield from db_manager.shutdown() loop.stop() def sighandler(): logger.critical("Received %s", signame) loop.create_task(terminate()) loop.add_signal_handler(getattr(signal, signame), sighandler) logger.debug("Installed signal handlers")
def _test_no_db_amqp(self): db_mamanger = DBManagerMock() heap = Heap() cfg = Config('test') worker = Worker(db_mamanger, heap, cfg, poll_interval=100500, default_timeout=10050, pending=tuple()) worker._amqp_channel_source = AMQPChannelMock() worker._amqp_channel_trigger = AMQPChannelMock() date = datetime.now(pytz.utc) date -= timedelta(microseconds=date.microsecond % 1000) def make_msg_bytes(msg): return json.dumps(msg, default=json_util.default).encode("utf-8") msg = {"action": "enqueue", "due_date": date, "data": "hello"} msg_data = make_msg_bytes(msg) yield from worker._amqp_callback_source( msg_data, EnvelopeMock(), PropertiesMock()) self.assertEqual(1, worker._heap.size()) self.assertEqual(date, worker._heap.min()[0]) self.assertEqual("", worker._heap.min()[1][0]) self.assertEqual(1, worker._heap.min()[1][1]) self.assertIsNone(worker._heap.min()[1][2]) self.assertEqual("hello", worker._heap.min()[1][4]) date -= timedelta(days=10) msg = {"action": "enqueue", "due_date": date, "data": "world", "id": "unique"} msg_data = make_msg_bytes(msg) yield from worker._amqp_callback_source( msg_data, EnvelopeMock(), PropertiesMock()) self.assertEqual(2, worker._heap.size()) self.assertEqual(date, worker._heap.min()[0]) self.assertEqual("unique", worker._heap.min()[1][0]) self.assertEqual(2, worker._heap.min()[1][1]) self.assertIsNone(worker._heap.min()[1][2]) self.assertEqual("world", worker._heap.min()[1][4]) date += timedelta(days=5) msg = {"action": "enqueue", "due_date": date, "data": "other", "expire_in": 1} msg_data = make_msg_bytes(msg) yield from worker._amqp_callback_source( msg_data, EnvelopeMock(), PropertiesMock()) self.assertEqual(3, worker._heap.size()) self.assertEqual("world", worker._heap.min()[1][4]) self.assertEqual(date, worker._heap._list[2][0]) self.assertEqual("", worker._heap._list[2][1][0]) self.assertEqual(3, worker._heap._list[2][1][1]) self.assertEqual(1, worker._heap._list[2][1][2]) self.assertEqual("other", worker._heap._list[2][1][4]) tasks = worker._poll() for task in tasks: yield from task self.assertEqual(0, worker._heap.size()) self.assertEqual(2, len(worker._amqp_channel_trigger.published)) self.assertEqual(2, len(worker._pending_tasks)) msg = {"action": "enqueue", "due_date": date, "data": "hello"} msg_data = make_msg_bytes(msg) yield from worker._amqp_callback_source( msg_data, EnvelopeMock(), PropertiesMock()) self.assertEqual(1, worker._heap.size()) msg = {"action": "cancel", "id": 4} msg_data = make_msg_bytes(msg) yield from worker._amqp_callback_source( msg_data, EnvelopeMock(), PropertiesMock()) self.assertEqual(1, worker._heap.size()) self.assertIn(4, worker._cancelled_tasks) tasks = worker._poll() self.assertEqual(0, len(tasks))