def _set_span_err(self, err: Exception) -> None:
if not span:
return
span.tag('error', 'true')
span.annotate(span.ANN_TRACEBACK, traceback.format_exc())
if hasattr(err, 'jsonrpc_error_code'):
span.tag(
SPAN_TAG_JSONRPC_CODE,
str(err.jsonrpc_error_code), # type: ignore
)
else:
code, _, _ = _get_code_message_and_data(err)
span.tag(SPAN_TAG_JSONRPC_CODE, str(code))
async def _transport(self, request: bytes,
timeout: Optional[float]) -> bytes:
correlation_id = str(uuid.uuid4())
fut: asyncio.Future = asyncio.Future()
self._futs[correlation_id] = (fut, span)
try:
with self.amqp.app.logger.capture_span(AmqpSpan) as trap:
headers: Dict[str, str] = {}
if self.cfg.propagate_trace and span:
headers = span.to_headers()
await self.publish(
'',
self.cfg.queue,
request,
Properties(
correlation_id=correlation_id,
reply_to=self._queue,
headers=headers,
),
propagate_trace=False,
)
if span:
anns = trap.span.annotations.get(AmqpSpan.ANN_OUT_PROPS)
if anns is not None and len(anns) > 0:
ann_body, ann_stamp = anns[0]
span.annotate(AmqpSpan.ANN_IN_PROPS, ann_body,
ann_stamp)
anns = trap.span.annotations.get(AmqpSpan.ANN_OUT_BODY)
if anns is not None and len(anns) > 0:
ann_body, ann_stamp = anns[0]
span.annotate(AmqpSpan.ANN_IN_BODY, ann_body,
ann_stamp)
trap.span.copy_to(span,
annotations=False,
tags=True,
error=True)
trap.span.skip()
return await asyncio.wait_for(fut,
timeout=timeout or self.cfg.timeout)
finally:
del self._futs[correlation_id]
async def _scan(self) -> None:
if self.app is None or self._lock is None: # pragma: no cover
raise UserWarning
while True:
if self._stopping:
return
async with self._lock:
delay = 1.0 # default: 1 second
try:
with wrap2span(
name=TaskManagerSpan.NAME_SCAN,
kind=Span.KIND_SERVER,
# ignore_ctx=True,
cls=TaskManagerSpan,
app=self.app,
) as span:
try:
tasks, delay = await self._search_and_exec()
if len(tasks) == 0:
span.skip()
except Exception as err:
span.error(err)
self.app.log_err(err)
finally:
if not self._stopping:
span.annotate(
TaskManagerSpan.ANN_NEXT_SCAN,
'next: %s' % delay,
)
finally:
if not self._stopping:
self._scan_fut = None
eta = time() + delay
self.stamp_early = eta
sleep = self.stamp_early - time()
if sleep > 0:
try:
self._scan_sleep_fut = asyncio.create_task(
asyncio.sleep(sleep))
await self._scan_sleep_fut
except asyncio.CancelledError:
pass
self._scan_sleep_fut = None
async def _search_and_exec(self) -> Tuple[List[Task], float]:
if self._db is None: # pragma: no cover
raise UserWarning
async with self._db.transaction():
tasks = await self._db.task_search(self.cfg.batch_size, lock=False)
span.annotate(TaskManagerSpan.ANN_TASKS, repr(tasks))
if len(tasks) == 0:
next_delay = await self._db.task_next_delay(lock=False)
if (next_delay is None
or next_delay >= self.cfg.max_scan_interval):
return tasks, self.cfg.max_scan_interval
if next_delay <= 0:
return tasks, 0
return tasks, next_delay
coros = [self._exec(span.trace_id, task) for task in tasks]
await asyncio.gather(*coros)
return tasks, 0
async def _scan(self) -> List[int]:
if self.app is None or self._lock is None: # pragma: no cover
raise UserWarning
if self._stopping:
return []
async with self._lock:
delay = 1.0 # default: 1 second
try:
with wrap2span(
name=TaskManagerSpan.NAME_SCAN,
kind=Span.KIND_SERVER,
# ignore_ctx=True,
cls=TaskManagerSpan,
app=self.app,
) as span:
try:
tasks, delay = await self._search_and_exec()
if len(tasks) == 0:
span.skip()
return [task.id for task in tasks]
except Exception as err:
span.error(err)
self.app.log_err(err)
finally:
if not self._stopping:
span.annotate(
TaskManagerSpan.ANN_NEXT_SCAN,
'next: %s' % delay,
)
return []
finally:
if not self._stopping:
self._scan_fut = None
eta = self.loop.time() + delay
self.stamp_early = eta
self.loop.call_at(eta, self._scan_later, eta)
async def read_root() -> dict:
span.annotate('k1', 'hello world')
return {"Hello": "World"}
async def schedule(
self,
func: TaskHandler,
params: dict,
reference: Optional[str] = None,
eta: Optional[ETA] = None,
max_retries: Optional[int] = None,
retry_delay: Optional[float] = None,
propagate_trace: bool = False,
) -> int:
with wrap2span(
name=TaskManagerSpan.NAME_SCHEDULE,
kind=Span.KIND_CLIENT,
cls=TaskManagerSpan,
app=self.app,
) as span:
if self._db is None: # pragma: no cover
raise UserWarning
if not isinstance(func, str):
if not hasattr(func, '__rpc_name__'): # pragma: no cover
raise UserWarning('Invalid task handler')
func_name = getattr(func, '__rpc_name__')
else:
func_name = func
if max_retries is None:
max_retries = getattr(func, '__task_max_retries__', 0)
if retry_delay is None:
retry_delay = getattr(func, '__task_retry_delay__', 60.0)
span.name = '%s::%s' % (TaskManagerSpan.NAME_SCHEDULE, func_name)
eta_dt: Optional[datetime] = None
if isinstance(eta, int) or isinstance(eta, float):
eta_dt = datetime.fromtimestamp(eta, tz=timezone.utc)
elif isinstance(eta, datetime):
eta_dt = eta
elif eta is not None: # pragma: no cover
raise UserWarning
if eta_dt is not None:
span.annotate(TaskManagerSpan.ANN_ETA,
'ETA: %s' % eta_dt.isoformat())
add_params: List[Any] = [
eta_dt,
func_name,
params,
reference,
max_retries,
retry_delay,
]
if propagate_trace:
add_params.append(span.trace_id)
add_params.append(span.id)
task_id, task_delay = await self._db.task_add(
*add_params,
lock=True,
)
span.annotate(TaskManagerSpan.ANN_DELAY, 'Delay: %s' % task_delay)
eta_float = self.loop.time() + task_delay
self.stamp_early = eta_float
self.loop.call_at(eta_float, self._scan_later, eta_float)
return task_id
async def send_wraper(self, send: Callable[..., Awaitable],
event: Any) -> None:
self._i += 1
span.annotate(str(self._i), str(event))
await send(event)
async def receive_wraper(self, receive: Callable[..., Awaitable]) -> Any:
data = await receive()
self._i += 1
span.annotate(str(self._i), str(data))
return data
