xref: /openbmc/qemu/python/qemu/qmp/protocol.py (revision acf87387)
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            self.logger.debug("Server stopped.")
499        finally:
500            self._server = None
501
502    @bottom_half  # However, it does not run from the R/W tasks.
503    async def _incoming(self,
504                        reader: asyncio.StreamReader,
505                        writer: asyncio.StreamWriter) -> None:
506        """
507        Accept an incoming connection and signal the upper_half.
508
509        This method does the minimum necessary to accept a single
510        incoming connection. It signals back to the upper_half ASAP so
511        that any errors during session initialization can occur
512        naturally in the caller's stack.
513
514        :param reader: Incoming `asyncio.StreamReader`
515        :param writer: Incoming `asyncio.StreamWriter`
516        """
517        peer = writer.get_extra_info('peername', 'Unknown peer')
518        self.logger.debug("Incoming connection from %s", peer)
519
520        if self._reader or self._writer:
521            # Sadly, we can have more than one pending connection
522            # because of https://bugs.python.org/issue46715
523            # Close any extra connections we don't actually want.
524            self.logger.warning("Extraneous connection inadvertently accepted")
525            writer.close()
526            return
527
528        # A connection has been accepted; stop listening for new ones.
529        assert self._accepted is not None
530        await self._stop_server()
531        self._reader, self._writer = (reader, writer)
532        self._accepted.set()
533
534    @upper_half
535    async def _do_start_server(self, address: SocketAddrT,
536                               ssl: Optional[SSLContext] = None) -> None:
537        """
538        Start listening for an incoming connection, but do not wait for a peer.
539
540        This method starts listening for an incoming connection, but does not
541        block waiting for a peer. This call will return immediately after
542        binding and listening to a socket. A later call to accept() must be
543        made in order to finalize the incoming connection.
544
545        :param address:
546            Address to listen on; UNIX socket path or TCP address/port.
547        :param ssl: SSL context to use, if any.
548
549        :raise OSError: For stream-related errors.
550        """
551        assert self.runstate == Runstate.IDLE
552        self._set_state(Runstate.CONNECTING)
553
554        self.logger.debug("Awaiting connection on %s ...", address)
555        self._accepted = asyncio.Event()
556
557        if isinstance(address, tuple):
558            coro = asyncio.start_server(
559                self._incoming,
560                host=address[0],
561                port=address[1],
562                ssl=ssl,
563                backlog=1,
564                limit=self._limit,
565            )
566        else:
567            coro = asyncio.start_unix_server(
568                self._incoming,
569                path=address,
570                ssl=ssl,
571                backlog=1,
572                limit=self._limit,
573            )
574
575        # Allow runstate watchers to witness 'CONNECTING' state; some
576        # failures in the streaming layer are synchronous and will not
577        # otherwise yield.
578        await asyncio.sleep(0)
579
580        # This will start the server (bind(2), listen(2)). It will also
581        # call accept(2) if we yield, but we don't block on that here.
582        self._server = await coro
583        self.logger.debug("Server listening on %s", address)
584
585    @upper_half
586    async def _do_accept(self) -> None:
587        """
588        Wait for and accept an incoming connection.
589
590        Requires that we have not yet accepted an incoming connection
591        from the upper_half, but it's OK if the server is no longer
592        running because the bottom_half has already accepted the
593        connection.
594        """
595        assert self._accepted is not None
596        await self._accepted.wait()
597        assert self._server is None
598        self._accepted = None
599
600        self.logger.debug("Connection accepted.")
601
602    @upper_half
603    async def _do_connect(self, address: Union[SocketAddrT, socket.socket],
604                          ssl: Optional[SSLContext] = None) -> None:
605        """
606        Acting as the transport client, initiate a connection to a server.
607
608        :param address:
609            Address to connect to; UNIX socket path or TCP address/port.
610        :param ssl: SSL context to use, if any.
611
612        :raise OSError: For stream-related errors.
613        """
614        assert self.runstate == Runstate.IDLE
615        self._set_state(Runstate.CONNECTING)
616
617        # Allow runstate watchers to witness 'CONNECTING' state; some
618        # failures in the streaming layer are synchronous and will not
619        # otherwise yield.
620        await asyncio.sleep(0)
621
622        if isinstance(address, socket.socket):
623            self.logger.debug("Connecting with existing socket: "
624                              "fd=%d, family=%r, type=%r",
625                              address.fileno(), address.family, address.type)
626            connect = asyncio.open_connection(
627                limit=self._limit,
628                ssl=ssl,
629                sock=address,
630            )
631        elif isinstance(address, tuple):
632            self.logger.debug("Connecting to %s ...", address)
633            connect = asyncio.open_connection(
634                address[0],
635                address[1],
636                ssl=ssl,
637                limit=self._limit,
638            )
639        else:
640            self.logger.debug("Connecting to file://%s ...", address)
641            connect = asyncio.open_unix_connection(
642                path=address,
643                ssl=ssl,
644                limit=self._limit,
645            )
646
647        self._reader, self._writer = await connect
648        self.logger.debug("Connected.")
649
650    @upper_half
651    async def _establish_session(self) -> None:
652        """
653        Establish a new session.
654
655        Starts the readers/writer tasks; subclasses may perform their
656        own negotiations here. The Runstate will be RUNNING upon
657        successful conclusion.
658        """
659        assert self.runstate == Runstate.CONNECTING
660
661        self._outgoing = asyncio.Queue()
662
663        reader_coro = self._bh_loop_forever(self._bh_recv_message, 'Reader')
664        writer_coro = self._bh_loop_forever(self._bh_send_message, 'Writer')
665
666        self._reader_task = create_task(reader_coro)
667        self._writer_task = create_task(writer_coro)
668
669        self._bh_tasks = asyncio.gather(
670            self._reader_task,
671            self._writer_task,
672        )
673
674        self._set_state(Runstate.RUNNING)
675        await asyncio.sleep(0)  # Allow runstate_event to process
676
677    @upper_half
678    @bottom_half
679    def _schedule_disconnect(self) -> None:
680        """
681        Initiate a disconnect; idempotent.
682
683        This method is used both in the upper-half as a direct
684        consequence of `disconnect()`, and in the bottom-half in the
685        case of unhandled exceptions in the reader/writer tasks.
686
687        It can be invoked no matter what the `runstate` is.
688        """
689        if not self._dc_task:
690            self._set_state(Runstate.DISCONNECTING)
691            self.logger.debug("Scheduling disconnect.")
692            self._dc_task = create_task(self._bh_disconnect())
693
694    @upper_half
695    async def _wait_disconnect(self) -> None:
696        """
697        Waits for a previously scheduled disconnect to finish.
698
699        This method will gather any bottom half exceptions and re-raise
700        the one that occurred first; presuming it to be the root cause
701        of any subsequent Exceptions. It is intended to be used in the
702        upper half of the call chain.
703
704        :raise Exception:
705            Arbitrary exception re-raised on behalf of the reader/writer.
706        """
707        assert self.runstate == Runstate.DISCONNECTING
708        assert self._dc_task
709
710        aws: List[Awaitable[object]] = [self._dc_task]
711        if self._bh_tasks:
712            aws.insert(0, self._bh_tasks)
713        all_defined_tasks = asyncio.gather(*aws)
714
715        # Ensure disconnect is done; Exception (if any) is not raised here:
716        await asyncio.wait((self._dc_task,))
717
718        try:
719            await all_defined_tasks  # Raise Exceptions from the bottom half.
720        finally:
721            self._cleanup()
722            self._set_state(Runstate.IDLE)
723
724    @upper_half
725    def _cleanup(self) -> None:
726        """
727        Fully reset this object to a clean state and return to `IDLE`.
728        """
729        def _paranoid_task_erase(task: Optional['asyncio.Future[_U]']
730                                 ) -> Optional['asyncio.Future[_U]']:
731            # Help to erase a task, ENSURING it is fully quiesced first.
732            assert (task is None) or task.done()
733            return None if (task and task.done()) else task
734
735        assert self.runstate == Runstate.DISCONNECTING
736        self._dc_task = _paranoid_task_erase(self._dc_task)
737        self._reader_task = _paranoid_task_erase(self._reader_task)
738        self._writer_task = _paranoid_task_erase(self._writer_task)
739        self._bh_tasks = _paranoid_task_erase(self._bh_tasks)
740
741        self._reader = None
742        self._writer = None
743        self._accepted = None
744
745        # NB: _runstate_changed cannot be cleared because we still need it to
746        # send the final runstate changed event ...!
747
748    # ----------------------------
749    # Section: Bottom Half methods
750    # ----------------------------
751
752    @bottom_half
753    async def _bh_disconnect(self) -> None:
754        """
755        Disconnect and cancel all outstanding tasks.
756
757        It is designed to be called from its task context,
758        :py:obj:`~AsyncProtocol._dc_task`. By running in its own task,
759        it is free to wait on any pending actions that may still need to
760        occur in either the reader or writer tasks.
761        """
762        assert self.runstate == Runstate.DISCONNECTING
763
764        def _done(task: Optional['asyncio.Future[Any]']) -> bool:
765            return task is not None and task.done()
766
767        # If the server is running, stop it.
768        await self._stop_server()
769
770        # Are we already in an error pathway? If either of the tasks are
771        # already done, or if we have no tasks but a reader/writer; we
772        # must be.
773        #
774        # NB: We can't use _bh_tasks to check for premature task
775        # completion, because it may not yet have had a chance to run
776        # and gather itself.
777        tasks = tuple(filter(None, (self._writer_task, self._reader_task)))
778        error_pathway = _done(self._reader_task) or _done(self._writer_task)
779        if not tasks:
780            error_pathway |= bool(self._reader) or bool(self._writer)
781
782        try:
783            # Try to flush the writer, if possible.
784            # This *may* cause an error and force us over into the error path.
785            if not error_pathway:
786                await self._bh_flush_writer()
787        except BaseException as err:
788            error_pathway = True
789            emsg = "Failed to flush the writer"
790            self.logger.error("%s: %s", emsg, exception_summary(err))
791            self.logger.debug("%s:\n%s\n", emsg, pretty_traceback())
792            raise
793        finally:
794            # Cancel any still-running tasks (Won't raise):
795            if self._writer_task is not None and not self._writer_task.done():
796                self.logger.debug("Cancelling writer task.")
797                self._writer_task.cancel()
798            if self._reader_task is not None and not self._reader_task.done():
799                self.logger.debug("Cancelling reader task.")
800                self._reader_task.cancel()
801
802            # Close out the tasks entirely (Won't raise):
803            if tasks:
804                self.logger.debug("Waiting for tasks to complete ...")
805                await asyncio.wait(tasks)
806
807            # Lastly, close the stream itself. (*May raise*!):
808            await self._bh_close_stream(error_pathway)
809            self.logger.debug("Disconnected.")
810
811    @bottom_half
812    async def _bh_flush_writer(self) -> None:
813        if not self._writer_task:
814            return
815
816        self.logger.debug("Draining the outbound queue ...")
817        await self._outgoing.join()
818        if self._writer is not None:
819            self.logger.debug("Flushing the StreamWriter ...")
820            await flush(self._writer)
821
822    @bottom_half
823    async def _bh_close_stream(self, error_pathway: bool = False) -> None:
824        # NB: Closing the writer also implicitly closes the reader.
825        if not self._writer:
826            return
827
828        if not is_closing(self._writer):
829            self.logger.debug("Closing StreamWriter.")
830            self._writer.close()
831
832        self.logger.debug("Waiting for StreamWriter to close ...")
833        try:
834            await wait_closed(self._writer)
835        except Exception:  # pylint: disable=broad-except
836            # It's hard to tell if the Stream is already closed or
837            # not. Even if one of the tasks has failed, it may have
838            # failed for a higher-layered protocol reason. The
839            # stream could still be open and perfectly fine.
840            # I don't know how to discern its health here.
841
842            if error_pathway:
843                # We already know that *something* went wrong. Let's
844                # just trust that the Exception we already have is the
845                # better one to present to the user, even if we don't
846                # genuinely *know* the relationship between the two.
847                self.logger.debug(
848                    "Discarding Exception from wait_closed:\n%s\n",
849                    pretty_traceback(),
850                )
851            else:
852                # Oops, this is a brand-new error!
853                raise
854        finally:
855            self.logger.debug("StreamWriter closed.")
856
857    @bottom_half
858    async def _bh_loop_forever(self, async_fn: _TaskFN, name: str) -> None:
859        """
860        Run one of the bottom-half methods in a loop forever.
861
862        If the bottom half ever raises any exception, schedule a
863        disconnect that will terminate the entire loop.
864
865        :param async_fn: The bottom-half method to run in a loop.
866        :param name: The name of this task, used for logging.
867        """
868        try:
869            while True:
870                await async_fn()
871        except asyncio.CancelledError:
872            # We have been cancelled by _bh_disconnect, exit gracefully.
873            self.logger.debug("Task.%s: cancelled.", name)
874            return
875        except BaseException as err:
876            self.logger.log(
877                logging.INFO if isinstance(err, EOFError) else logging.ERROR,
878                "Task.%s: %s",
879                name, exception_summary(err)
880            )
881            self.logger.debug("Task.%s: failure:\n%s\n",
882                              name, pretty_traceback())
883            self._schedule_disconnect()
884            raise
885        finally:
886            self.logger.debug("Task.%s: exiting.", name)
887
888    @bottom_half
889    async def _bh_send_message(self) -> None:
890        """
891        Wait for an outgoing message, then send it.
892
893        Designed to be run in `_bh_loop_forever()`.
894        """
895        msg = await self._outgoing.get()
896        try:
897            await self._send(msg)
898        finally:
899            self._outgoing.task_done()
900
901    @bottom_half
902    async def _bh_recv_message(self) -> None:
903        """
904        Wait for an incoming message and call `_on_message` to route it.
905
906        Designed to be run in `_bh_loop_forever()`.
907        """
908        msg = await self._recv()
909        await self._on_message(msg)
910
911    # --------------------
912    # Section: Message I/O
913    # --------------------
914
915    @upper_half
916    @bottom_half
917    def _cb_outbound(self, msg: T) -> T:
918        """
919        Callback: outbound message hook.
920
921        This is intended for subclasses to be able to add arbitrary
922        hooks to filter or manipulate outgoing messages. The base
923        implementation does nothing but log the message without any
924        manipulation of the message.
925
926        :param msg: raw outbound message
927        :return: final outbound message
928        """
929        self.logger.debug("--> %s", str(msg))
930        return msg
931
932    @upper_half
933    @bottom_half
934    def _cb_inbound(self, msg: T) -> T:
935        """
936        Callback: inbound message hook.
937
938        This is intended for subclasses to be able to add arbitrary
939        hooks to filter or manipulate incoming messages. The base
940        implementation does nothing but log the message without any
941        manipulation of the message.
942
943        This method does not "handle" incoming messages; it is a filter.
944        The actual "endpoint" for incoming messages is `_on_message()`.
945
946        :param msg: raw inbound message
947        :return: processed inbound message
948        """
949        self.logger.debug("<-- %s", str(msg))
950        return msg
951
952    @upper_half
953    @bottom_half
954    async def _readline(self) -> bytes:
955        """
956        Wait for a newline from the incoming reader.
957
958        This method is provided as a convenience for upper-layer
959        protocols, as many are line-based.
960
961        This method *may* return a sequence of bytes without a trailing
962        newline if EOF occurs, but *some* bytes were received. In this
963        case, the next call will raise `EOFError`. It is assumed that
964        the layer 5 protocol will decide if there is anything meaningful
965        to be done with a partial message.
966
967        :raise OSError: For stream-related errors.
968        :raise EOFError:
969            If the reader stream is at EOF and there are no bytes to return.
970        :return: bytes, including the newline.
971        """
972        assert self._reader is not None
973        msg_bytes = await self._reader.readline()
974
975        if not msg_bytes:
976            if self._reader.at_eof():
977                raise EOFError
978
979        return msg_bytes
980
981    @upper_half
982    @bottom_half
983    async def _do_recv(self) -> T:
984        """
985        Abstract: Read from the stream and return a message.
986
987        Very low-level; intended to only be called by `_recv()`.
988        """
989        raise NotImplementedError
990
991    @upper_half
992    @bottom_half
993    async def _recv(self) -> T:
994        """
995        Read an arbitrary protocol message.
996
997        .. warning::
998            This method is intended primarily for `_bh_recv_message()`
999            to use in an asynchronous task loop. Using it outside of
1000            this loop will "steal" messages from the normal routing
1001            mechanism. It is safe to use prior to `_establish_session()`,
1002            but should not be used otherwise.
1003
1004        This method uses `_do_recv()` to retrieve the raw message, and
1005        then transforms it using `_cb_inbound()`.
1006
1007        :return: A single (filtered, processed) protocol message.
1008        """
1009        message = await self._do_recv()
1010        return self._cb_inbound(message)
1011
1012    @upper_half
1013    @bottom_half
1014    def _do_send(self, msg: T) -> None:
1015        """
1016        Abstract: Write a message to the stream.
1017
1018        Very low-level; intended to only be called by `_send()`.
1019        """
1020        raise NotImplementedError
1021
1022    @upper_half
1023    @bottom_half
1024    async def _send(self, msg: T) -> None:
1025        """
1026        Send an arbitrary protocol message.
1027
1028        This method will transform any outgoing messages according to
1029        `_cb_outbound()`.
1030
1031        .. warning::
1032            Like `_recv()`, this method is intended to be called by
1033            the writer task loop that processes outgoing
1034            messages. Calling it directly may circumvent logic
1035            implemented by the caller meant to correlate outgoing and
1036            incoming messages.
1037
1038        :raise OSError: For problems with the underlying stream.
1039        """
1040        msg = self._cb_outbound(msg)
1041        self._do_send(msg)
1042
1043    @bottom_half
1044    async def _on_message(self, msg: T) -> None:
1045        """
1046        Called to handle the receipt of a new message.
1047
1048        .. caution::
1049            This is executed from within the reader loop, so be advised
1050            that waiting on either the reader or writer task will lead
1051            to deadlock. Additionally, any unhandled exceptions will
1052            directly cause the loop to halt, so logic may be best-kept
1053            to a minimum if at all possible.
1054
1055        :param msg: The incoming message, already logged/filtered.
1056        """
1057        # Nothing to do in the abstract case.
1058