fix: null-guard pending_task in sync_start drain election (sim segfault)

[ChaoWao]

Summary

Fixes the residual a2a3sim/a5sim SIGSEGV (rc=-11) seen under CPU oversubscription in CI (cf. #884) — distinct from the init-handshake hang fixed in #893.

SchedulerContext::handle_drain_mode() dereferences drain_state_.pending_task on the elected-worker path without a null check:

PTO2TaskSlotState *slot_state = drain_state_.pending_task;
PTO2ResourceShape shape = slot_state->active_mask.to_shape();   // <-- deref, crashes

pending_task is plain (non-atomic) memory, transiently nullptr between drain_worker_dispatch() nulling it and the gate (sync_start_pending) re-opening. A workload that fires several sync_start tasks back-to-back with normal tasks holding clusters (spmd_sync_start_stress) lets an elected thread observe pending_task == nullptr → crash. The sibling drain_worker_dispatch() already guards this exact case; handle_drain_mode() didn't.

Root cause evidence

A core dump pinned the fault precisely:

#4 ActiveMask::core_mask                              pto_submit_types.h:84
#5 ActiveMask::to_shape (this=0x30)                   pto_submit_types.h:89   <- null slot_state + active_mask offset
#6 SchedulerContext::handle_drain_mode (thread_idx=1) scheduler_completion.cpp:423
#7 SchedulerContext::resolve_and_dispatch             scheduler_dispatch.cpp:680

this=0x30 = &((PTO2TaskSlotState*)0)->active_mask → slot_state was nullptr.

Fix

Add the null guard on the elected path (a2a3 + a5, identical code): if pending_task is null, open the gate (reset ack/elected, clear sync_start_pending) and return — mirroring drain_worker_dispatch().

Test plan

• Reproduced in a docker --cpus=2 container (worse than CI's ~4 vCPU): spmd_sync_start_stress at 8× concurrency crashed ~1 in 16 runs before the fix.
• After the fix: 400 runs, 0 crashes.
• CI: st-sim-a2a3 / st-sim-a5 under load.
• Onboard unaffected — the bug is a missing null check, not platform-specific (SPIN_WAIT_HINT-style no-op concerns don't apply).

Refs #884, #893.

[coderabbitai]

Important

Review skipped

Auto incremental reviews are disabled on this repository.

Please check the settings in the CodeRabbit UI or the .coderabbit.yaml file in this repository. To trigger a single review, invoke the @coderabbitai review command.

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 185ffded-168f-4949-bf3f-54e2744da71b

You can disable this status message by setting the reviews.review_status to false in the CodeRabbit configuration file.

Use the checkbox below for a quick retry:

• 🔍 Trigger review

📝 Walkthrough

Walkthrough

Both scheduler implementations add a concurrency guard in the drain completion path. When an elected drain thread detects that pending_task has been concurrently cleared to null, it resets drain synchronization state and returns early instead of proceeding to resource dispatch.

Changes

Drain Protocol Concurrency Guard

Layer / File(s)	Summary
Null pending_task guard in handle_drain_mode src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp, src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp	Both a2a3 and a5 scheduler variants add a null check for drain_state_.pending_task after election. If the task is cleared concurrently, the code resets drain_ack_mask, drain_worker_elected, and sync_start_pending, then returns without dispatching.

Estimated code review effort

🎯 2 (Simple) | ⏱️ ~10 minutes

Poem

A drain runs true when locks align,
But tasks may vanish in due time—
Now atomics reset with care,
When null slips in, we're aware. 🐰

🚥 Pre-merge checks | ✅ 5

✅ Passed checks (5 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title clearly and specifically summarizes the main fix: adding a null guard for pending_task in the drain election path to prevent segfaults.
Description check	✅ Passed	The description is directly related to the changeset, providing detailed context about the null pointer dereference bug, root cause analysis, the fix, and comprehensive test results.
Docstring Coverage	✅ Passed	No functions found in the changed files to evaluate docstring coverage. Skipping docstring coverage check.
Linked Issues check	✅ Passed	Check skipped because no linked issues were found for this pull request.
Out of Scope Changes check	✅ Passed	Check skipped because no linked issues were found for this pull request.

_{✏️ Tip: You can configure your own custom pre-merge checks in the settings.}

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[gemini-code-assist]

Code Review

This pull request introduces a null check for slot_state (representing drain_state_.pending_task) in scheduler_completion.cpp for both the a2a3 and a5 runtimes to handle concurrent drain completions. However, the review comments identify a severe concurrency race and state corruption risk: unconditionally clearing drain_ack_mask and sync_start_pending when slot_state is null can overwrite and destroy the active state of a subsequent, concurrently started drain run, potentially leading to permanent hangs. The reviewer suggests only resetting drain_worker_elected to release the stale election lock, leaving the other state variables untouched.

[coderabbitai]

Actionable comments posted: 1

🤖 Prompt for all review comments with AI agents

Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

Inline comments:
In
`@src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp`:
- Around line 423-433: The guard race is due to
SyncStartDrainState::pending_task being a non-atomic pointer; make pending_task
an std::atomic<PTO2TaskSlotState*> and use explicit release/acquire semantics:
in drain_worker_dispatch() replace the non-atomic clear with
pending_task.store(nullptr, std::memory_order_release) (and any place that sets
it use store(..., std::memory_order_release)), and in the elected-path (the
CAS-success branch that reads pending_task around the null-check and
dereference) perform an acquire load
(pending_task.load(std::memory_order_acquire)) before checking for nullptr and
before dereferencing; alternatively, if you prefer fences, insert a
std::atomic_thread_fence(std::memory_order_acquire) in the elected read path and
a release fence in the clearing path to establish the happens-before with
drain_worker_elected/sync_start_pending updates (symbols: SyncStartDrainState,
pending_task, drain_worker_dispatch, drain_worker_elected, sync_start_pending).

🪄 Autofix (Beta)

Fix all unresolved CodeRabbit comments on this PR:

• Push a commit to this branch (recommended)
• Create a new PR with the fixes

---

ℹ️ Review info

⚙️ Run configuration

Configuration used: Organization UI

Review profile: CHILL

Plan: Pro

Run ID: 2c6caae0-6328-44a4-9b96-16d3ea6ddfef

📥 Commits

Reviewing files that changed from the base of the PR and between 0610fa3 and afd0f25.

📒 Files selected for processing (2)

• src/a2a3/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp
• src/a5/runtime/tensormap_and_ringbuffer/runtime/scheduler/scheduler_completion.cpp