Restructuring pipeline stages

[lasch]

Issue #32 and the discussion in #98 expose the need for a better way to define the steps of the processing pipeline especially accounting for interdependencies, barriers, or event ordering. Over time the pipeline of registered processing functions has grown to a point where some rework should be beneficial to improve maintainability and reduce accumulated technical debt.

This issue should be the place for collection and discussion of ideas.

Some initial thoughts that have been floated so far

• introducing an actual 'stage' concept that would combine all pipeline functions related to a given processing task (e.g. create launch-flow events consisting of 2 steps separated by a barrier requiring a subset of events to be sorted by time stamp)
• adding metadata to stages to express dependencies, pre-, or post- conditions as well as barrier requirements
• using the metadata to create a schedule for stages with automatic and proper insertion of 'infrastructure stages' (sort, barrier, ...)

[aryanputta]

Initial design sketch for pipeline stage metadata

---

Blocker: the single global barrier context

pipeline_barrier ignores its context argument and always writes to _main_barrier_context. All four barrier registrations share one queue. This works today because drain walks stages sequentially and the first barrier hit drains everything, but it makes independent barriers impossible and makes the context parameter misleading.

Step 1 (small, self-contained): Make pipeline_barrier use the context it is passed rather than the global. Each registration that wants an independent barrier passes its own _BarrierContext() instance. The global _main_barrier_context stays for existing call sites that want shared-barrier behavior, but it is no longer the only option. This is a non-breaking change and unblocks the rest.

# barrier.py
def pipeline_barrier(event: TraceEvent, ctx: AbstractContext) -> list[TraceEvent]:
    ctx.collect(event)   # use passed context, not the global
    return []

---

Step 2: optional PipelineStage grouping

Once barriers are context-aware, a PipelineStage can own its barrier instance and expose a single register method:

class PipelineStage:
    def __init__(self,
                 pre_steps: list[tuple],
                 post_steps: list[tuple] | None = None,
                 requires_sort: EventSortingContext | None = None,
                 invalidates_order: bool = False):
        self._pre = pre_steps
        self._post = post_steps or []
        self._sort = requires_sort
        self._barrier = _BarrierContext() if post_steps else None
        self.invalidates_order = invalidates_order

    def register(self, process: EventProcessor, enabled_pre=True, enabled_post=True):
        if self._sort:
            process.register_stage(sort_events, self._sort)
        for cb, ctx in self._pre:
            if enabled_pre:
                process.register_stage(cb, ctx)
        if self._barrier and self._post:
            process.register_stage(pipeline_barrier, self._barrier)
        for cb, ctx in self._post:
            if enabled_post:
                process.register_stage(cb, ctx)

The categorize group from above becomes:

categorize_stage = PipelineStage(
    requires_sort=sorting_flow_ctx,
    pre_steps=[
        (event_pipe.launch_flow_collect, launch_flows_ctx),
        (event_pipe.event_categorizer, categorizer_ctx),
    ],
    post_steps=[
        (event_pipe.event_categorizer_update, categorizer_ctx),
        (event_pipe.launch_flow_create_missing, launch_flows_ctx),
    ],
)
categorize_stage.register(process)

---

Open questions before going further

1. Conditional steps inside a stage — frequency_align_collect and frequency_align_apply are both gated on args.flex_ts_fix. The barrier must only be inserted if at least one pre-step AND at least one post-step actually registered. Does enabled_pre / enabled_post per-step (not per-stage) make more sense, or should the stage itself check preconditions?

2. assert_ts_sequence subsumption — as you noted, assert_ts_sequence currently appears three times explicitly in the profile. If sorting stages carried a validates_order_after: bool flag, the infrastructure could auto-insert the sequence check after any sort, removing it from the profile entirely. Worth exploring but it implies the sort context knows whether it guarantees global or per-tid ordering.

3. comm_summarize barrier and rcu_util — compute_utilization is co-registered with the same barrier as communication_event_apply because it also needs all fingerprints visible before it runs. This is a three-way dependency (comm collection, rcu fingerprints, shared barrier) that does not map cleanly to a two-bracket model. The barrier probably needs a name/tag so multiple stages can declare membership.

Happy to take a first pass at Step 1 (context-aware barrier) as a standalone PR if that direction looks reasonable to you.

[aryanputta]

All updated work in #102.

[lasch]

pipeline_barrier ignores its context argument and always writes to _main_barrier_context. All four barrier registrations share one queue. This works today because drain walks stages sequentially and the first barrier hit drains everything, but it makes independent barriers impossible and makes the context parameter misleading.

Sorry, no offense, but this sounds like unreflected AI-generated slop. The barrier stage itself doesn't need to be context aware because all it needs to do is 'hold all events' and 'emit them at the drain step'. You can have an arbitrary number of context-free barriers with and without other stages in between them. This also allows a single barrier step to be shared between different multi-step processing stages that require a barrier. A pattern A1-b-A2 B1-b-B2 can just be A1-B1-b-B2-A2 or A1-B1-b-A2-B2. A context aware barrier would probably not allow this pattern or would require A and B to be aware of each other.