handle EndTaskFrame, StopTaskFrame and CancelTaskFrame downstream
EndTaskFrame and StopTaskFrame are now ControlFrames instead of SystemFrames, so they flow through the pipeline and queue behind pending work. This prevents races where EndFrame could overtake in-flight frames (e.g. function call responses). CancelTaskFrame and InterruptionTaskFrame remain SystemFrames (via new TaskSystemFrame base): since they need immediate propagation. The sink now catches EndTaskFrame, StopTaskFrame and CancelTaskFrame downstream and re-queues it upstream to the task, ensuring the full pipeline drains before shutdown begins.
This commit is contained in:
@@ -1742,7 +1742,7 @@ class ServiceSwitcherRequestMetadataFrame(ControlFrame):
|
|||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class TaskFrame(SystemFrame):
|
class TaskFrame(ControlFrame):
|
||||||
"""Base frame for task frames.
|
"""Base frame for task frames.
|
||||||
|
|
||||||
This is a base class for frames that are meant to be sent and handled
|
This is a base class for frames that are meant to be sent and handled
|
||||||
@@ -1756,7 +1756,21 @@ class TaskFrame(SystemFrame):
|
|||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class EndTaskFrame(TaskFrame):
|
class TaskSystemFrame(SystemFrame):
|
||||||
|
"""Base frame for task system frames.
|
||||||
|
|
||||||
|
This is a base class for frames that are meant to be sent and handled
|
||||||
|
upstream by the pipeline task. This might result in a corresponding frame
|
||||||
|
sent downstream (e.g. `InterruptionTaskFrame` / `InterruptionFrame` or
|
||||||
|
`EndTaskFrame` / `EndFrame`).
|
||||||
|
|
||||||
|
"""
|
||||||
|
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class EndTaskFrame(TaskFrame, UninterruptibleFrame):
|
||||||
"""Frame to request graceful pipeline task closure.
|
"""Frame to request graceful pipeline task closure.
|
||||||
|
|
||||||
This is used to notify the pipeline task that the pipeline should be
|
This is used to notify the pipeline task that the pipeline should be
|
||||||
@@ -1774,7 +1788,20 @@ class EndTaskFrame(TaskFrame):
|
|||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class CancelTaskFrame(TaskFrame):
|
class StopTaskFrame(TaskFrame, UninterruptibleFrame):
|
||||||
|
"""Frame to request pipeline task stop while keeping processors running.
|
||||||
|
|
||||||
|
This is used to notify the pipeline task that it should be stopped as
|
||||||
|
soon as possible (flushing all the queued frames) but that the pipeline
|
||||||
|
processors should be kept in a running state. This frame should be pushed
|
||||||
|
upstream.
|
||||||
|
"""
|
||||||
|
|
||||||
|
pass
|
||||||
|
|
||||||
|
|
||||||
|
@dataclass
|
||||||
|
class CancelTaskFrame(TaskSystemFrame):
|
||||||
"""Frame to request immediate pipeline task cancellation.
|
"""Frame to request immediate pipeline task cancellation.
|
||||||
|
|
||||||
This is used to notify the pipeline task that the pipeline should be
|
This is used to notify the pipeline task that the pipeline should be
|
||||||
@@ -1792,20 +1819,7 @@ class CancelTaskFrame(TaskFrame):
|
|||||||
|
|
||||||
|
|
||||||
@dataclass
|
@dataclass
|
||||||
class StopTaskFrame(TaskFrame):
|
class InterruptionTaskFrame(TaskSystemFrame):
|
||||||
"""Frame to request pipeline task stop while keeping processors running.
|
|
||||||
|
|
||||||
This is used to notify the pipeline task that it should be stopped as
|
|
||||||
soon as possible (flushing all the queued frames) but that the pipeline
|
|
||||||
processors should be kept in a running state. This frame should be pushed
|
|
||||||
upstream.
|
|
||||||
"""
|
|
||||||
|
|
||||||
pass
|
|
||||||
|
|
||||||
|
|
||||||
@dataclass
|
|
||||||
class InterruptionTaskFrame(TaskFrame):
|
|
||||||
"""Frame indicating the pipeline should be interrupted.
|
"""Frame indicating the pipeline should be interrupted.
|
||||||
|
|
||||||
This frame should be pushed upstream to indicate the pipeline should be
|
This frame should be pushed upstream to indicate the pipeline should be
|
||||||
|
|||||||
@@ -876,22 +876,22 @@ class PipelineTask(BasePipelineTask):
|
|||||||
|
|
||||||
if isinstance(frame, EndTaskFrame):
|
if isinstance(frame, EndTaskFrame):
|
||||||
# Tell the task we should end nicely.
|
# Tell the task we should end nicely.
|
||||||
logger.debug(f"{self}: received end task frame {frame}")
|
logger.debug(f"{self}: received end task frame upstream {frame}")
|
||||||
await self.queue_frame(EndFrame(reason=frame.reason))
|
await self.queue_frame(EndFrame(reason=frame.reason))
|
||||||
elif isinstance(frame, CancelTaskFrame):
|
elif isinstance(frame, CancelTaskFrame):
|
||||||
# Tell the task we should end right away.
|
# Tell the task we should end right away.
|
||||||
logger.debug(f"{self}: received cancel task frame {frame}")
|
logger.debug(f"{self}: received cancel task frame upstream {frame}")
|
||||||
await self.queue_frame(CancelFrame(reason=frame.reason))
|
await self.queue_frame(CancelFrame(reason=frame.reason))
|
||||||
elif isinstance(frame, StopTaskFrame):
|
elif isinstance(frame, StopTaskFrame):
|
||||||
# Tell the task we should stop nicely.
|
# Tell the task we should stop nicely.
|
||||||
logger.debug(f"{self}: received stop task frame {frame}")
|
logger.debug(f"{self}: received stop task frame upstream {frame}")
|
||||||
await self.queue_frame(StopFrame())
|
await self.queue_frame(StopFrame())
|
||||||
elif isinstance(frame, InterruptionTaskFrame):
|
elif isinstance(frame, InterruptionTaskFrame):
|
||||||
# Tell the task we should interrupt the pipeline. Note that we are
|
# Tell the task we should interrupt the pipeline. Note that we are
|
||||||
# bypassing the push queue and directly queue into the
|
# bypassing the push queue and directly queue into the
|
||||||
# pipeline. This is in case the push task is blocked waiting for a
|
# pipeline. This is in case the push task is blocked waiting for a
|
||||||
# pipeline-ending frame to finish traversing the pipeline.
|
# pipeline-ending frame to finish traversing the pipeline.
|
||||||
logger.debug(f"{self}: received interruption task frame {frame}")
|
logger.debug(f"{self}: received interruption task frame upstream {frame}")
|
||||||
await self._pipeline.queue_frame(InterruptionFrame())
|
await self._pipeline.queue_frame(InterruptionFrame())
|
||||||
elif isinstance(frame, ErrorFrame):
|
elif isinstance(frame, ErrorFrame):
|
||||||
await self._call_event_handler("on_pipeline_error", frame)
|
await self._call_event_handler("on_pipeline_error", frame)
|
||||||
@@ -934,6 +934,18 @@ class PipelineTask(BasePipelineTask):
|
|||||||
self._pipeline_end_event.set()
|
self._pipeline_end_event.set()
|
||||||
elif isinstance(frame, HeartbeatFrame):
|
elif isinstance(frame, HeartbeatFrame):
|
||||||
await self._heartbeat_queue.put(frame)
|
await self._heartbeat_queue.put(frame)
|
||||||
|
elif isinstance(frame, EndTaskFrame):
|
||||||
|
logger.debug(f"{self}: received end task frame downstream {frame}")
|
||||||
|
await self.queue_frame(EndTaskFrame(reason=frame.reason), FrameDirection.UPSTREAM)
|
||||||
|
elif isinstance(frame, StopTaskFrame):
|
||||||
|
logger.debug(f"{self}: received stop task frame downstream {frame}")
|
||||||
|
await self.queue_frame(StopTaskFrame(), FrameDirection.UPSTREAM)
|
||||||
|
elif isinstance(frame, CancelTaskFrame):
|
||||||
|
logger.debug(f"{self}: received cancel task frame downstream {frame}")
|
||||||
|
await self.queue_frame(CancelTaskFrame(reason=frame.reason), FrameDirection.UPSTREAM)
|
||||||
|
elif isinstance(frame, InterruptionTaskFrame):
|
||||||
|
logger.debug(f"{self}: received interruption task frame downstream {frame}")
|
||||||
|
await self.queue_frame(InterruptionTaskFrame(), FrameDirection.UPSTREAM)
|
||||||
|
|
||||||
async def _heartbeat_push_handler(self):
|
async def _heartbeat_push_handler(self):
|
||||||
"""Push heartbeat frames at regular intervals."""
|
"""Push heartbeat frames at regular intervals."""
|
||||||
|
|||||||
Reference in New Issue
Block a user