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.

631 lines
20 KiB

4 years ago
  1. import asyncio
  2. import collections
  3. import warnings
  4. from typing import List # noqa
  5. from typing import Awaitable, Callable, Generic, Optional, Tuple, TypeVar
  6. from .base_protocol import BaseProtocol
  7. from .helpers import BaseTimerContext, set_exception, set_result
  8. from .log import internal_logger
  9. try: # pragma: no cover
  10. from typing import Deque # noqa
  11. except ImportError:
  12. from typing_extensions import Deque # noqa
  13. __all__ = (
  14. 'EMPTY_PAYLOAD', 'EofStream', 'StreamReader', 'DataQueue',
  15. 'FlowControlDataQueue')
  16. DEFAULT_LIMIT = 2 ** 16
  17. _T = TypeVar('_T')
  18. class EofStream(Exception):
  19. """eof stream indication."""
  20. class AsyncStreamIterator(Generic[_T]):
  21. def __init__(self, read_func: Callable[[], Awaitable[_T]]) -> None:
  22. self.read_func = read_func
  23. def __aiter__(self) -> 'AsyncStreamIterator[_T]':
  24. return self
  25. async def __anext__(self) -> _T:
  26. try:
  27. rv = await self.read_func()
  28. except EofStream:
  29. raise StopAsyncIteration # NOQA
  30. if rv == b'':
  31. raise StopAsyncIteration # NOQA
  32. return rv
  33. class ChunkTupleAsyncStreamIterator:
  34. def __init__(self, stream: 'StreamReader') -> None:
  35. self._stream = stream
  36. def __aiter__(self) -> 'ChunkTupleAsyncStreamIterator':
  37. return self
  38. async def __anext__(self) -> Tuple[bytes, bool]:
  39. rv = await self._stream.readchunk()
  40. if rv == (b'', False):
  41. raise StopAsyncIteration # NOQA
  42. return rv
  43. class AsyncStreamReaderMixin:
  44. def __aiter__(self) -> AsyncStreamIterator[bytes]:
  45. return AsyncStreamIterator(self.readline) # type: ignore
  46. def iter_chunked(self, n: int) -> AsyncStreamIterator[bytes]:
  47. """Returns an asynchronous iterator that yields chunks of size n.
  48. Python-3.5 available for Python 3.5+ only
  49. """
  50. return AsyncStreamIterator(lambda: self.read(n)) # type: ignore
  51. def iter_any(self) -> AsyncStreamIterator[bytes]:
  52. """Returns an asynchronous iterator that yields all the available
  53. data as soon as it is received
  54. Python-3.5 available for Python 3.5+ only
  55. """
  56. return AsyncStreamIterator(self.readany) # type: ignore
  57. def iter_chunks(self) -> ChunkTupleAsyncStreamIterator:
  58. """Returns an asynchronous iterator that yields chunks of data
  59. as they are received by the server. The yielded objects are tuples
  60. of (bytes, bool) as returned by the StreamReader.readchunk method.
  61. Python-3.5 available for Python 3.5+ only
  62. """
  63. return ChunkTupleAsyncStreamIterator(self) # type: ignore
  64. class StreamReader(AsyncStreamReaderMixin):
  65. """An enhancement of asyncio.StreamReader.
  66. Supports asynchronous iteration by line, chunk or as available::
  67. async for line in reader:
  68. ...
  69. async for chunk in reader.iter_chunked(1024):
  70. ...
  71. async for slice in reader.iter_any():
  72. ...
  73. """
  74. total_bytes = 0
  75. def __init__(self, protocol: BaseProtocol,
  76. *, limit: int=DEFAULT_LIMIT,
  77. timer: Optional[BaseTimerContext]=None,
  78. loop: Optional[asyncio.AbstractEventLoop]=None) -> None:
  79. self._protocol = protocol
  80. self._low_water = limit
  81. self._high_water = limit * 2
  82. if loop is None:
  83. loop = asyncio.get_event_loop()
  84. self._loop = loop
  85. self._size = 0
  86. self._cursor = 0
  87. self._http_chunk_splits = None # type: Optional[List[int]]
  88. self._buffer = collections.deque() # type: Deque[bytes]
  89. self._buffer_offset = 0
  90. self._eof = False
  91. self._waiter = None # type: Optional[asyncio.Future[bool]]
  92. self._eof_waiter = None # type: Optional[asyncio.Future[bool]]
  93. self._exception = None # type: Optional[BaseException]
  94. self._timer = timer
  95. self._eof_callbacks = [] # type: List[Callable[[], None]]
  96. def __repr__(self) -> str:
  97. info = [self.__class__.__name__]
  98. if self._size:
  99. info.append('%d bytes' % self._size)
  100. if self._eof:
  101. info.append('eof')
  102. if self._low_water != DEFAULT_LIMIT:
  103. info.append('low=%d high=%d' % (self._low_water, self._high_water))
  104. if self._waiter:
  105. info.append('w=%r' % self._waiter)
  106. if self._exception:
  107. info.append('e=%r' % self._exception)
  108. return '<%s>' % ' '.join(info)
  109. def exception(self) -> Optional[BaseException]:
  110. return self._exception
  111. def set_exception(self, exc: BaseException) -> None:
  112. self._exception = exc
  113. self._eof_callbacks.clear()
  114. waiter = self._waiter
  115. if waiter is not None:
  116. self._waiter = None
  117. set_exception(waiter, exc)
  118. waiter = self._eof_waiter
  119. if waiter is not None:
  120. set_exception(waiter, exc)
  121. self._eof_waiter = None
  122. def on_eof(self, callback: Callable[[], None]) -> None:
  123. if self._eof:
  124. try:
  125. callback()
  126. except Exception:
  127. internal_logger.exception('Exception in eof callback')
  128. else:
  129. self._eof_callbacks.append(callback)
  130. def feed_eof(self) -> None:
  131. self._eof = True
  132. waiter = self._waiter
  133. if waiter is not None:
  134. self._waiter = None
  135. set_result(waiter, True)
  136. waiter = self._eof_waiter
  137. if waiter is not None:
  138. self._eof_waiter = None
  139. set_result(waiter, True)
  140. for cb in self._eof_callbacks:
  141. try:
  142. cb()
  143. except Exception:
  144. internal_logger.exception('Exception in eof callback')
  145. self._eof_callbacks.clear()
  146. def is_eof(self) -> bool:
  147. """Return True if 'feed_eof' was called."""
  148. return self._eof
  149. def at_eof(self) -> bool:
  150. """Return True if the buffer is empty and 'feed_eof' was called."""
  151. return self._eof and not self._buffer
  152. async def wait_eof(self) -> None:
  153. if self._eof:
  154. return
  155. assert self._eof_waiter is None
  156. self._eof_waiter = self._loop.create_future()
  157. try:
  158. await self._eof_waiter
  159. finally:
  160. self._eof_waiter = None
  161. def unread_data(self, data: bytes) -> None:
  162. """ rollback reading some data from stream, inserting it to buffer head.
  163. """
  164. warnings.warn("unread_data() is deprecated "
  165. "and will be removed in future releases (#3260)",
  166. DeprecationWarning,
  167. stacklevel=2)
  168. if not data:
  169. return
  170. if self._buffer_offset:
  171. self._buffer[0] = self._buffer[0][self._buffer_offset:]
  172. self._buffer_offset = 0
  173. self._size += len(data)
  174. self._cursor -= len(data)
  175. self._buffer.appendleft(data)
  176. self._eof_counter = 0
  177. # TODO: size is ignored, remove the param later
  178. def feed_data(self, data: bytes, size: int=0) -> None:
  179. assert not self._eof, 'feed_data after feed_eof'
  180. if not data:
  181. return
  182. self._size += len(data)
  183. self._buffer.append(data)
  184. self.total_bytes += len(data)
  185. waiter = self._waiter
  186. if waiter is not None:
  187. self._waiter = None
  188. set_result(waiter, False)
  189. if (self._size > self._high_water and
  190. not self._protocol._reading_paused):
  191. self._protocol.pause_reading()
  192. def begin_http_chunk_receiving(self) -> None:
  193. if self._http_chunk_splits is None:
  194. if self.total_bytes:
  195. raise RuntimeError("Called begin_http_chunk_receiving when"
  196. "some data was already fed")
  197. self._http_chunk_splits = []
  198. def end_http_chunk_receiving(self) -> None:
  199. if self._http_chunk_splits is None:
  200. raise RuntimeError("Called end_chunk_receiving without calling "
  201. "begin_chunk_receiving first")
  202. # self._http_chunk_splits contains logical byte offsets from start of
  203. # the body transfer. Each offset is the offset of the end of a chunk.
  204. # "Logical" means bytes, accessible for a user.
  205. # If no chunks containig logical data were received, current position
  206. # is difinitely zero.
  207. pos = self._http_chunk_splits[-1] if self._http_chunk_splits else 0
  208. if self.total_bytes == pos:
  209. # We should not add empty chunks here. So we check for that.
  210. # Note, when chunked + gzip is used, we can receive a chunk
  211. # of compressed data, but that data may not be enough for gzip FSM
  212. # to yield any uncompressed data. That's why current position may
  213. # not change after receiving a chunk.
  214. return
  215. self._http_chunk_splits.append(self.total_bytes)
  216. # wake up readchunk when end of http chunk received
  217. waiter = self._waiter
  218. if waiter is not None:
  219. self._waiter = None
  220. set_result(waiter, False)
  221. async def _wait(self, func_name: str) -> None:
  222. # StreamReader uses a future to link the protocol feed_data() method
  223. # to a read coroutine. Running two read coroutines at the same time
  224. # would have an unexpected behaviour. It would not possible to know
  225. # which coroutine would get the next data.
  226. if self._waiter is not None:
  227. raise RuntimeError('%s() called while another coroutine is '
  228. 'already waiting for incoming data' % func_name)
  229. waiter = self._waiter = self._loop.create_future()
  230. try:
  231. if self._timer:
  232. with self._timer:
  233. await waiter
  234. else:
  235. await waiter
  236. finally:
  237. self._waiter = None
  238. async def readline(self) -> bytes:
  239. if self._exception is not None:
  240. raise self._exception
  241. line = []
  242. line_size = 0
  243. not_enough = True
  244. while not_enough:
  245. while self._buffer and not_enough:
  246. offset = self._buffer_offset
  247. ichar = self._buffer[0].find(b'\n', offset) + 1
  248. # Read from current offset to found b'\n' or to the end.
  249. data = self._read_nowait_chunk(ichar - offset if ichar else -1)
  250. line.append(data)
  251. line_size += len(data)
  252. if ichar:
  253. not_enough = False
  254. if line_size > self._high_water:
  255. raise ValueError('Line is too long')
  256. if self._eof:
  257. break
  258. if not_enough:
  259. await self._wait('readline')
  260. return b''.join(line)
  261. async def read(self, n: int=-1) -> bytes:
  262. if self._exception is not None:
  263. raise self._exception
  264. # migration problem; with DataQueue you have to catch
  265. # EofStream exception, so common way is to run payload.read() inside
  266. # infinite loop. what can cause real infinite loop with StreamReader
  267. # lets keep this code one major release.
  268. if __debug__:
  269. if self._eof and not self._buffer:
  270. self._eof_counter = getattr(self, '_eof_counter', 0) + 1
  271. if self._eof_counter > 5:
  272. internal_logger.warning(
  273. 'Multiple access to StreamReader in eof state, '
  274. 'might be infinite loop.', stack_info=True)
  275. if not n:
  276. return b''
  277. if n < 0:
  278. # This used to just loop creating a new waiter hoping to
  279. # collect everything in self._buffer, but that would
  280. # deadlock if the subprocess sends more than self.limit
  281. # bytes. So just call self.readany() until EOF.
  282. blocks = []
  283. while True:
  284. block = await self.readany()
  285. if not block:
  286. break
  287. blocks.append(block)
  288. return b''.join(blocks)
  289. # TODO: should be `if` instead of `while`
  290. # because waiter maybe triggered on chunk end,
  291. # without feeding any data
  292. while not self._buffer and not self._eof:
  293. await self._wait('read')
  294. return self._read_nowait(n)
  295. async def readany(self) -> bytes:
  296. if self._exception is not None:
  297. raise self._exception
  298. # TODO: should be `if` instead of `while`
  299. # because waiter maybe triggered on chunk end,
  300. # without feeding any data
  301. while not self._buffer and not self._eof:
  302. await self._wait('readany')
  303. return self._read_nowait(-1)
  304. async def readchunk(self) -> Tuple[bytes, bool]:
  305. """Returns a tuple of (data, end_of_http_chunk). When chunked transfer
  306. encoding is used, end_of_http_chunk is a boolean indicating if the end
  307. of the data corresponds to the end of a HTTP chunk , otherwise it is
  308. always False.
  309. """
  310. if self._exception is not None:
  311. raise self._exception
  312. if not self._buffer and not self._eof:
  313. if (self._http_chunk_splits and
  314. self._cursor == self._http_chunk_splits[0]):
  315. # end of http chunk without available data
  316. self._http_chunk_splits = self._http_chunk_splits[1:]
  317. return (b"", True)
  318. await self._wait('readchunk')
  319. if not self._buffer and not self._http_chunk_splits:
  320. # end of file
  321. return (b"", False)
  322. elif self._http_chunk_splits is not None:
  323. while self._http_chunk_splits:
  324. pos = self._http_chunk_splits[0]
  325. self._http_chunk_splits = self._http_chunk_splits[1:]
  326. if pos == self._cursor:
  327. return (b"", True)
  328. if pos > self._cursor:
  329. return (self._read_nowait(pos-self._cursor), True)
  330. return (self._read_nowait(-1), False)
  331. else:
  332. return (self._read_nowait_chunk(-1), False)
  333. async def readexactly(self, n: int) -> bytes:
  334. if self._exception is not None:
  335. raise self._exception
  336. blocks = [] # type: List[bytes]
  337. while n > 0:
  338. block = await self.read(n)
  339. if not block:
  340. partial = b''.join(blocks)
  341. raise asyncio.streams.IncompleteReadError(
  342. partial, len(partial) + n)
  343. blocks.append(block)
  344. n -= len(block)
  345. return b''.join(blocks)
  346. def read_nowait(self, n: int=-1) -> bytes:
  347. # default was changed to be consistent with .read(-1)
  348. #
  349. # I believe the most users don't know about the method and
  350. # they are not affected.
  351. if self._exception is not None:
  352. raise self._exception
  353. if self._waiter and not self._waiter.done():
  354. raise RuntimeError(
  355. 'Called while some coroutine is waiting for incoming data.')
  356. return self._read_nowait(n)
  357. def _read_nowait_chunk(self, n: int) -> bytes:
  358. first_buffer = self._buffer[0]
  359. offset = self._buffer_offset
  360. if n != -1 and len(first_buffer) - offset > n:
  361. data = first_buffer[offset:offset + n]
  362. self._buffer_offset += n
  363. elif offset:
  364. self._buffer.popleft()
  365. data = first_buffer[offset:]
  366. self._buffer_offset = 0
  367. else:
  368. data = self._buffer.popleft()
  369. self._size -= len(data)
  370. self._cursor += len(data)
  371. if self._size < self._low_water and self._protocol._reading_paused:
  372. self._protocol.resume_reading()
  373. return data
  374. def _read_nowait(self, n: int) -> bytes:
  375. chunks = []
  376. while self._buffer:
  377. chunk = self._read_nowait_chunk(n)
  378. chunks.append(chunk)
  379. if n != -1:
  380. n -= len(chunk)
  381. if n == 0:
  382. break
  383. return b''.join(chunks) if chunks else b''
  384. class EmptyStreamReader(AsyncStreamReaderMixin):
  385. def exception(self) -> Optional[BaseException]:
  386. return None
  387. def set_exception(self, exc: BaseException) -> None:
  388. pass
  389. def on_eof(self, callback: Callable[[], None]) -> None:
  390. try:
  391. callback()
  392. except Exception:
  393. internal_logger.exception('Exception in eof callback')
  394. def feed_eof(self) -> None:
  395. pass
  396. def is_eof(self) -> bool:
  397. return True
  398. def at_eof(self) -> bool:
  399. return True
  400. async def wait_eof(self) -> None:
  401. return
  402. def feed_data(self, data: bytes, n: int=0) -> None:
  403. pass
  404. async def readline(self) -> bytes:
  405. return b''
  406. async def read(self, n: int=-1) -> bytes:
  407. return b''
  408. async def readany(self) -> bytes:
  409. return b''
  410. async def readchunk(self) -> Tuple[bytes, bool]:
  411. return (b'', True)
  412. async def readexactly(self, n: int) -> bytes:
  413. raise asyncio.streams.IncompleteReadError(b'', n)
  414. def read_nowait(self) -> bytes:
  415. return b''
  416. EMPTY_PAYLOAD = EmptyStreamReader()
  417. class DataQueue(Generic[_T]):
  418. """DataQueue is a general-purpose blocking queue with one reader."""
  419. def __init__(self, loop: asyncio.AbstractEventLoop) -> None:
  420. self._loop = loop
  421. self._eof = False
  422. self._waiter = None # type: Optional[asyncio.Future[bool]]
  423. self._exception = None # type: Optional[BaseException]
  424. self._size = 0
  425. self._buffer = collections.deque() # type: Deque[Tuple[_T, int]]
  426. def __len__(self) -> int:
  427. return len(self._buffer)
  428. def is_eof(self) -> bool:
  429. return self._eof
  430. def at_eof(self) -> bool:
  431. return self._eof and not self._buffer
  432. def exception(self) -> Optional[BaseException]:
  433. return self._exception
  434. def set_exception(self, exc: BaseException) -> None:
  435. self._eof = True
  436. self._exception = exc
  437. waiter = self._waiter
  438. if waiter is not None:
  439. set_exception(waiter, exc)
  440. self._waiter = None
  441. def feed_data(self, data: _T, size: int=0) -> None:
  442. self._size += size
  443. self._buffer.append((data, size))
  444. waiter = self._waiter
  445. if waiter is not None:
  446. self._waiter = None
  447. set_result(waiter, True)
  448. def feed_eof(self) -> None:
  449. self._eof = True
  450. waiter = self._waiter
  451. if waiter is not None:
  452. self._waiter = None
  453. set_result(waiter, False)
  454. async def read(self) -> _T:
  455. if not self._buffer and not self._eof:
  456. assert not self._waiter
  457. self._waiter = self._loop.create_future()
  458. try:
  459. await self._waiter
  460. except (asyncio.CancelledError, asyncio.TimeoutError):
  461. self._waiter = None
  462. raise
  463. if self._buffer:
  464. data, size = self._buffer.popleft()
  465. self._size -= size
  466. return data
  467. else:
  468. if self._exception is not None:
  469. raise self._exception
  470. else:
  471. raise EofStream
  472. def __aiter__(self) -> AsyncStreamIterator[_T]:
  473. return AsyncStreamIterator(self.read)
  474. class FlowControlDataQueue(DataQueue[_T]):
  475. """FlowControlDataQueue resumes and pauses an underlying stream.
  476. It is a destination for parsed data."""
  477. def __init__(self, protocol: BaseProtocol, *,
  478. limit: int=DEFAULT_LIMIT,
  479. loop: asyncio.AbstractEventLoop) -> None:
  480. super().__init__(loop=loop)
  481. self._protocol = protocol
  482. self._limit = limit * 2
  483. def feed_data(self, data: _T, size: int=0) -> None:
  484. super().feed_data(data, size)
  485. if self._size > self._limit and not self._protocol._reading_paused:
  486. self._protocol.pause_reading()
  487. async def read(self) -> _T:
  488. try:
  489. return await super().read()
  490. finally:
  491. if self._size < self._limit and self._protocol._reading_paused:
  492. self._protocol.resume_reading()