Let `MsgStream.receive_nowait()` take in msg key list

Call it `allow_msg_keys: list[str] = ['yield']` and set it to accept
`['yield', 'return']` from the drain loop in `.aclose()`. Only pass the
last key error to `_raise_from_no_key_in_msg()` in the fall-through
case.

Somehow this seems to prevent all the intermittent test failures i was
seeing in local runs including when running the entire suite all in
sequence; i ain't complaining B)
modden_spawn_from_client_req
Tyler Goodlet 2024-03-11 10:20:55 -04:00
parent f067cf48a7
commit 37ee477aee
1 changed files with 26 additions and 14 deletions

View File

@ -90,19 +90,29 @@ class MsgStream(trio.abc.Channel):
self._closed: bool|trio.ClosedResourceError = False self._closed: bool|trio.ClosedResourceError = False
# delegate directly to underlying mem channel # delegate directly to underlying mem channel
def receive_nowait(self): def receive_nowait(
msg = self._rx_chan.receive_nowait() self,
try: allow_msg_keys: list[str] = ['yield'],
return msg['yield'] ):
except KeyError as kerr: msg: dict = self._rx_chan.receive_nowait()
_raise_from_no_key_in_msg( for (
ctx=self._ctx, i,
msg=msg, key,
src_err=kerr, ) in enumerate(allow_msg_keys):
log=log, try:
expect_key='yield', return msg[key]
stream=self, except KeyError as kerr:
) if i < (len(allow_msg_keys) - 1):
continue
_raise_from_no_key_in_msg(
ctx=self._ctx,
msg=msg,
src_err=kerr,
log=log,
expect_key=key,
stream=self,
)
async def receive(self): async def receive(self):
''' '''
@ -263,7 +273,9 @@ class MsgStream(trio.abc.Channel):
drained: list[Exception|dict] = [] drained: list[Exception|dict] = []
while not drained: while not drained:
try: try:
maybe_final_msg = self.receive_nowait() maybe_final_msg = self.receive_nowait(
allow_msg_keys=['yield', 'return'],
)
if maybe_final_msg: if maybe_final_msg:
log.debug( log.debug(
'Drained un-processed stream msg:\n' 'Drained un-processed stream msg:\n'