You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

369 lines
12 KiB

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