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