async def connect(
self,
on_close: Callable[['_Connection', Exception], Awaitable[None]],
) -> None:
with wrap2span(
name=AmqpSpan.NAME_CONNECT,
kind=AmqpSpan.KIND_CLIENT,
cls=AmqpSpan,
app=self.pika.app,
) as span:
span.tag(AmqpSpan.TAG_URL, self.pika._masked_url)
async with self._lock:
self._state = self.STATE_CONNECTING
self._on_close = on_close
self._fut = asyncio.Future()
self._conn = AsyncioConnection(
parameters=pika.URLParameters(self.cfg.url),
on_open_callback=self._on_connection_open,
on_open_error_callback=self._on_connection_open_error,
on_close_callback=self._on_connection_closed,
)
try:
await asyncio.wait_for(self._fut,
timeout=self.cfg.connect_timeout)
if self._fut.result() != 1:
raise RuntimeError()
self._state = self.STATE_CONNECTED
except Exception:
self._state = self.STATE_CLOSED
raise
finally:
self._fut = None
def connect(self):
"""
This method connects to RabbitMQ, returning the connection handle.
When the connection is established, the on_connection_open method
will be invoked by pika.
------------------
:rtype: pika.SelectConnection
"""
if self.async_processing:
asyncio.set_event_loop(asyncio.new_event_loop())
return AsyncioConnection(
parameters=pika.URLParameters("{}{}".format(
self._url, self.heartbeat)),
on_open_callback=self.on_connection_open,
on_open_error_callback=self.on_open_error_callback,
# on_close_callback=self.on_connection_closed
)
else:
self._connection = pika.BlockingConnection(
parameters=pika.URLParameters("{}{}".format(
self._url, self.heartbeat)),
# on_open_callback=self.on_connection_open,
# on_open_error_callback=self.on_open_error_callback,
# on_close_callback=self.on_connection_closed
)
self.on_connection_open(self._connection)
return self._connection
def connect(self):
"""This method connects to RabbitMQ, returning the connection handle.
When the connection is established, the on_connection_open method
will be invoked by pika.
:rtype: pika.adapters.asyncio_connection.AsyncioConnection
"""
LOGGER.info('Connecting to %s', self._url)
return AsyncioConnection(
parameters=pika.URLParameters(self._url),
on_open_callback=self.on_connection_open,
on_open_error_callback=self.on_connection_open_error,
on_close_callback=self.on_connection_closed)
async def listen_to_rabbitmq(queue_name):
loop = asyncio.get_running_loop()
fut = loop.create_future()
AsyncioConnection(
pika.ConnectionParameters('localhost'),
on_open_callback=lambda c: fut.set_result(c),
on_open_error_callback=lambda c, exc: fut.set_exception(exc),
on_close_callback=lambda c, exc: fut.set_exception(exc))
conn = await fut
fut = loop.create_future()
conn.channel(on_open_callback=lambda ch: fut.set_result(ch))
channel = await fut
fut = loop.create_future()
channel.queue_declare(queue=queue_name,
durable=True,
auto_delete=True,
arguments={'x-expires': 86400000},
callback=lambda ch: fut.set_result(ch))
cnt = await fut
if cnt:
messages = asyncio.Queue()
channel.basic_consume(
queue_name,
on_message_callback=lambda *args: messages.put_nowait(args[3]),
auto_ack=True)
fut = loop.create_future()
def on_close_callback(ch, reason):
try:
raise reason
except ChannelClosedByClient:
fut.cancel()
except:
fut.set_exception(reason)
channel.add_on_close_callback(on_close_callback)
while True:
msg = asyncio.ensure_future(messages.get())
await asyncio.wait([msg, fut], return_when=asyncio.FIRST_COMPLETED)
if fut.done():
if fut.cancelled():
break
raise fut.exception()
decoded_message = msg.result().decode('utf-8').replace(
'null', 'None')
decoded_message_dict = ast.literal_eval(decoded_message)
return decoded_message_dict['result']
def connect(self):
LOGGER.info('Connecting to %s', self._url)
url = urlparse(self._url)
creds = pika.PlainCredentials(url.username, url.password)
parameters = pika.ConnectionParameters(
host=url.hostname,
port=url.port,
virtual_host=url.path,
credentials=creds
)
return AsyncioConnection(
parameters=parameters,
on_open_callback=self.on_connection_open,
on_open_error_callback=self.on_connection_open_error,
on_close_callback=self.on_connection_closed)
def connect(self):
"""This method connects to RabbitMQ, returning the connection handle.
When the connection is established, the on_connection_open method
will be invoked by pika.
Returns
-------
AsyncioConnection
"""
LOGGER.info('Connecting...')
return AsyncioConnection(
parameters=pika.URLParameters(self._url),
on_open_callback=self.on_connection_open,
on_open_error_callback=self.on_connection_open_error)
def try_connect(self):
self.logger.info("Creating connection to RabbitMQ")
self._io_loop.call_later(self._timeout, self._on_timeout)
if isinstance(self._io_loop, IOLoop):
TornadoConnection(self._parameters,
on_open_callback=self._open_callback,
on_open_error_callback=self._open_error_callback,
on_close_callback=self._close_callback,
custom_ioloop=self._io_loop)
else:
AsyncioConnection(self._parameters,
on_open_callback=self._open_callback,
on_open_error_callback=self._open_error_callback,
on_close_callback=self._close_callback,
custom_ioloop=self._io_loop)
def connect(self):
self._connection = AsyncioConnection(
pika.URLParameters(self._url),
on_open_callback=self.on_connection_open,
on_close_callback=self.on_connection_closed,
on_open_error_callback=self.on_connection_open_error)
class Publisher(object):
"""RabbitMQ publisher
Parameters
----------
amqp_url : `str`
Broker url
csc_parent : `lsst.dm.csc.base.dm_csc`
CSC service using this publisher
logger_level : `int`
log level
"""
def __init__(self, amqp_url, csc_parent=None, logger_level=LOGGER.info):
# only emit logging messages from pika and WARNING and above
logging.getLogger("pika").setLevel(logging.WARNING)
self._connection = None
self._channel = None
self._url = amqp_url
self.setup_complete_event = asyncio.Event()
self.setup_complete_event.clear()
self._message_handler = YamlHandler()
self._stopping = False
self.csc_parent = csc_parent
self.logger_level = logger_level
def connect(self):
self._connection = AsyncioConnection(
pika.URLParameters(self._url),
on_open_callback=self.on_connection_open,
on_close_callback=self.on_connection_closed,
on_open_error_callback=self.on_connection_open_error)
def on_connection_open(self, connection):
LOGGER.info("connection opened")
self._connection = connection
self.open_channel()
def on_connection_open_error(self, _unused_connection, err):
"""This method is called by pika if the connection to RabbitMQ
can't be established.
:param pika.adapters.asyncio_connection.AsyncioConnection _unused_connection:
The connection
:param Exception err: The error
"""
LOGGER.error(f'Connection open failed: {err}')
if self.csc_parent is not None:
self.csc_parent.fault(5071, f'Connection open failed: {err}')
def open_channel(self):
LOGGER.info("creating a new channel")
self._connection.channel(on_open_callback=self.on_channel_open)
def on_channel_open(self, channel):
LOGGER.info("channel opened")
self._channel = channel
self.add_on_channel_close_callback()
self.setup_complete_event.set()
def add_on_channel_close_callback(self):
LOGGER.info('adding channel close callback')
self._channel.add_on_close_callback(self.on_channel_closed)
def on_channel_closed(self, channel, reason):
LOGGER.info('Channel %i was closed %s' % (channel, reason))
self._channel = None
self._connection.close()
def on_connection_closed(self, connection, reason):
"""
Params:
connection - closed object connection
reason - reason for closure
"""
LOGGER.info("on_connection_closed called")
self._channel = None
self._connect = None
async def close(self):
if self._channel is not None:
self._channel.close()
async def publish_message(self, route_key, msg):
encoded_data = self._message_handler.encode_message(msg)
self.logger_level("Sending msg to %s", route_key)
# Since this is asynchronous, it's possible to still be in the
# process of getting setup and having self._channel be None when
# publish_message is being called from another thread, so we wait
# here until setup is completed.
await self.setup_complete_event.wait()
self._channel.basic_publish(exchange='message',
routing_key=route_key,
body=encoded_data)
self.logger_level(f'message sent message body is: {msg}')
async def stop(self):
self._stopping = True
await self.close()
LOGGER.info('Stopped')
async def start(self):
try:
await self.connect()
except Exception:
LOGGER.error('No channel - connection channel is None')
class _Connection:
STATE_NONE = 0
STATE_CONNECTING = 1
STATE_CONNECTED = 2
STATE_CLOSING = 3
STATE_CLOSED = 4
def __init__(
self,
pika: 'Pika',
cfg: PikaConfig,
*,
loop: Optional[asyncio.AbstractEventLoop] = None,
) -> None:
self.pika = pika
self.cfg = cfg
self._conn: Optional[AsyncioConnection] = None
self._loop: asyncio.AbstractEventLoop = (loop
or asyncio.get_event_loop())
self._fut: Optional[asyncio.Future] = None
self._lock = asyncio.Lock()
self._on_close: Optional[Callable[['_Connection', Exception],
Awaitable[None]]] = None
self._state = self.STATE_NONE
@property
def state(self) -> int:
return self._state
async def connect(
self,
on_close: Callable[['_Connection', Exception], Awaitable[None]],
) -> None:
with wrap2span(
name=AmqpSpan.NAME_CONNECT,
kind=AmqpSpan.KIND_CLIENT,
cls=AmqpSpan,
app=self.pika.app,
) as span:
span.tag(AmqpSpan.TAG_URL, self.pika._masked_url)
async with self._lock:
self._state = self.STATE_CONNECTING
self._on_close = on_close
self._fut = asyncio.Future()
self._conn = AsyncioConnection(
parameters=pika.URLParameters(self.cfg.url),
on_open_callback=self._on_connection_open,
on_open_error_callback=self._on_connection_open_error,
on_close_callback=self._on_connection_closed,
)
try:
await asyncio.wait_for(self._fut,
timeout=self.cfg.connect_timeout)
if self._fut.result() != 1:
raise RuntimeError()
self._state = self.STATE_CONNECTED
except Exception:
self._state = self.STATE_CLOSED
raise
finally:
self._fut = None
async def close(self) -> None:
async with self._lock:
if self._conn is None:
self._state = self.STATE_CLOSED
return
self._state = self.STATE_CLOSING
self._on_close = None
self._fut = asyncio.Future()
self._conn.close()
try:
await asyncio.wait_for(self._fut,
timeout=self.cfg.connect_timeout)
except pika.exceptions.ConnectionClosed:
self._state = self.STATE_CLOSED
return
finally:
self._fut = None
def _on_connection_open(self,
_unused_connection: AsyncioConnection) -> None:
if self._fut is not None:
self._fut.set_result(1)
def _on_connection_open_error(self, _unused_connection: AsyncioConnection,
err: Exception) -> None:
if self._fut is not None:
self._fut.set_exception(err)
elif self._on_close is not None:
asyncio.ensure_future(self._on_close(self, err))
def _on_connection_closed(self, _unused_connection: AsyncioConnection,
reason: Exception) -> None:
if self._fut is not None:
self._fut.set_exception(reason)
elif self._on_close is not None:
asyncio.ensure_future(self._on_close(self, reason))
async def channel(
self,
on_close: Optional[Callable[[pika.channel.Channel, Exception],
Awaitable[None]]],
name: Optional[str],
) -> pika.channel.Channel:
with wrap2span(
name=AmqpSpan.NAME_CHANNEL,
kind=AmqpSpan.KIND_CLIENT,
cls=AmqpSpan,
app=self.pika.app,
) as span:
if name is not None:
span.tag(AmqpSpan.TAG_CHANNEL_NAME, name)
async with self._lock:
if self._conn is None:
raise pika.exceptions.AMQPConnectionError()
self._fut = asyncio.Future()
self._conn.channel(
on_open_callback=partial(self._on_channel_open, on_close))
try:
channel: Any = await asyncio.wait_for(
self._fut, timeout=self.cfg.channel_open_timeout)
if not isinstance(channel, pika.channel.Channel):
raise RuntimeError()
finally:
self._fut = None
span.tag(AmqpSpan.TAG_CHANNEL_NUMBER,
str(channel.channel_number))
return channel
def _on_channel_open(
self,
on_close: Optional[Callable[[pika.channel.Channel, Exception],
Awaitable[None]]],
channel: pika.channel.Channel,
) -> None:
channel.add_on_close_callback(
partial(self._on_channel_closed, on_close))
if self._fut is None:
raise RuntimeError()
self._fut.set_result(channel)
def _on_channel_closed(
self,
on_close: Optional[Callable[[pika.channel.Channel, Exception],
Awaitable[None]]],
channel: pika.channel.Channel,
reason: Exception,
) -> None:
if on_close is not None:
asyncio.ensure_future(on_close(channel, reason))
def _connect(self) -> None:
self._conn = AsyncioConnection(
parameters=get_mq_conn_params(),
on_open_callback=self._conn_on_open,
on_open_error_callback=self._conn_on_open_error,
on_close_callback=self._conn_on_close)
class Worker(object):
def __init__(self, db_config: dict, max_processes: int) -> None:
self._max_processes = max_processes
self._executor = ProcessPoolExecutor(
max_workers=max_processes,
mp_context=mp.get_context('spawn'),
initializer=partial(_initializer, db_config))
self._queue_name = 'judge_queue'
self._conn: AsyncioConnection = None
self._ch: Channel = None
self._hostname: Optional[str] = None
self._pid = os.getpid()
self._task_processed, self._task_errors = 0, 0
self._maint_interval = timedelta(seconds=60)
def __enter__(self) -> 'Worker':
return self
def __exit__(self, exc_type: Any, exc_value: Any, traceback: Any) -> None:
self._executor.shutdown(wait=False)
if self._ch:
self._ch.close()
if self._conn:
self._conn.close()
def start(self) -> None:
self._connect()
asyncio.get_event_loop().call_soon_threadsafe(self._update_status)
asyncio.get_event_loop().run_forever()
def _connect(self) -> None:
self._conn = AsyncioConnection(
parameters=get_mq_conn_params(),
on_open_callback=self._conn_on_open,
on_open_error_callback=self._conn_on_open_error,
on_close_callback=self._conn_on_close)
def _schedule_update_status(self) -> None:
asyncio.get_event_loop().call_later(
uniform(
self._maint_interval.total_seconds() - 1,
self._maint_interval.total_seconds() + 1,
), self._update_status)
def _update_status(self) -> None:
updates = dict(
last_contact=func.now(),
processed=self._task_processed,
errors=self._task_errors,
)
try:
hostname = self._hostname or gethostname()
with transaction() as s:
if self._hostname:
s.query(WorkerTable).filter(
WorkerTable.hostname == self._hostname,
WorkerTable.pid == self._pid).update(
updates, synchronize_session=False)
else:
updates.update(
dict(hostname=hostname,
pid=self._pid,
max_processes=self._max_processes,
startup_time=func.now()))
s.add(WorkerTable(**updates))
if uniform(0, 1) <= 0.01 or not self._hostname:
threshold = self._maint_interval * 10
s.query(WorkerTable).filter(
func.now() -
WorkerTable.last_contact > threshold).delete(
synchronize_session=False)
self._hostname = hostname
except Exception:
pass
self._schedule_update_status()
def _conn_on_open(self, _: AsyncioConnection) -> None:
self._conn.channel(on_open_callback=self._ch_on_open)
def _conn_on_open_error(self, _: AsyncioConnection,
err: AMQPError) -> None:
LOGGER.warning(
'Cannot open RabbitMQ connection({}). retrying...'.format(err))
asyncio.get_event_loop().call_later(uniform(1, 5), self._connect)
def _conn_on_close(self, _: AsyncioConnection, reason: AMQPError) -> None:
LOGGER.warning(
'RabbitMQ connection closed({}). retrying...'.format(reason))
asyncio.get_event_loop().call_later(uniform(1, 5), self._connect)
def _ch_on_open(self, ch: Channel) -> None:
self._ch = ch
ch.add_on_close_callback(self._ch_on_close)
ch.queue_declare(queue=self._queue_name,
callback=self._on_queue_declared)
def _ch_on_close(self, ch: Channel, reason: AMQPError) -> None:
LOGGER.warning('RabbitMQ channel closed ({})'.format(reason))
try:
self._conn.close()
except Exception:
pass
def _on_queue_declared(self, method: pika.frame.Method) -> None:
self._ch.basic_qos(prefetch_count=self._max_processes,
callback=self._on_basic_qos_ok)
def _on_basic_qos_ok(self, method: pika.frame.Method) -> None:
self._ch.basic_consume(self._queue_name,
on_message_callback=self._recv_message)
LOGGER.info('Worker started')
def _recv_message(self, ch: Channel, method: pika.spec.Basic.Return,
_: pika.spec.BasicProperties, body: bytes) -> None:
try:
self._process(ch, method, body)
except Exception:
LOGGER.warning('', exc_info=True)
def _process(self, ch: Channel, method: pika.spec.Basic.Return,
body: bytes) -> None:
def _done(fut: Optional[Future]) -> None:
ch.basic_ack(delivery_tag=method.delivery_tag)
if fut is None:
return
self._task_processed += 1
if fut.exception() is not None:
self._task_errors += 1
elif fut.result() == JudgeStatus.InternalError:
self._task_errors += 1
try:
contest_id, problem_id, submission_id = pickle.loads(body)
except Exception:
LOGGER.warning('Received invalid message. ignored.', exc_info=True)
_done(None)
return
with transaction() as s:
submission = s.query(Submission).with_for_update().filter(
Submission.contest_id == contest_id,
Submission.problem_id == problem_id,
Submission.id == submission_id).first()
if not submission:
LOGGER.warning(
'Submission.id "{}" is not found'.format(submission_id))
_done(None)
if submission.status not in (JudgeStatus.Waiting,
JudgeStatus.Running,
JudgeStatus.InternalError):
_done(None)
env = s.query(Environment).filter(
Environment.id == submission.environment_id).first()
problem = s.query(Problem).filter(
Problem.contest_id == contest_id,
Problem.id == problem_id).first()
assert env and problem # env/problemは外部キー制約によって常に取得可能
# ワーカーダウン等ですべてのテストのジャッジが完了していない場合は
# ジャッジ済みのテストは結果を流用する
existed_results = {
jr.test_id: jr
for jr in s.query(JudgeResult).filter(
JudgeResult.contest_id == contest_id,
JudgeResult.problem_id == problem_id,
JudgeResult.submission_id == submission_id)
}
task = JudgeTask(id=submission_id,
contest_id=contest_id,
problem_id=problem_id,
user_id=submission.user_id,
code=submission.code,
compile_image_name=env.compile_image_name,
test_image_name=env.test_image_name,
time_limit=problem.time_limit,
memory_limit=problem.memory_limit,
tests=[])
submission.status = JudgeStatus.Running
testcases = s.query(TestCase).filter(
TestCase.contest_id == contest_id,
TestCase.problem_id == problem_id).all()
for test in testcases:
jr = existed_results.get(test.id, None)
if not jr:
s.add(
JudgeResult(contest_id=contest_id,
problem_id=problem_id,
submission_id=submission_id,
test_id=test.id))
if not jr or jr.status in (JudgeStatus.Waiting,
JudgeStatus.Running,
JudgeStatus.InternalError):
task.tests.append(
JudgeTestInfo(id=test.id,
input=test.input,
output=test.output))
# テストの実行順序をシャッフルする
shuffle(task.tests)
def _submit() -> None:
LOGGER.info('submit to child process (submission.id={})'.format(
submission_id))
future = self._executor.submit(run, DockerJudgeDriver, task)
future.add_done_callback(_done)
asyncio.get_event_loop().call_soon_threadsafe(_submit)
def connect(self):
return AsyncioConnection(
on_open_callback=self.on_connect_open,
)
def create_connection(self, parameters):
LOGGER.info('Creating a new connection')
return AsyncioConnection(parameters=parameters,
on_open_callback=self.on_connection_opened,
on_open_error_callback=self.on_connection_open_error,
on_close_callback=self.on_connection_closed)
