Clean up parallel SFTP requests on cancellation

asyncssh/sftp.py

@@ -753,33 +753,38 @@ async def iter(self) -> AsyncIterator[Tuple[int, _T]]:
 
         self._start_tasks()
 
-        while self._pending:
-            done, self._pending = await asyncio.wait(
-                self._pending, return_when=asyncio.FIRST_COMPLETED)
+        try:
+            while self._pending:
+                done, _ = await asyncio.wait(
+                    self._pending, return_when=asyncio.FIRST_COMPLETED)
 
-            exceptions = []
+                exceptions = []
 
-            for task in done:
-                try:
-                    offset, size, count, result = task.result()
-                    yield offset, result
+                for task in done:
+                    self._pending.remove(task)
 
-                    if count and count < size:
-                        self._pending.add(asyncio.ensure_future(
-                            self._start_task(offset+count, size-count)))
-                except SFTPEOFError:
-                    self._bytes_left = 0
-                except (OSError, SFTPError) as exc:
-                    exceptions.append(exc)
+                    try:
+                        offset, size, count, result = task.result()
+                        yield offset, result
 
-            if exceptions:
-                for task in self._pending:
-                    task.cancel()
+                        if count and count < size:
+                            self._pending.add(asyncio.ensure_future(
+                                self._start_task(offset+count, size-count)))
+                    except SFTPEOFError:
+                        self._bytes_left = 0
+                    except (OSError, SFTPError) as exc:
+                        exceptions.append(exc)
 
-                await asyncio.gather(*self._pending, return_exceptions=True)
-                raise exceptions[0]
+                if exceptions:
+                    raise exceptions[0]
+
+                self._start_tasks()
+        finally:
+            for task in self._pending:
+                task.cancel()
 
-            self._start_tasks()
+            await asyncio.gather(*self._pending, return_exceptions=True)
+            self._pending.clear()
 
 
 class _SFTPFileReader(_SFTPParallelIO[bytes]):

tests/test_sftp.py

@@ -5007,6 +5007,94 @@ async def test_open_error_v6(self, sftp):
             await sftp.open('ENXIO')
 
 
+class _TestSFTPCancelParallelIO(_CheckSFTP):
+    """Unit tests for cancellation of parallel SFTP requests"""
+
+    @classmethod
+    async def start_server(cls):
+        """Start an SFTP server for the tests to use"""
+
+        return await cls.create_server(sftp_factory=True)
+
+    async def _check_cancel(self, sftp, operation):
+        """Cancel blocked requests, then reuse the file and connection"""
+
+        entered = asyncio.Event()
+        release = asyncio.Event()
+        pending = set()
+        data = b'0123456789abcdef'
+        method = 'write' if operation == 'write' else 'read'
+        original = getattr(SFTPServer, method)
+
+        async def delayed_io(server, *args):
+            """Hold the response until the client cancels its requests"""
+
+            entered.set()
+            await release.wait()
+            return original(server, *args)
+
+        self._create_file('file', data)
+
+        try:
+            async with sftp.open('file', 'r+b', block_size=4,
+                                 max_requests=4) as file:
+                with patch.object(SFTPServer, method, delayed_io):
+                    if operation == 'write':
+                        request = file.write(data)
+                    elif operation == 'read_parallel':
+                        reader = await file.read_parallel(len(data))
+                        request = reader.__anext__()
+                    else:
+                        request = file.read(len(data))
+
+                    try:
+                        task = asyncio.ensure_future(request)
+                        await asyncio.wait_for(entered.wait(), 5)
+                        pending = {task for task in asyncio.all_tasks()
+                                   if getattr(task.get_coro(), '__qualname__', '') ==
+                                   '_SFTPParallelIO._start_task'}
+                        self.assertEqual(len(pending), 4)
+                        task.cancel()
+
+                        with self.assertRaises(asyncio.CancelledError):
+                            await task
+
+                        completed = all(task.done() for task in pending)
+                        release.set()
+                        self.assertTrue(completed)
+                    finally:
+                        release.set()
+
+                await file.seek(0)
+                self.assertEqual(await file.read(), data)
+        finally:
+            release.set()
+
+            for task in pending:
+                task.cancel()
+
+            await asyncio.gather(*pending, return_exceptions=True)
+            remove('file')
+
+    @sftp_test
+    async def test_read_cancel(self, sftp):
+        """Test cancelling a parallel file read"""
+
+        await self._check_cancel(sftp, 'read')
+
+    @sftp_test
+    async def test_write_cancel(self, sftp):
+        """Test cancelling a parallel file write"""
+
+        await self._check_cancel(sftp, 'write')
+
+    @sftp_test
+    async def test_read_parallel_cancel(self, sftp):
+        """Test cancelling an asynchronous iterator of file reads"""
+
+        await self._check_cancel(sftp, 'read_parallel')
+
+
 class _TestSFTPIOError(_CheckSFTP):
     """Unit test for SFTP server returning file I/O error"""