streams.py 20 KB

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