forked from goodboy/tractor
1
0
Fork 0

Only send stop msg if not received from far end

transport_hardening
Tyler Goodlet 2021-05-11 23:42:34 -04:00
parent 9d5f193dfd
commit 7192654d51
1 changed files with 8 additions and 2 deletions

View File

@ -53,7 +53,8 @@ class ReceiveMsgStream(trio.abc.ReceiveChannel):
# delegate directly to underlying mem channel
def receive_nowait(self):
return self._rx_chan.receive_nowait()
msg = self._rx_chan.receive_nowait()
return msg['yield']
async def receive(self):
try:
@ -106,6 +107,7 @@ class ReceiveMsgStream(trio.abc.ReceiveChannel):
# when the send is closed we assume the stream has
# terminated and signal this local iterator to stop
await self.aclose()
# await self._ctx.send_stop()
raise StopAsyncIteration
@ -326,10 +328,14 @@ class Context:
try:
yield rchan
finally:
except trio.EndOfChannel:
raise
else:
# signal ``StopAsyncIteration`` on far end.
await self.send_stop()
finally:
if self._portal:
self._portal._streams.remove(rchan)