Fix asyncio future cancellation in SSE stream to prevent shutdown hangs

Co-authored-by: enzofrnt <[email protected]>

Author: Copilot

django_eventstream/views.py

@@ -223,59 +223,64 @@ async def stream(event_request, listener):
 
             while True:
                 f = asyncio.ensure_future(listener.aevent.wait())
-                while True:
-                    done, _ = await asyncio.wait([f], timeout=20)
-                    if f in done:
-                        break
-                    body = "event: keep-alive\ndata:\n\n"
-                    yield body
+                try:
+                    while True:
+                        done, _ = await asyncio.wait([f], timeout=20)
+                        if f in done:
+                            break
+                        body = "event: keep-alive\ndata:\n\n"
+                        yield body
+
+                    lm.lock.acquire()
+
+                    channel_items = listener.channel_items
+                    overflow = listener.overflow
+                    error_data = listener.error
+
+                    listener.aevent.clear()
+                    listener.channel_items = {}
+                    listener.overflow = False
+
+                    lm.lock.release()
+
+                    body = ""
+                    for channel, items in channel_items.items():
+                        for item in items:
+                            if channel in last_ids:
+                                if item.id is not None:
+                                    last_ids[channel] = item.id
+                                else:
+                                    del last_ids[channel]
+                            if last_ids:
+                                event_id = make_id(last_ids)
+                            else:
+                                event_id = None
+                            body += sse_encode_event(
+                                item.type, item.data, event_id=event_id
+                            )
 
-                lm.lock.acquire()
+                    more = True
 
-                channel_items = listener.channel_items
-                overflow = listener.overflow
-                error_data = listener.error
+                    if error_data:
+                        condition = error_data["condition"]
+                        text = error_data["text"]
+                        extra = error_data.get("extra")
+                        body += sse_encode_error(condition, text, extra=extra)
+                        more = False
 
-                listener.aevent.clear()
-                listener.channel_items = {}
-                listener.overflow = False
+                    if body or not more:
+                        yield body
 
-                lm.lock.release()
+                    if not more:
+                        break
 
-                body = ""
-                for channel, items in channel_items.items():
-                    for item in items:
-                        if channel in last_ids:
-                            if item.id is not None:
-                                last_ids[channel] = item.id
-                            else:
-                                del last_ids[channel]
-                        if last_ids:
-                            event_id = make_id(last_ids)
-                        else:
-                            event_id = None
-                        body += sse_encode_event(
-                            item.type, item.data, event_id=event_id
-                        )
-
-                more = True
-
-                if error_data:
-                    condition = error_data["condition"]
-                    text = error_data["text"]
-                    extra = error_data.get("extra")
-                    body += sse_encode_error(condition, text, extra=extra)
-                    more = False
-
-                if body or not more:
-                    yield body
-
-                if not more:
-                    break
-
-                if overflow:
-                    # check db
-                    break
+                    if overflow:
+                        # check db
+                        break
+                finally:
+                    # Always cancel the future to prevent it from lingering
+                    if not f.done():
+                        f.cancel()
 
             event_request.channel_last_ids = last_ids
     finally:

tests/test_stream.py

@@ -58,6 +58,56 @@ def __assert_all_events_are_retrieved_only_once(self):
             CHANNEL_NAME, EVENTS_LIMIT, limit=EVENTS_LIMIT + 1
         )
 
+    @patch("django_eventstream.eventstream.get_storage")
+    async def test_stream_cancellation_during_wait(self, mock_get_storage):
+        """Test that stream properly cleans up when cancelled during event wait."""
+        mock_get_storage.return_value = self.storage
+        
+        # Create a real listener (not mocked) to test actual wait behavior
+        listener = Listener()
+        
+        request = EventRequest()
+        request.is_next = False
+        request.is_recover = False
+        request.channels = [CHANNEL_NAME]
+        
+        # Get current ID using sync_to_async
+        get_current_id = sync_to_async(self.storage.get_current_id)
+        current_id = await get_current_id(CHANNEL_NAME)
+        request.channel_last_ids = {CHANNEL_NAME: str(current_id)}
+        
+        # Start streaming - this will wait for events since we're caught up
+        stream_task = asyncio.create_task(
+            self.__collect_response(stream(request, listener))
+        )
+        
+        # Give it time to enter the wait loop
+        await asyncio.sleep(0.5)
+        
+        # Cancel the stream
+        stream_task.cancel()
+        
+        try:
+            await stream_task
+            raise ValueError("stream completed unexpectedly")
+        except asyncio.CancelledError:
+            pass
+        
+        # Verify no tasks are left running
+        pending_tasks = [task for task in asyncio.all_tasks() 
+                        if not task.done() and task != asyncio.current_task()]
+        
+        # Allow brief time for cleanup
+        await asyncio.sleep(0.1)
+        
+        # Check again after cleanup time
+        pending_tasks_after = [task for task in asyncio.all_tasks() 
+                              if not task.done() and task != asyncio.current_task()]
+        
+        # The number of pending tasks should not increase after cancellation
+        self.assertLessEqual(len(pending_tasks_after), len(pending_tasks),
+                            "Stream cancellation should not leave lingering tasks")
+
     async def __initialise_test(self, mock_get_storage):
         mock_get_storage.return_value = self.storage