[Enhancement] Add the spill-event scheduler infrastructure (default off)

Foundation for moving spillable pipeline operators off the busy-poller onto the
pipeline event scheduler. Inert on its own: no operator opts in yet, so every
plan keeps running exactly as before.

- Spill event-notification layer: a SpillEventObservable on the spiller with
  sink/source subscription lists, gated on enable_event_scheduler(); complete_io
  fires point notifications with a query-lifetime guard (GuardedCompletion) that
  closes the decrement->notify use-after-free window.
- block_reason primitive: the BlockReason enum, block_reason_bit(), the
  named<R, Mask>() compile guard, and verify_block_reason_covered() (the
  park-time safety net) with its fail-soft SpillMetrics counter.
- Driver-core interior-block support: the INTERMEDIATE_BLOCK driver state, the
  supports_intermediate_wakeup() gate, the no-workable-pair release predicate,
  and all-operator pending_finish aggregation -- inert unless an interior
  operator declares a wakeup.
- Shared SpillableFlatSinkMixin and the spill-process pump (channel + operator)
  reworked for one-shot notify.

Each per-family operator PR (agg/join/sort) adds its own enable_spill_*_events
kill-switch config flag on top of this base.

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
Signed-off-by: Evgeniy Shishkin <eshishki@gmail.com>

Author: eshishki

be/src/compute_env/spill/spill_components.cpp

