tracing.py 14 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442
  1. from types import SimpleNamespace
  2. from typing import TYPE_CHECKING, Awaitable, Optional, Type, TypeVar
  3. import attr
  4. from multidict import CIMultiDict
  5. from yarl import URL
  6. from .client_reqrep import ClientResponse
  7. from .signals import Signal
  8. if TYPE_CHECKING: # pragma: no cover
  9. from typing_extensions import Protocol
  10. from .client import ClientSession
  11. _ParamT_contra = TypeVar("_ParamT_contra", contravariant=True)
  12. class _SignalCallback(Protocol[_ParamT_contra]):
  13. def __call__(
  14. self,
  15. __client_session: ClientSession,
  16. __trace_config_ctx: SimpleNamespace,
  17. __params: _ParamT_contra,
  18. ) -> Awaitable[None]:
  19. ...
  20. __all__ = (
  21. "TraceConfig",
  22. "TraceRequestStartParams",
  23. "TraceRequestEndParams",
  24. "TraceRequestExceptionParams",
  25. "TraceConnectionQueuedStartParams",
  26. "TraceConnectionQueuedEndParams",
  27. "TraceConnectionCreateStartParams",
  28. "TraceConnectionCreateEndParams",
  29. "TraceConnectionReuseconnParams",
  30. "TraceDnsResolveHostStartParams",
  31. "TraceDnsResolveHostEndParams",
  32. "TraceDnsCacheHitParams",
  33. "TraceDnsCacheMissParams",
  34. "TraceRequestRedirectParams",
  35. "TraceRequestChunkSentParams",
  36. "TraceResponseChunkReceivedParams",
  37. )
  38. class TraceConfig:
  39. """First-class used to trace requests launched via ClientSession
  40. objects."""
  41. def __init__(
  42. self, trace_config_ctx_factory: Type[SimpleNamespace] = SimpleNamespace
  43. ) -> None:
  44. self._on_request_start = Signal(
  45. self
  46. ) # type: Signal[_SignalCallback[TraceRequestStartParams]]
  47. self._on_request_chunk_sent = Signal(
  48. self
  49. ) # type: Signal[_SignalCallback[TraceRequestChunkSentParams]]
  50. self._on_response_chunk_received = Signal(
  51. self
  52. ) # type: Signal[_SignalCallback[TraceResponseChunkReceivedParams]]
  53. self._on_request_end = Signal(
  54. self
  55. ) # type: Signal[_SignalCallback[TraceRequestEndParams]]
  56. self._on_request_exception = Signal(
  57. self
  58. ) # type: Signal[_SignalCallback[TraceRequestExceptionParams]]
  59. self._on_request_redirect = Signal(
  60. self
  61. ) # type: Signal[_SignalCallback[TraceRequestRedirectParams]]
  62. self._on_connection_queued_start = Signal(
  63. self
  64. ) # type: Signal[_SignalCallback[TraceConnectionQueuedStartParams]]
  65. self._on_connection_queued_end = Signal(
  66. self
  67. ) # type: Signal[_SignalCallback[TraceConnectionQueuedEndParams]]
  68. self._on_connection_create_start = Signal(
  69. self
  70. ) # type: Signal[_SignalCallback[TraceConnectionCreateStartParams]]
  71. self._on_connection_create_end = Signal(
  72. self
  73. ) # type: Signal[_SignalCallback[TraceConnectionCreateEndParams]]
  74. self._on_connection_reuseconn = Signal(
  75. self
  76. ) # type: Signal[_SignalCallback[TraceConnectionReuseconnParams]]
  77. self._on_dns_resolvehost_start = Signal(
  78. self
  79. ) # type: Signal[_SignalCallback[TraceDnsResolveHostStartParams]]
  80. self._on_dns_resolvehost_end = Signal(
  81. self
  82. ) # type: Signal[_SignalCallback[TraceDnsResolveHostEndParams]]
  83. self._on_dns_cache_hit = Signal(
  84. self
  85. ) # type: Signal[_SignalCallback[TraceDnsCacheHitParams]]
  86. self._on_dns_cache_miss = Signal(
  87. self
  88. ) # type: Signal[_SignalCallback[TraceDnsCacheMissParams]]
  89. self._trace_config_ctx_factory = trace_config_ctx_factory
  90. def trace_config_ctx(
  91. self, trace_request_ctx: Optional[SimpleNamespace] = None
  92. ) -> SimpleNamespace:
  93. """ Return a new trace_config_ctx instance """
  94. return self._trace_config_ctx_factory(trace_request_ctx=trace_request_ctx)
  95. def freeze(self) -> None:
  96. self._on_request_start.freeze()
  97. self._on_request_chunk_sent.freeze()
  98. self._on_response_chunk_received.freeze()
  99. self._on_request_end.freeze()
  100. self._on_request_exception.freeze()
  101. self._on_request_redirect.freeze()
  102. self._on_connection_queued_start.freeze()
  103. self._on_connection_queued_end.freeze()
  104. self._on_connection_create_start.freeze()
  105. self._on_connection_create_end.freeze()
  106. self._on_connection_reuseconn.freeze()
  107. self._on_dns_resolvehost_start.freeze()
  108. self._on_dns_resolvehost_end.freeze()
  109. self._on_dns_cache_hit.freeze()
  110. self._on_dns_cache_miss.freeze()
  111. @property
  112. def on_request_start(self) -> "Signal[_SignalCallback[TraceRequestStartParams]]":
  113. return self._on_request_start
  114. @property
  115. def on_request_chunk_sent(
  116. self,
  117. ) -> "Signal[_SignalCallback[TraceRequestChunkSentParams]]":
  118. return self._on_request_chunk_sent
  119. @property
  120. def on_response_chunk_received(
  121. self,
  122. ) -> "Signal[_SignalCallback[TraceResponseChunkReceivedParams]]":
  123. return self._on_response_chunk_received
  124. @property
  125. def on_request_end(self) -> "Signal[_SignalCallback[TraceRequestEndParams]]":
  126. return self._on_request_end
  127. @property
  128. def on_request_exception(
  129. self,
  130. ) -> "Signal[_SignalCallback[TraceRequestExceptionParams]]":
  131. return self._on_request_exception
  132. @property
  133. def on_request_redirect(
  134. self,
  135. ) -> "Signal[_SignalCallback[TraceRequestRedirectParams]]":
  136. return self._on_request_redirect
  137. @property
  138. def on_connection_queued_start(
  139. self,
  140. ) -> "Signal[_SignalCallback[TraceConnectionQueuedStartParams]]":
  141. return self._on_connection_queued_start
  142. @property
  143. def on_connection_queued_end(
  144. self,
  145. ) -> "Signal[_SignalCallback[TraceConnectionQueuedEndParams]]":
  146. return self._on_connection_queued_end
  147. @property
  148. def on_connection_create_start(
  149. self,
  150. ) -> "Signal[_SignalCallback[TraceConnectionCreateStartParams]]":
  151. return self._on_connection_create_start
  152. @property
  153. def on_connection_create_end(
  154. self,
  155. ) -> "Signal[_SignalCallback[TraceConnectionCreateEndParams]]":
  156. return self._on_connection_create_end
  157. @property
  158. def on_connection_reuseconn(
  159. self,
  160. ) -> "Signal[_SignalCallback[TraceConnectionReuseconnParams]]":
  161. return self._on_connection_reuseconn
  162. @property
  163. def on_dns_resolvehost_start(
  164. self,
  165. ) -> "Signal[_SignalCallback[TraceDnsResolveHostStartParams]]":
  166. return self._on_dns_resolvehost_start
  167. @property
  168. def on_dns_resolvehost_end(
  169. self,
  170. ) -> "Signal[_SignalCallback[TraceDnsResolveHostEndParams]]":
  171. return self._on_dns_resolvehost_end
  172. @property
  173. def on_dns_cache_hit(self) -> "Signal[_SignalCallback[TraceDnsCacheHitParams]]":
  174. return self._on_dns_cache_hit
  175. @property
  176. def on_dns_cache_miss(self) -> "Signal[_SignalCallback[TraceDnsCacheMissParams]]":
  177. return self._on_dns_cache_miss
  178. @attr.s(auto_attribs=True, frozen=True, slots=True)
  179. class TraceRequestStartParams:
  180. """ Parameters sent by the `on_request_start` signal"""
  181. method: str
  182. url: URL
  183. headers: "CIMultiDict[str]"
  184. @attr.s(auto_attribs=True, frozen=True, slots=True)
  185. class TraceRequestChunkSentParams:
  186. """ Parameters sent by the `on_request_chunk_sent` signal"""
  187. method: str
  188. url: URL
  189. chunk: bytes
  190. @attr.s(auto_attribs=True, frozen=True, slots=True)
  191. class TraceResponseChunkReceivedParams:
  192. """ Parameters sent by the `on_response_chunk_received` signal"""
  193. method: str
  194. url: URL
  195. chunk: bytes
  196. @attr.s(auto_attribs=True, frozen=True, slots=True)
  197. class TraceRequestEndParams:
  198. """ Parameters sent by the `on_request_end` signal"""
  199. method: str
  200. url: URL
  201. headers: "CIMultiDict[str]"
  202. response: ClientResponse
  203. @attr.s(auto_attribs=True, frozen=True, slots=True)
  204. class TraceRequestExceptionParams:
  205. """ Parameters sent by the `on_request_exception` signal"""
  206. method: str
  207. url: URL
  208. headers: "CIMultiDict[str]"
  209. exception: BaseException
  210. @attr.s(auto_attribs=True, frozen=True, slots=True)
  211. class TraceRequestRedirectParams:
  212. """ Parameters sent by the `on_request_redirect` signal"""
  213. method: str
  214. url: URL
  215. headers: "CIMultiDict[str]"
  216. response: ClientResponse
  217. @attr.s(auto_attribs=True, frozen=True, slots=True)
  218. class TraceConnectionQueuedStartParams:
  219. """ Parameters sent by the `on_connection_queued_start` signal"""
  220. @attr.s(auto_attribs=True, frozen=True, slots=True)
  221. class TraceConnectionQueuedEndParams:
  222. """ Parameters sent by the `on_connection_queued_end` signal"""
  223. @attr.s(auto_attribs=True, frozen=True, slots=True)
  224. class TraceConnectionCreateStartParams:
  225. """ Parameters sent by the `on_connection_create_start` signal"""
  226. @attr.s(auto_attribs=True, frozen=True, slots=True)
  227. class TraceConnectionCreateEndParams:
  228. """ Parameters sent by the `on_connection_create_end` signal"""
  229. @attr.s(auto_attribs=True, frozen=True, slots=True)
  230. class TraceConnectionReuseconnParams:
  231. """ Parameters sent by the `on_connection_reuseconn` signal"""
  232. @attr.s(auto_attribs=True, frozen=True, slots=True)
  233. class TraceDnsResolveHostStartParams:
  234. """ Parameters sent by the `on_dns_resolvehost_start` signal"""
  235. host: str
  236. @attr.s(auto_attribs=True, frozen=True, slots=True)
  237. class TraceDnsResolveHostEndParams:
  238. """ Parameters sent by the `on_dns_resolvehost_end` signal"""
  239. host: str
  240. @attr.s(auto_attribs=True, frozen=True, slots=True)
  241. class TraceDnsCacheHitParams:
  242. """ Parameters sent by the `on_dns_cache_hit` signal"""
  243. host: str
  244. @attr.s(auto_attribs=True, frozen=True, slots=True)
  245. class TraceDnsCacheMissParams:
  246. """ Parameters sent by the `on_dns_cache_miss` signal"""
  247. host: str
  248. class Trace:
  249. """Internal class used to keep together the main dependencies used
  250. at the moment of send a signal."""
  251. def __init__(
  252. self,
  253. session: "ClientSession",
  254. trace_config: TraceConfig,
  255. trace_config_ctx: SimpleNamespace,
  256. ) -> None:
  257. self._trace_config = trace_config
  258. self._trace_config_ctx = trace_config_ctx
  259. self._session = session
  260. async def send_request_start(
  261. self, method: str, url: URL, headers: "CIMultiDict[str]"
  262. ) -> None:
  263. return await self._trace_config.on_request_start.send(
  264. self._session,
  265. self._trace_config_ctx,
  266. TraceRequestStartParams(method, url, headers),
  267. )
  268. async def send_request_chunk_sent(
  269. self, method: str, url: URL, chunk: bytes
  270. ) -> None:
  271. return await self._trace_config.on_request_chunk_sent.send(
  272. self._session,
  273. self._trace_config_ctx,
  274. TraceRequestChunkSentParams(method, url, chunk),
  275. )
  276. async def send_response_chunk_received(
  277. self, method: str, url: URL, chunk: bytes
  278. ) -> None:
  279. return await self._trace_config.on_response_chunk_received.send(
  280. self._session,
  281. self._trace_config_ctx,
  282. TraceResponseChunkReceivedParams(method, url, chunk),
  283. )
  284. async def send_request_end(
  285. self,
  286. method: str,
  287. url: URL,
  288. headers: "CIMultiDict[str]",
  289. response: ClientResponse,
  290. ) -> None:
  291. return await self._trace_config.on_request_end.send(
  292. self._session,
  293. self._trace_config_ctx,
  294. TraceRequestEndParams(method, url, headers, response),
  295. )
  296. async def send_request_exception(
  297. self,
  298. method: str,
  299. url: URL,
  300. headers: "CIMultiDict[str]",
  301. exception: BaseException,
  302. ) -> None:
  303. return await self._trace_config.on_request_exception.send(
  304. self._session,
  305. self._trace_config_ctx,
  306. TraceRequestExceptionParams(method, url, headers, exception),
  307. )
  308. async def send_request_redirect(
  309. self,
  310. method: str,
  311. url: URL,
  312. headers: "CIMultiDict[str]",
  313. response: ClientResponse,
  314. ) -> None:
  315. return await self._trace_config._on_request_redirect.send(
  316. self._session,
  317. self._trace_config_ctx,
  318. TraceRequestRedirectParams(method, url, headers, response),
  319. )
  320. async def send_connection_queued_start(self) -> None:
  321. return await self._trace_config.on_connection_queued_start.send(
  322. self._session, self._trace_config_ctx, TraceConnectionQueuedStartParams()
  323. )
  324. async def send_connection_queued_end(self) -> None:
  325. return await self._trace_config.on_connection_queued_end.send(
  326. self._session, self._trace_config_ctx, TraceConnectionQueuedEndParams()
  327. )
  328. async def send_connection_create_start(self) -> None:
  329. return await self._trace_config.on_connection_create_start.send(
  330. self._session, self._trace_config_ctx, TraceConnectionCreateStartParams()
  331. )
  332. async def send_connection_create_end(self) -> None:
  333. return await self._trace_config.on_connection_create_end.send(
  334. self._session, self._trace_config_ctx, TraceConnectionCreateEndParams()
  335. )
  336. async def send_connection_reuseconn(self) -> None:
  337. return await self._trace_config.on_connection_reuseconn.send(
  338. self._session, self._trace_config_ctx, TraceConnectionReuseconnParams()
  339. )
  340. async def send_dns_resolvehost_start(self, host: str) -> None:
  341. return await self._trace_config.on_dns_resolvehost_start.send(
  342. self._session, self._trace_config_ctx, TraceDnsResolveHostStartParams(host)
  343. )
  344. async def send_dns_resolvehost_end(self, host: str) -> None:
  345. return await self._trace_config.on_dns_resolvehost_end.send(
  346. self._session, self._trace_config_ctx, TraceDnsResolveHostEndParams(host)
  347. )
  348. async def send_dns_cache_hit(self, host: str) -> None:
  349. return await self._trace_config.on_dns_cache_hit.send(
  350. self._session, self._trace_config_ctx, TraceDnsCacheHitParams(host)
  351. )
  352. async def send_dns_cache_miss(self, host: str) -> None:
  353. return await self._trace_config.on_dns_cache_miss.send(
  354. self._session, self._trace_config_ctx, TraceDnsCacheMissParams(host)
  355. )