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