forked from goodboy/tractor
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
parent
f067cf48a7
commit
37ee477aee
|
@ -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'
|
||||||
|
|
Loading…
Reference in New Issue