Refactor

Author: qianduoduo0904

mars/oscar/backends/pool.py

@@ -494,36 +494,44 @@ async def stop(self):
     def stopped(self) -> bool:
         return self._stopped.is_set()
 
+    @classmethod
+    async def _recv_message(cls, channel: Channel):
+        try:
+            return await channel.recv()
+        except EOFError:
+            # no data to read, check channel
+            try:
+                await channel.close()
+            except (ConnectionError, EOFError):
+                # close failed, ignore
+                pass
+
     async def on_new_channel(self, channel: Channel):
-        is_transfer = False
+        message = await self._recv_message(channel)
+        if (
+            message.message_type == MessageType.control
+            and message.control_message_type == ControlMessageType.switch_to_transfer
+        ):
+            # switch this channel to data transfer channel
+            # the channel will be handed over to TransferServer
+            # and this loop will exit
+            return await TransferServer.handle_transfer_channel(channel, self._stopped)
+        else:
+            asyncio.create_task(self.process_message(message, channel))
+            # delete to release the reference of message
+            del message
+            await asyncio.sleep(0)
+
+        # continue to keep processing messages
         while not self._stopped.is_set():
-            try:
-                message = await channel.recv()
-            except EOFError:
-                # no data to read, check channel
-                try:
-                    await channel.close()
-                except (ConnectionError, EOFError):
-                    # close failed, ignore
-                    pass
-                return
-            if (
-                message.message_type == MessageType.control
-                and message.control_message_type
-                == ControlMessageType.switch_to_transfer
-            ):
-                # switch this channel to data transfer channel
-                # the channel will be handed over to TansferServer
-                is_transfer = True
+            message = await self._recv_message(channel)
+            if message is None:
                 break
             asyncio.create_task(self.process_message(message, channel))
             # delete to release the reference of message
             del message
             await asyncio.sleep(0)
 
-        if is_transfer:
-            await TransferServer.handle_transfer_channel(channel)
-
     async def __aenter__(self):
         await self.start()
         return self

mars/oscar/backends/transfer.py

@@ -193,13 +193,18 @@ async def _catch_error(channel: Channel, message_id: bytes):
 
 class TransferServer:
     @classmethod
-    async def handle_transfer_channel(cls, channel: Channel):
-        while True:
+    async def handle_transfer_channel(cls, channel: Channel, stopped: asyncio.Event):
+        while not stopped.is_set():
             try:
                 message = await channel.recv()
-            except EOFError:
-                # closed
-                break
+            except EOFError:  # pragma: no cover
+                # no data to read, check channel
+                try:
+                    await channel.close()
+                except (ConnectionError, EOFError):
+                    # close failed, ignore
+                    pass
+                return
             assert message.message_type in (
                 MessageType.copyto_buffers,
                 MessageType.copyto_fileobjects,