1""" 2Generic Asynchronous Message-based Protocol Support 3 4This module provides a generic framework for sending and receiving 5messages over an asyncio stream. `AsyncProtocol` is an abstract class 6that implements the core mechanisms of a simple send/receive protocol, 7and is designed to be extended. 8 9In this package, it is used as the implementation for the `QMPClient` 10class. 11""" 12 13# It's all the docstrings ... ! It's long for a good reason ^_^; 14# pylint: disable=too-many-lines 15 16import asyncio 17from asyncio import StreamReader, StreamWriter 18from enum import Enum 19from functools import wraps 20import logging 21import socket 22from ssl import SSLContext 23from typing import ( 24 Any, 25 Awaitable, 26 Callable, 27 Generic, 28 List, 29 Optional, 30 Tuple, 31 TypeVar, 32 Union, 33 cast, 34) 35 36from .error import QMPError 37from .util import ( 38 bottom_half, 39 create_task, 40 exception_summary, 41 flush, 42 is_closing, 43 pretty_traceback, 44 upper_half, 45 wait_closed, 46) 47 48 49T = TypeVar('T') 50_U = TypeVar('_U') 51_TaskFN = Callable[[], Awaitable[None]] # aka ``async def func() -> None`` 52 53InternetAddrT = Tuple[str, int] 54UnixAddrT = str 55SocketAddrT = Union[UnixAddrT, InternetAddrT] 56 57 58class Runstate(Enum): 59 """Protocol session runstate.""" 60 61 #: Fully quiesced and disconnected. 62 IDLE = 0 63 #: In the process of connecting or establishing a session. 64 CONNECTING = 1 65 #: Fully connected and active session. 66 RUNNING = 2 67 #: In the process of disconnecting. 68 #: Runstate may be returned to `IDLE` by calling `disconnect()`. 69 DISCONNECTING = 3 70 71 72class ConnectError(QMPError): 73 """ 74 Raised when the initial connection process has failed. 75 76 This Exception always wraps a "root cause" exception that can be 77 interrogated for additional information. 78 79 :param error_message: Human-readable string describing the error. 80 :param exc: The root-cause exception. 81 """ 82 def __init__(self, error_message: str, exc: Exception): 83 super().__init__(error_message) 84 #: Human-readable error string 85 self.error_message: str = error_message 86 #: Wrapped root cause exception 87 self.exc: Exception = exc 88 89 def __str__(self) -> str: 90 cause = str(self.exc) 91 if not cause: 92 # If there's no error string, use the exception name. 93 cause = exception_summary(self.exc) 94 return f"{self.error_message}: {cause}" 95 96 97class StateError(QMPError): 98 """ 99 An API command (connect, execute, etc) was issued at an inappropriate time. 100 101 This error is raised when a command like 102 :py:meth:`~AsyncProtocol.connect()` is issued at an inappropriate 103 time. 104 105 :param error_message: Human-readable string describing the state violation. 106 :param state: The actual `Runstate` seen at the time of the violation. 107 :param required: The `Runstate` required to process this command. 108 """ 109 def __init__(self, error_message: str, 110 state: Runstate, required: Runstate): 111 super().__init__(error_message) 112 self.error_message = error_message 113 self.state = state 114 self.required = required 115 116 117F = TypeVar('F', bound=Callable[..., Any]) # pylint: disable=invalid-name 118 119 120# Don't Panic. 121def require(required_state: Runstate) -> Callable[[F], F]: 122 """ 123 Decorator: protect a method so it can only be run in a certain `Runstate`. 124 125 :param required_state: The `Runstate` required to invoke this method. 126 :raise StateError: When the required `Runstate` is not met. 127 """ 128 def _decorator(func: F) -> F: 129 # _decorator is the decorator that is built by calling the 130 # require() decorator factory; e.g.: 131 # 132 # @require(Runstate.IDLE) def foo(): ... 133 # will replace 'foo' with the result of '_decorator(foo)'. 134 135 @wraps(func) 136 def _wrapper(proto: 'AsyncProtocol[Any]', 137 *args: Any, **kwargs: Any) -> Any: 138 # _wrapper is the function that gets executed prior to the 139 # decorated method. 140 141 name = type(proto).__name__ 142 143 if proto.runstate != required_state: 144 if proto.runstate == Runstate.CONNECTING: 145 emsg = f"{name} is currently connecting." 146 elif proto.runstate == Runstate.DISCONNECTING: 147 emsg = (f"{name} is disconnecting." 148 " Call disconnect() to return to IDLE state.") 149 elif proto.runstate == Runstate.RUNNING: 150 emsg = f"{name} is already connected and running." 151 elif proto.runstate == Runstate.IDLE: 152 emsg = f"{name} is disconnected and idle." 153 else: 154 assert False 155 raise StateError(emsg, proto.runstate, required_state) 156 # No StateError, so call the wrapped method. 157 return func(proto, *args, **kwargs) 158 159 # Return the decorated method; 160 # Transforming Func to Decorated[Func]. 161 return cast(F, _wrapper) 162 163 # Return the decorator instance from the decorator factory. Phew! 164 return _decorator 165 166 167class AsyncProtocol(Generic[T]): 168 """ 169 AsyncProtocol implements a generic async message-based protocol. 170 171 This protocol assumes the basic unit of information transfer between 172 client and server is a "message", the details of which are left up 173 to the implementation. It assumes the sending and receiving of these 174 messages is full-duplex and not necessarily correlated; i.e. it 175 supports asynchronous inbound messages. 176 177 It is designed to be extended by a specific protocol which provides 178 the implementations for how to read and send messages. These must be 179 defined in `_do_recv()` and `_do_send()`, respectively. 180 181 Other callbacks have a default implementation, but are intended to be 182 either extended or overridden: 183 184 - `_establish_session`: 185 The base implementation starts the reader/writer tasks. 186 A protocol implementation can override this call, inserting 187 actions to be taken prior to starting the reader/writer tasks 188 before the super() call; actions needing to occur afterwards 189 can be written after the super() call. 190 - `_on_message`: 191 Actions to be performed when a message is received. 192 - `_cb_outbound`: 193 Logging/Filtering hook for all outbound messages. 194 - `_cb_inbound`: 195 Logging/Filtering hook for all inbound messages. 196 This hook runs *before* `_on_message()`. 197 198 :param name: 199 Name used for logging messages, if any. By default, messages 200 will log to 'qemu.qmp.protocol', but each individual connection 201 can be given its own logger by giving it a name; messages will 202 then log to 'qemu.qmp.protocol.${name}'. 203 """ 204 # pylint: disable=too-many-instance-attributes 205 206 #: Logger object for debugging messages from this connection. 207 logger = logging.getLogger(__name__) 208 209 # Maximum allowable size of read buffer 210 _limit = 64 * 1024 211 212 # ------------------------- 213 # Section: Public interface 214 # ------------------------- 215 216 def __init__(self, name: Optional[str] = None) -> None: 217 #: The nickname for this connection, if any. 218 self.name: Optional[str] = name 219 if self.name is not None: 220 self.logger = self.logger.getChild(self.name) 221 222 # stream I/O 223 self._reader: Optional[StreamReader] = None 224 self._writer: Optional[StreamWriter] = None 225 226 # Outbound Message queue 227 self._outgoing: asyncio.Queue[T] 228 229 # Special, long-running tasks: 230 self._reader_task: Optional[asyncio.Future[None]] = None 231 self._writer_task: Optional[asyncio.Future[None]] = None 232 233 # Aggregate of the above two tasks, used for Exception management. 234 self._bh_tasks: Optional[asyncio.Future[Tuple[None, None]]] = None 235 236 #: Disconnect task. The disconnect implementation runs in a task 237 #: so that asynchronous disconnects (initiated by the 238 #: reader/writer) are allowed to wait for the reader/writers to 239 #: exit. 240 self._dc_task: Optional[asyncio.Future[None]] = None 241 242 self._runstate = Runstate.IDLE 243 self._runstate_changed: Optional[asyncio.Event] = None 244 245 # Server state for start_server() and _incoming() 246 self._server: Optional[asyncio.AbstractServer] = None 247 self._accepted: Optional[asyncio.Event] = None 248 249 def __repr__(self) -> str: 250 cls_name = type(self).__name__ 251 tokens = [] 252 if self.name is not None: 253 tokens.append(f"name={self.name!r}") 254 tokens.append(f"runstate={self.runstate.name}") 255 return f"<{cls_name} {' '.join(tokens)}>" 256 257 @property # @upper_half 258 def runstate(self) -> Runstate: 259 """The current `Runstate` of the connection.""" 260 return self._runstate 261 262 @upper_half 263 async def runstate_changed(self) -> Runstate: 264 """ 265 Wait for the `runstate` to change, then return that runstate. 266 """ 267 await self._runstate_event.wait() 268 return self.runstate 269 270 @upper_half 271 @require(Runstate.IDLE) 272 async def start_server_and_accept( 273 self, address: SocketAddrT, 274 ssl: Optional[SSLContext] = None 275 ) -> None: 276 """ 277 Accept a connection and begin processing message queues. 278 279 If this call fails, `runstate` is guaranteed to be set back to `IDLE`. 280 This method is precisely equivalent to calling `start_server()` 281 followed by `accept()`. 282 283 :param address: 284 Address to listen on; UNIX socket path or TCP address/port. 285 :param ssl: SSL context to use, if any. 286 287 :raise StateError: When the `Runstate` is not `IDLE`. 288 :raise ConnectError: 289 When a connection or session cannot be established. 290 291 This exception will wrap a more concrete one. In most cases, 292 the wrapped exception will be `OSError` or `EOFError`. If a 293 protocol-level failure occurs while establishing a new 294 session, the wrapped error may also be an `QMPError`. 295 """ 296 await self.start_server(address, ssl) 297 await self.accept() 298 assert self.runstate == Runstate.RUNNING 299 300 @upper_half 301 @require(Runstate.IDLE) 302 async def start_server(self, address: SocketAddrT, 303 ssl: Optional[SSLContext] = None) -> None: 304 """ 305 Start listening for an incoming connection, but do not wait for a peer. 306 307 This method starts listening for an incoming connection, but 308 does not block waiting for a peer. This call will return 309 immediately after binding and listening on a socket. A later 310 call to `accept()` must be made in order to finalize the 311 incoming connection. 312 313 :param address: 314 Address to listen on; UNIX socket path or TCP address/port. 315 :param ssl: SSL context to use, if any. 316 317 :raise StateError: When the `Runstate` is not `IDLE`. 318 :raise ConnectError: 319 When the server could not start listening on this address. 320 321 This exception will wrap a more concrete one. In most cases, 322 the wrapped exception will be `OSError`. 323 """ 324 await self._session_guard( 325 self._do_start_server(address, ssl), 326 'Failed to establish connection') 327 assert self.runstate == Runstate.CONNECTING 328 329 @upper_half 330 @require(Runstate.CONNECTING) 331 async def accept(self) -> None: 332 """ 333 Accept an incoming connection and begin processing message queues. 334 335 If this call fails, `runstate` is guaranteed to be set back to `IDLE`. 336 337 :raise StateError: When the `Runstate` is not `CONNECTING`. 338 :raise QMPError: When `start_server()` was not called yet. 339 :raise ConnectError: 340 When a connection or session cannot be established. 341 342 This exception will wrap a more concrete one. In most cases, 343 the wrapped exception will be `OSError` or `EOFError`. If a 344 protocol-level failure occurs while establishing a new 345 session, the wrapped error may also be an `QMPError`. 346 """ 347 if self._accepted is None: 348 raise QMPError("Cannot call accept() before start_server().") 349 await self._session_guard( 350 self._do_accept(), 351 'Failed to establish connection') 352 await self._session_guard( 353 self._establish_session(), 354 'Failed to establish session') 355 assert self.runstate == Runstate.RUNNING 356 357 @upper_half 358 @require(Runstate.IDLE) 359 async def connect(self, address: Union[SocketAddrT, socket.socket], 360 ssl: Optional[SSLContext] = None) -> None: 361 """ 362 Connect to the server and begin processing message queues. 363 364 If this call fails, `runstate` is guaranteed to be set back to `IDLE`. 365 366 :param address: 367 Address to connect to; UNIX socket path or TCP address/port. 368 :param ssl: SSL context to use, if any. 369 370 :raise StateError: When the `Runstate` is not `IDLE`. 371 :raise ConnectError: 372 When a connection or session cannot be established. 373 374 This exception will wrap a more concrete one. In most cases, 375 the wrapped exception will be `OSError` or `EOFError`. If a 376 protocol-level failure occurs while establishing a new 377 session, the wrapped error may also be an `QMPError`. 378 """ 379 await self._session_guard( 380 self._do_connect(address, ssl), 381 'Failed to establish connection') 382 await self._session_guard( 383 self._establish_session(), 384 'Failed to establish session') 385 assert self.runstate == Runstate.RUNNING 386 387 @upper_half 388 async def disconnect(self) -> None: 389 """ 390 Disconnect and wait for all tasks to fully stop. 391 392 If there was an exception that caused the reader/writers to 393 terminate prematurely, it will be raised here. 394 395 :raise Exception: When the reader or writer terminate unexpectedly. 396 """ 397 self.logger.debug("disconnect() called.") 398 self._schedule_disconnect() 399 await self._wait_disconnect() 400 401 # -------------------------- 402 # Section: Session machinery 403 # -------------------------- 404 405 async def _session_guard(self, coro: Awaitable[None], emsg: str) -> None: 406 """ 407 Async guard function used to roll back to `IDLE` on any error. 408 409 On any Exception, the state machine will be reset back to 410 `IDLE`. Most Exceptions will be wrapped with `ConnectError`, but 411 `BaseException` events will be left alone (This includes 412 asyncio.CancelledError, even prior to Python 3.8). 413 414 :param error_message: 415 Human-readable string describing what connection phase failed. 416 417 :raise BaseException: 418 When `BaseException` occurs in the guarded block. 419 :raise ConnectError: 420 When any other error is encountered in the guarded block. 421 """ 422 # Note: After Python 3.6 support is removed, this should be an 423 # @asynccontextmanager instead of accepting a callback. 424 try: 425 await coro 426 except BaseException as err: 427 self.logger.error("%s: %s", emsg, exception_summary(err)) 428 self.logger.debug("%s:\n%s\n", emsg, pretty_traceback()) 429 try: 430 # Reset the runstate back to IDLE. 431 await self.disconnect() 432 except: 433 # We don't expect any Exceptions from the disconnect function 434 # here, because we failed to connect in the first place. 435 # The disconnect() function is intended to perform 436 # only cannot-fail cleanup here, but you never know. 437 emsg = ( 438 "Unexpected bottom half exception. " 439 "This is a bug in the QMP library. " 440 "Please report it to <qemu-devel@nongnu.org> and " 441 "CC: John Snow <jsnow@redhat.com>." 442 ) 443 self.logger.critical("%s:\n%s\n", emsg, pretty_traceback()) 444 raise 445 446 # CancelledError is an Exception with special semantic meaning; 447 # We do NOT want to wrap it up under ConnectError. 448 # NB: CancelledError is not a BaseException before Python 3.8 449 if isinstance(err, asyncio.CancelledError): 450 raise 451 452 # Any other kind of error can be treated as some kind of connection 453 # failure broadly. Inspect the 'exc' field to explore the root 454 # cause in greater detail. 455 if isinstance(err, Exception): 456 raise ConnectError(emsg, err) from err 457 458 # Raise BaseExceptions un-wrapped, they're more important. 459 raise 460 461 @property 462 def _runstate_event(self) -> asyncio.Event: 463 # asyncio.Event() objects should not be created prior to entrance into 464 # an event loop, so we can ensure we create it in the correct context. 465 # Create it on-demand *only* at the behest of an 'async def' method. 466 if not self._runstate_changed: 467 self._runstate_changed = asyncio.Event() 468 return self._runstate_changed 469 470 @upper_half 471 @bottom_half 472 def _set_state(self, state: Runstate) -> None: 473 """ 474 Change the `Runstate` of the protocol connection. 475 476 Signals the `runstate_changed` event. 477 """ 478 if state == self._runstate: 479 return 480 481 self.logger.debug("Transitioning from '%s' to '%s'.", 482 str(self._runstate), str(state)) 483 self._runstate = state 484 self._runstate_event.set() 485 self._runstate_event.clear() 486 487 @bottom_half 488 async def _stop_server(self) -> None: 489 """ 490 Stop listening for / accepting new incoming connections. 491 """ 492 if self._server is None: 493 return 494 495 try: 496 self.logger.debug("Stopping server.") 497 self._server.close() 498 await self._server.wait_closed() 499 self.logger.debug("Server stopped.") 500 finally: 501 self._server = None 502 503 @bottom_half # However, it does not run from the R/W tasks. 504 async def _incoming(self, 505 reader: asyncio.StreamReader, 506 writer: asyncio.StreamWriter) -> None: 507 """ 508 Accept an incoming connection and signal the upper_half. 509 510 This method does the minimum necessary to accept a single 511 incoming connection. It signals back to the upper_half ASAP so 512 that any errors during session initialization can occur 513 naturally in the caller's stack. 514 515 :param reader: Incoming `asyncio.StreamReader` 516 :param writer: Incoming `asyncio.StreamWriter` 517 """ 518 peer = writer.get_extra_info('peername', 'Unknown peer') 519 self.logger.debug("Incoming connection from %s", peer) 520 521 if self._reader or self._writer: 522 # Sadly, we can have more than one pending connection 523 # because of https://bugs.python.org/issue46715 524 # Close any extra connections we don't actually want. 525 self.logger.warning("Extraneous connection inadvertently accepted") 526 writer.close() 527 return 528 529 # A connection has been accepted; stop listening for new ones. 530 assert self._accepted is not None 531 await self._stop_server() 532 self._reader, self._writer = (reader, writer) 533 self._accepted.set() 534 535 @upper_half 536 async def _do_start_server(self, address: SocketAddrT, 537 ssl: Optional[SSLContext] = None) -> None: 538 """ 539 Start listening for an incoming connection, but do not wait for a peer. 540 541 This method starts listening for an incoming connection, but does not 542 block waiting for a peer. This call will return immediately after 543 binding and listening to a socket. A later call to accept() must be 544 made in order to finalize the incoming connection. 545 546 :param address: 547 Address to listen on; UNIX socket path or TCP address/port. 548 :param ssl: SSL context to use, if any. 549 550 :raise OSError: For stream-related errors. 551 """ 552 assert self.runstate == Runstate.IDLE 553 self._set_state(Runstate.CONNECTING) 554 555 self.logger.debug("Awaiting connection on %s ...", address) 556 self._accepted = asyncio.Event() 557 558 if isinstance(address, tuple): 559 coro = asyncio.start_server( 560 self._incoming, 561 host=address[0], 562 port=address[1], 563 ssl=ssl, 564 backlog=1, 565 limit=self._limit, 566 ) 567 else: 568 coro = asyncio.start_unix_server( 569 self._incoming, 570 path=address, 571 ssl=ssl, 572 backlog=1, 573 limit=self._limit, 574 ) 575 576 # Allow runstate watchers to witness 'CONNECTING' state; some 577 # failures in the streaming layer are synchronous and will not 578 # otherwise yield. 579 await asyncio.sleep(0) 580 581 # This will start the server (bind(2), listen(2)). It will also 582 # call accept(2) if we yield, but we don't block on that here. 583 self._server = await coro 584 self.logger.debug("Server listening on %s", address) 585 586 @upper_half 587 async def _do_accept(self) -> None: 588 """ 589 Wait for and accept an incoming connection. 590 591 Requires that we have not yet accepted an incoming connection 592 from the upper_half, but it's OK if the server is no longer 593 running because the bottom_half has already accepted the 594 connection. 595 """ 596 assert self._accepted is not None 597 await self._accepted.wait() 598 assert self._server is None 599 self._accepted = None 600 601 self.logger.debug("Connection accepted.") 602 603 @upper_half 604 async def _do_connect(self, address: Union[SocketAddrT, socket.socket], 605 ssl: Optional[SSLContext] = None) -> None: 606 """ 607 Acting as the transport client, initiate a connection to a server. 608 609 :param address: 610 Address to connect to; UNIX socket path or TCP address/port. 611 :param ssl: SSL context to use, if any. 612 613 :raise OSError: For stream-related errors. 614 """ 615 assert self.runstate == Runstate.IDLE 616 self._set_state(Runstate.CONNECTING) 617 618 # Allow runstate watchers to witness 'CONNECTING' state; some 619 # failures in the streaming layer are synchronous and will not 620 # otherwise yield. 621 await asyncio.sleep(0) 622 623 if isinstance(address, socket.socket): 624 self.logger.debug("Connecting with existing socket: " 625 "fd=%d, family=%r, type=%r", 626 address.fileno(), address.family, address.type) 627 connect = asyncio.open_connection( 628 limit=self._limit, 629 ssl=ssl, 630 sock=address, 631 ) 632 elif isinstance(address, tuple): 633 self.logger.debug("Connecting to %s ...", address) 634 connect = asyncio.open_connection( 635 address[0], 636 address[1], 637 ssl=ssl, 638 limit=self._limit, 639 ) 640 else: 641 self.logger.debug("Connecting to file://%s ...", address) 642 connect = asyncio.open_unix_connection( 643 path=address, 644 ssl=ssl, 645 limit=self._limit, 646 ) 647 648 self._reader, self._writer = await connect 649 self.logger.debug("Connected.") 650 651 @upper_half 652 async def _establish_session(self) -> None: 653 """ 654 Establish a new session. 655 656 Starts the readers/writer tasks; subclasses may perform their 657 own negotiations here. The Runstate will be RUNNING upon 658 successful conclusion. 659 """ 660 assert self.runstate == Runstate.CONNECTING 661 662 self._outgoing = asyncio.Queue() 663 664 reader_coro = self._bh_loop_forever(self._bh_recv_message, 'Reader') 665 writer_coro = self._bh_loop_forever(self._bh_send_message, 'Writer') 666 667 self._reader_task = create_task(reader_coro) 668 self._writer_task = create_task(writer_coro) 669 670 self._bh_tasks = asyncio.gather( 671 self._reader_task, 672 self._writer_task, 673 ) 674 675 self._set_state(Runstate.RUNNING) 676 await asyncio.sleep(0) # Allow runstate_event to process 677 678 @upper_half 679 @bottom_half 680 def _schedule_disconnect(self) -> None: 681 """ 682 Initiate a disconnect; idempotent. 683 684 This method is used both in the upper-half as a direct 685 consequence of `disconnect()`, and in the bottom-half in the 686 case of unhandled exceptions in the reader/writer tasks. 687 688 It can be invoked no matter what the `runstate` is. 689 """ 690 if not self._dc_task: 691 self._set_state(Runstate.DISCONNECTING) 692 self.logger.debug("Scheduling disconnect.") 693 self._dc_task = create_task(self._bh_disconnect()) 694 695 @upper_half 696 async def _wait_disconnect(self) -> None: 697 """ 698 Waits for a previously scheduled disconnect to finish. 699 700 This method will gather any bottom half exceptions and re-raise 701 the one that occurred first; presuming it to be the root cause 702 of any subsequent Exceptions. It is intended to be used in the 703 upper half of the call chain. 704 705 :raise Exception: 706 Arbitrary exception re-raised on behalf of the reader/writer. 707 """ 708 assert self.runstate == Runstate.DISCONNECTING 709 assert self._dc_task 710 711 aws: List[Awaitable[object]] = [self._dc_task] 712 if self._bh_tasks: 713 aws.insert(0, self._bh_tasks) 714 all_defined_tasks = asyncio.gather(*aws) 715 716 # Ensure disconnect is done; Exception (if any) is not raised here: 717 await asyncio.wait((self._dc_task,)) 718 719 try: 720 await all_defined_tasks # Raise Exceptions from the bottom half. 721 finally: 722 self._cleanup() 723 self._set_state(Runstate.IDLE) 724 725 @upper_half 726 def _cleanup(self) -> None: 727 """ 728 Fully reset this object to a clean state and return to `IDLE`. 729 """ 730 def _paranoid_task_erase(task: Optional['asyncio.Future[_U]'] 731 ) -> Optional['asyncio.Future[_U]']: 732 # Help to erase a task, ENSURING it is fully quiesced first. 733 assert (task is None) or task.done() 734 return None if (task and task.done()) else task 735 736 assert self.runstate == Runstate.DISCONNECTING 737 self._dc_task = _paranoid_task_erase(self._dc_task) 738 self._reader_task = _paranoid_task_erase(self._reader_task) 739 self._writer_task = _paranoid_task_erase(self._writer_task) 740 self._bh_tasks = _paranoid_task_erase(self._bh_tasks) 741 742 self._reader = None 743 self._writer = None 744 self._accepted = None 745 746 # NB: _runstate_changed cannot be cleared because we still need it to 747 # send the final runstate changed event ...! 748 749 # ---------------------------- 750 # Section: Bottom Half methods 751 # ---------------------------- 752 753 @bottom_half 754 async def _bh_disconnect(self) -> None: 755 """ 756 Disconnect and cancel all outstanding tasks. 757 758 It is designed to be called from its task context, 759 :py:obj:`~AsyncProtocol._dc_task`. By running in its own task, 760 it is free to wait on any pending actions that may still need to 761 occur in either the reader or writer tasks. 762 """ 763 assert self.runstate == Runstate.DISCONNECTING 764 765 def _done(task: Optional['asyncio.Future[Any]']) -> bool: 766 return task is not None and task.done() 767 768 # If the server is running, stop it. 769 await self._stop_server() 770 771 # Are we already in an error pathway? If either of the tasks are 772 # already done, or if we have no tasks but a reader/writer; we 773 # must be. 774 # 775 # NB: We can't use _bh_tasks to check for premature task 776 # completion, because it may not yet have had a chance to run 777 # and gather itself. 778 tasks = tuple(filter(None, (self._writer_task, self._reader_task))) 779 error_pathway = _done(self._reader_task) or _done(self._writer_task) 780 if not tasks: 781 error_pathway |= bool(self._reader) or bool(self._writer) 782 783 try: 784 # Try to flush the writer, if possible. 785 # This *may* cause an error and force us over into the error path. 786 if not error_pathway: 787 await self._bh_flush_writer() 788 except BaseException as err: 789 error_pathway = True 790 emsg = "Failed to flush the writer" 791 self.logger.error("%s: %s", emsg, exception_summary(err)) 792 self.logger.debug("%s:\n%s\n", emsg, pretty_traceback()) 793 raise 794 finally: 795 # Cancel any still-running tasks (Won't raise): 796 if self._writer_task is not None and not self._writer_task.done(): 797 self.logger.debug("Cancelling writer task.") 798 self._writer_task.cancel() 799 if self._reader_task is not None and not self._reader_task.done(): 800 self.logger.debug("Cancelling reader task.") 801 self._reader_task.cancel() 802 803 # Close out the tasks entirely (Won't raise): 804 if tasks: 805 self.logger.debug("Waiting for tasks to complete ...") 806 await asyncio.wait(tasks) 807 808 # Lastly, close the stream itself. (*May raise*!): 809 await self._bh_close_stream(error_pathway) 810 self.logger.debug("Disconnected.") 811 812 @bottom_half 813 async def _bh_flush_writer(self) -> None: 814 if not self._writer_task: 815 return 816 817 self.logger.debug("Draining the outbound queue ...") 818 await self._outgoing.join() 819 if self._writer is not None: 820 self.logger.debug("Flushing the StreamWriter ...") 821 await flush(self._writer) 822 823 @bottom_half 824 async def _bh_close_stream(self, error_pathway: bool = False) -> None: 825 # NB: Closing the writer also implicitly closes the reader. 826 if not self._writer: 827 return 828 829 if not is_closing(self._writer): 830 self.logger.debug("Closing StreamWriter.") 831 self._writer.close() 832 833 self.logger.debug("Waiting for StreamWriter to close ...") 834 try: 835 await wait_closed(self._writer) 836 except Exception: # pylint: disable=broad-except 837 # It's hard to tell if the Stream is already closed or 838 # not. Even if one of the tasks has failed, it may have 839 # failed for a higher-layered protocol reason. The 840 # stream could still be open and perfectly fine. 841 # I don't know how to discern its health here. 842 843 if error_pathway: 844 # We already know that *something* went wrong. Let's 845 # just trust that the Exception we already have is the 846 # better one to present to the user, even if we don't 847 # genuinely *know* the relationship between the two. 848 self.logger.debug( 849 "Discarding Exception from wait_closed:\n%s\n", 850 pretty_traceback(), 851 ) 852 else: 853 # Oops, this is a brand-new error! 854 raise 855 finally: 856 self.logger.debug("StreamWriter closed.") 857 858 @bottom_half 859 async def _bh_loop_forever(self, async_fn: _TaskFN, name: str) -> None: 860 """ 861 Run one of the bottom-half methods in a loop forever. 862 863 If the bottom half ever raises any exception, schedule a 864 disconnect that will terminate the entire loop. 865 866 :param async_fn: The bottom-half method to run in a loop. 867 :param name: The name of this task, used for logging. 868 """ 869 try: 870 while True: 871 await async_fn() 872 except asyncio.CancelledError: 873 # We have been cancelled by _bh_disconnect, exit gracefully. 874 self.logger.debug("Task.%s: cancelled.", name) 875 return 876 except BaseException as err: 877 self.logger.log( 878 logging.INFO if isinstance(err, EOFError) else logging.ERROR, 879 "Task.%s: %s", 880 name, exception_summary(err) 881 ) 882 self.logger.debug("Task.%s: failure:\n%s\n", 883 name, pretty_traceback()) 884 self._schedule_disconnect() 885 raise 886 finally: 887 self.logger.debug("Task.%s: exiting.", name) 888 889 @bottom_half 890 async def _bh_send_message(self) -> None: 891 """ 892 Wait for an outgoing message, then send it. 893 894 Designed to be run in `_bh_loop_forever()`. 895 """ 896 msg = await self._outgoing.get() 897 try: 898 await self._send(msg) 899 finally: 900 self._outgoing.task_done() 901 902 @bottom_half 903 async def _bh_recv_message(self) -> None: 904 """ 905 Wait for an incoming message and call `_on_message` to route it. 906 907 Designed to be run in `_bh_loop_forever()`. 908 """ 909 msg = await self._recv() 910 await self._on_message(msg) 911 912 # -------------------- 913 # Section: Message I/O 914 # -------------------- 915 916 @upper_half 917 @bottom_half 918 def _cb_outbound(self, msg: T) -> T: 919 """ 920 Callback: outbound message hook. 921 922 This is intended for subclasses to be able to add arbitrary 923 hooks to filter or manipulate outgoing messages. The base 924 implementation does nothing but log the message without any 925 manipulation of the message. 926 927 :param msg: raw outbound message 928 :return: final outbound message 929 """ 930 self.logger.debug("--> %s", str(msg)) 931 return msg 932 933 @upper_half 934 @bottom_half 935 def _cb_inbound(self, msg: T) -> T: 936 """ 937 Callback: inbound message hook. 938 939 This is intended for subclasses to be able to add arbitrary 940 hooks to filter or manipulate incoming messages. The base 941 implementation does nothing but log the message without any 942 manipulation of the message. 943 944 This method does not "handle" incoming messages; it is a filter. 945 The actual "endpoint" for incoming messages is `_on_message()`. 946 947 :param msg: raw inbound message 948 :return: processed inbound message 949 """ 950 self.logger.debug("<-- %s", str(msg)) 951 return msg 952 953 @upper_half 954 @bottom_half 955 async def _readline(self) -> bytes: 956 """ 957 Wait for a newline from the incoming reader. 958 959 This method is provided as a convenience for upper-layer 960 protocols, as many are line-based. 961 962 This method *may* return a sequence of bytes without a trailing 963 newline if EOF occurs, but *some* bytes were received. In this 964 case, the next call will raise `EOFError`. It is assumed that 965 the layer 5 protocol will decide if there is anything meaningful 966 to be done with a partial message. 967 968 :raise OSError: For stream-related errors. 969 :raise EOFError: 970 If the reader stream is at EOF and there are no bytes to return. 971 :return: bytes, including the newline. 972 """ 973 assert self._reader is not None 974 msg_bytes = await self._reader.readline() 975 976 if not msg_bytes: 977 if self._reader.at_eof(): 978 raise EOFError 979 980 return msg_bytes 981 982 @upper_half 983 @bottom_half 984 async def _do_recv(self) -> T: 985 """ 986 Abstract: Read from the stream and return a message. 987 988 Very low-level; intended to only be called by `_recv()`. 989 """ 990 raise NotImplementedError 991 992 @upper_half 993 @bottom_half 994 async def _recv(self) -> T: 995 """ 996 Read an arbitrary protocol message. 997 998 .. warning:: 999 This method is intended primarily for `_bh_recv_message()` 1000 to use in an asynchronous task loop. Using it outside of 1001 this loop will "steal" messages from the normal routing 1002 mechanism. It is safe to use prior to `_establish_session()`, 1003 but should not be used otherwise. 1004 1005 This method uses `_do_recv()` to retrieve the raw message, and 1006 then transforms it using `_cb_inbound()`. 1007 1008 :return: A single (filtered, processed) protocol message. 1009 """ 1010 message = await self._do_recv() 1011 return self._cb_inbound(message) 1012 1013 @upper_half 1014 @bottom_half 1015 def _do_send(self, msg: T) -> None: 1016 """ 1017 Abstract: Write a message to the stream. 1018 1019 Very low-level; intended to only be called by `_send()`. 1020 """ 1021 raise NotImplementedError 1022 1023 @upper_half 1024 @bottom_half 1025 async def _send(self, msg: T) -> None: 1026 """ 1027 Send an arbitrary protocol message. 1028 1029 This method will transform any outgoing messages according to 1030 `_cb_outbound()`. 1031 1032 .. warning:: 1033 Like `_recv()`, this method is intended to be called by 1034 the writer task loop that processes outgoing 1035 messages. Calling it directly may circumvent logic 1036 implemented by the caller meant to correlate outgoing and 1037 incoming messages. 1038 1039 :raise OSError: For problems with the underlying stream. 1040 """ 1041 msg = self._cb_outbound(msg) 1042 self._do_send(msg) 1043 1044 @bottom_half 1045 async def _on_message(self, msg: T) -> None: 1046 """ 1047 Called to handle the receipt of a new message. 1048 1049 .. caution:: 1050 This is executed from within the reader loop, so be advised 1051 that waiting on either the reader or writer task will lead 1052 to deadlock. Additionally, any unhandled exceptions will 1053 directly cause the loop to halt, so logic may be best-kept 1054 to a minimum if at all possible. 1055 1056 :param msg: The incoming message, already logged/filtered. 1057 """ 1058 # Nothing to do in the abstract case. 1059