tracing.py 14 KB

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