@@ -45,6 +45,12 @@ namespace starrocks::spill {
 // implements for SpillerWriter
 Status SpillerWriter::_decrease_running_flush_tasks() {
     if (_running_flush_tasks.fetch_sub(1) == 1) {
+        // Flush-all completion: all flushes are done and the read phase may begin. The callback wrapper
+        // (Spiller::set_flush_all_call_back) runs the client callback, acquires the input stream and
+        // triggers the first restore. The source wakeup must fire after the whole wrapper (so a woken
+        // reader sees the stream), and on the error path too (an error is also a wakeup). That is why the
+        // notify is in a defer around the callback body, not a straight-line call.
+        auto notify = DeferOp([this]() { _spiller->notify_source_observers(); });
         if (_flush_all_callback) {
             RETURN_IF_ERROR(_flush_all_callback());
         }

be/src/compute_env/spill/spill_metrics.cpp

@@ -68,6 +68,9 @@ void SpillMetrics::install(MetricRegistry* registry) {
     registry->register_metric("spill_disk_bytes_used", MetricLabels().add("storage_type", "remote"),
                               _remote_disk_bytes_used.get());
 
+    _parked_with_uncovered_reason_total = std::make_unique<IntCounter>(MetricUnit::OPERATIONS);
+    registry->register_metric("spill_parked_with_uncovered_reason_total", _parked_with_uncovered_reason_total.get());
+
     register_labeled(registry, "local", &_local);
     register_labeled(registry, "remote", &_remote);
 }

be/src/compute_env/spill/spill_metrics.h

@@ -48,11 +48,18 @@ class SpillMetrics {
     IntGauge* local_disk_bytes_used() { return _local_disk_bytes_used.get(); }
     IntGauge* remote_disk_bytes_used() { return _remote_disk_bytes_used.get(); }
 
+    // Counter for the BlockReason check that runs when an operator parks. It increments when a wakeable
+    // operator parks with a block_reason() that its covered_wakeups() mask does not list. In debug builds
+    // this is a DCHECK instead. A non-zero value means the wakeup table has a gap: the operator may sleep
+    // with nobody left to wake it. This is the counter to alert on.
+    IntCounter* parked_with_uncovered_reason_total() { return _parked_with_uncovered_reason_total.get(); }
+
 private:
     LabeledCounters _local;
     LabeledCounters _remote;
     std::unique_ptr<IntGauge> _local_disk_bytes_used;
     std::unique_ptr<IntGauge> _remote_disk_bytes_used;
+    std::unique_ptr<IntCounter> _parked_with_uncovered_reason_total;
     MetricRegistry* _registry = nullptr;
 };
 

be/src/compute_env/spill/spill_observable.h

@@ -0,0 +1,102 @@
+// Copyright 2021-present StarRocks, Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#pragma once
+
+#include <shared_mutex>
+#include <vector>
+
+#include "exec/pipeline/primitives/pipeline_observer.h"
+#include "runtime/runtime_state.h"
+
+namespace starrocks::spill {
+
+// The notification layer a Spiller owns. It keeps two observer lists -- sink-side and source-side -- and
+// wakes them when spill IO completes. A sink-side sleeper (an operator parked OUTPUT_FULL, e.g. on a full
+// writer or a busy spill channel) is woken with sink_trigger(); a source-side sleeper (an operator parked
+// INPUT_EMPTY waiting on restore/flush IO) is woken with source_trigger(). Each list must be woken with its
+// own trigger; the wrong trigger does not wake the sleeper.
+//
+// notify_*_observers() runs under a shared lock, so several notifies can run at once. detach_observers()
+// takes the exclusive lock to clear the lists, so it cannot run while a notify is still iterating -- a
+// notify never touches an observer that detach is removing. The lock is taken once per IO completion (not
+// per chunk) and the fan-out is small (<= 3 observers per spiller), so it is cheap.
+//
+// This is deliberately not pipeline::PipeObservable: its defer_notify_sink() wakes the sink list with
+// source_trigger(), the wrong trigger here, so reusing it would lose every OUTPUT_FULL wakeup.
+class SpillEventObservable {
+public:
+    // Subscribes a sink observer. A no-op in poller mode (event scheduler off): there the driver's observer
+    // is never set, so subscribing a null pointer would crash the later notify. This gate lets the operator
+    // call subscribe unconditionally. Prepare-phase, single-threaded.
+    void subscribe_sink(RuntimeState* state, pipeline::PipelineObserver* o) {
+        if (!state->enable_event_scheduler()) {
+            return;
+        }
+        DCHECK(o != nullptr);
+        std::unique_lock l(_mu);
+        _sink_observers.push_back(o);
+    }
+
+    void subscribe_source(RuntimeState* state, pipeline::PipelineObserver* o) {
+        if (!state->enable_event_scheduler()) {
+            return;
+        }
+        DCHECK(o != nullptr);
+        std::unique_lock l(_mu);
+        _source_observers.push_back(o);
+    }
+
+    // Callable from any thread, but only after the state this predicate reads has been published.
+    void notify_sink_observers() {
+        std::shared_lock l(_mu);
+        for (auto* o : _sink_observers) {
+            o->sink_trigger();
+        }
+    }
+
+    void notify_source_observers() {
+        std::shared_lock l(_mu);
+        for (auto* o : _source_observers) {
+            o->source_trigger();
+        }
+    }
+
+    // Takes the exclusive lock so it cannot race a live notify.
+    //
+    // On the normal spill path detach_observers is deliberately not called: a notify can fire on the
+    // observer of an already-finished driver, and that is safe. Memory-wise the observer is alive because
+    // every notify runs inside the IO task's query pin (the ResourceMemTrackerGuard window), and drivers
+    // are torn down only with the QueryContext; behavior-wise a trigger on a non-blocked driver only sets
+    // need_check_reschedule. The exclusive lock here protects the off-path callers (reset_state / explicit
+    // teardown); the normal spill path never reaches it.
+    void detach_observers() {
+        std::unique_lock l(_mu);
+        _sink_observers.clear();
+        _source_observers.clear();
+    }
+
+    // test accessor: total subscriptions across both lists.
+    size_t observer_count() const {
+        std::shared_lock l(_mu);
+        return _sink_observers.size() + _source_observers.size();
+    }
+
+private:
+    mutable std::shared_mutex _mu;
+    std::vector<pipeline::PipelineObserver*> _sink_observers;
+    std::vector<pipeline::PipelineObserver*> _source_observers;
+};
+
+} // namespace starrocks::spill

be/src/compute_env/spill/spiller.cpp

@@ -43,6 +43,20 @@
 namespace starrocks::spill {
 DEFINE_FAIL_POINT(spill_restore_sleep);
 DEFINE_FAIL_POINT(spill_restore_error);
+DEFINE_FAIL_POINT(spill_submit_error);
+DEFINE_FAIL_POINT(spill_flush_block);
+DEFINE_FAIL_POINT(spill_restore_block);
+
+#ifdef FIU_ENABLE
+failpoint::OneToAnyBarrier& spill_flush_block_barrier() {
+    static failpoint::OneToAnyBarrier barrier;
+    return barrier;
+}
+failpoint::OneToAnyBarrier& spill_restore_block_barrier() {
+    static failpoint::OneToAnyBarrier barrier;
+    return barrier;
+}
+#endif
 
 TQueryType::type spill_query_type(RuntimeState* state) {
     return state->query_options().query_type;
@@ -182,6 +196,12 @@ void Spiller::update_spilled_task_status(Status&& st) {
 }
 
 Status Spiller::reset_state(RuntimeState* state) {
+    // prepare() reallocates the writer/reader, but in-flight flush/restore tasks captured the previous raw
+    // pointers. Refuse the reset while IO is in flight instead of swapping objects out from under a live
+    // task; the caller (query-cache refill) retries once the tasks drain.
+    if (has_running_io_tasks()) {
+        return Status::InternalError("reset_state called while spill IO tasks are still in flight");
+    }
     _spilled_append_rows = 0;
     _restore_read_rows = 0;
     RETURN_IF_ERROR(prepare(state));

be/src/compute_env/spill/spiller.h

@@ -35,10 +35,12 @@
 #include "compute_env/spill/serde.h"
 #include "compute_env/spill/spill_components.h"
 #include "compute_env/spill/spill_metrics.h"
+#include "compute_env/spill/spill_observable.h"
 #include "compute_env/spill/spiller_factory.h"
 #include "compute_env/spill/task_executor.h"
 #include "fs/fs.h"
 #include "gen_cpp/InternalService_types.h"
+#include "gutil/macros.h"
 #include "runtime/mem_tracker.h"
 #include "runtime/runtime_state_fwd.h"
 
@@ -169,6 +171,45 @@ class Spiller : public std::enable_shared_from_this<Spiller> {
 
     const SpillProcessMetrics& metrics() { return _metrics; }
 
+    // Notification layer. It lives on the Spiller and survives reset_state, so the observer lists outlive
+    // reallocation of the writer/reader.
+    SpillEventObservable& observable() { return _observable; }
+
+    void notify_sink_observers() { _observable.notify_sink_observers(); }
+    void notify_source_observers() { _observable.notify_source_observers(); }
+
+    // The single place that owns the completion order for spiller IO:
+    //   publish -> decrement the in-flight count -> notify.
+    // The notify runs after both the predicates (published by publish_fn) and the count are visible, so one
+    // notify serves every sleeper: the OUTPUT_FULL/INPUT_EMPTY predicate sleepers re-check predicates that
+    // are already open, and a PENDING_FINISH driver gating on the count sees it already at its new value.
+    // The caller must keep the notified observers alive across the call: a task body calls this through
+    // defer_complete_io (the notify runs inside the task's query-lifetime guard window), and the
+    // submit-failure compensation calls it directly on the pipeline thread, where the executing driver
+    // itself keeps the query alive. Every completion path goes through here.
+    template <class PublishFn>
+    void complete_io(PublishFn&& publish_fn, bool wake_sink) {
+        std::forward<PublishFn>(publish_fn)();
+        decrease_in_flight_io();
+        if (wake_sink) {
+            notify_sink_observers();
+        }
+        notify_source_observers();
+    }
+
+    // RAII form of complete_io for IO task bodies; see IOCompletion below. The guard must already be
+    // scoped_begin'd; the returned object runs the completion and then guard.scoped_end() at scope exit.
+    template <class PublishFn, class Guard>
+    [[nodiscard]] auto defer_complete_io(PublishFn publish_fn, bool wake_sink, const Guard& guard);
+
+    // Counts in-flight IO across writer (flush) and reader (restore) tasks. It is in addition to the
+    // per-writer/reader counters, which stay as lifetime gates for transient readers. It is incremented
+    // before submit at every choke point, and decremented in the same defer as the per-task counter.
+    void increase_in_flight_io() { _in_flight_io.fetch_add(1, std::memory_order_acq_rel); }
+    void decrease_in_flight_io() { _in_flight_io.fetch_sub(1, std::memory_order_acq_rel); }
+
+    bool has_running_io_tasks() const { return _in_flight_io.load(std::memory_order_acquire) > 0; }
+
     // set partitions for spiller only works when spiller has partitioned spill writer
     void set_partition(const std::vector<const SpillPartitionInfo*>& parititons);
     // init partition by `num_partitions`
@@ -303,6 +344,54 @@ class Spiller : public std::enable_shared_from_this<Spiller> {
     size_t _max_sorted_block_cnt = 0;
     std::atomic_bool _is_cancel = false;
     std::atomic_bool _global_spill_triggered{false};
+
+    SpillEventObservable _observable;
+    // Count of in-flight flush/restore IO tasks. It survives reset_state (it lives on the Spiller, not on
+    // the writer/reader), so it can gate teardown and re-prepare against tasks that captured raw
+    // writer/reader pointers.
+    std::atomic<int64_t> _in_flight_io{0};
+};
+
+// RAII completion for an IO task body, returned by Spiller::defer_complete_io. At scope exit it runs the
+// completion (complete_io: publish -> decrement -> notify) and then the guard's scoped_end, in that order.
+// The order matters for correctness, not for style: scoped_end releases the query-lifetime pin the guard
+// took in scoped_begin, and that pin is what keeps the notified observers alive -- if scoped_end ran first,
+// the notify could dereference freed observers. Do not reorder. cancel_for_yield() drops the completion AND
+// marks the yield context unfinished in one call (on a yield-resubmit the task lives on, keeps its
+// in-flight increment, and re-enters scoped_begin on its next run); scoped_end still runs to balance this
+// invocation's begin. The two cancels are coupled on purpose: cancelling the completion without cancelling
+// the yield defer would let the executor drop the task with the in-flight count still held.
+template <class PublishFn, class Guard>
+class IOCompletion {
+public:
+    IOCompletion(Spiller* spiller, PublishFn publish_fn, bool wake_sink, const Guard& guard)
+            : _spiller(spiller), _publish_fn(std::move(publish_fn)), _wake_sink(wake_sink), _guard(guard) {}
+    ~IOCompletion() noexcept {
+        if (!_cancelled) {
+            _spiller->complete_io(_publish_fn, _wake_sink);
+        }
+        _guard.scoped_end();
+    }
+    template <class YieldDefer>
+    void cancel_for_yield(YieldDefer& yield_defer) {
+        _cancelled = true;
+        yield_defer.cancel();
+    }
+    DISALLOW_COPY_AND_MOVE(IOCompletion);
+
+private:
+    Spiller* _spiller;
+    PublishFn _publish_fn;
+    const bool _wake_sink;
+    const Guard& _guard;
+    bool _cancelled = false;
 };
 
+template <class PublishFn, class Guard>
+[[nodiscard]] auto Spiller::defer_complete_io(PublishFn publish_fn, bool wake_sink, const Guard& guard) {
+    // Returned as a prvalue: C++17 guaranteed elision constructs it in the caller's variable, so the
+    // deleted copy/move constructors are never used.
+    return IOCompletion<PublishFn, Guard>(this, std::move(publish_fn), wake_sink, guard);
+}
+
 } // namespace starrocks::spill