Repository navigation
dxf: add StepPrepared prepare-mode transition in scheduler framework - #68581
Conversation
ti-chi-bot
Bot
added
the
release-note-none
|
@D3Hunter I've received your pull request and will start the review. I'll conduct a thorough review covering code quality, potential issues, and implementation details. ⏳ This process typically takes 10-30 minutes depending on the complexity of the changes. ℹ️ Learn more details on Pantheon AI. |
ti-chi-bot
Bot
added
the
size/XL
|
Hi @D3Hunter. Thanks for your PR. PRs from untrusted users cannot be marked as trusted with I understand the commands that are listed here. DetailsInstructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
|
Note Reviews pausedIt looks like this branch is under active development. To avoid overwhelming you with review comments due to an influx of new commits, CodeRabbit has automatically paused this review. You can configure this behavior by changing the Use the following commands to manage reviews:
Use the checkboxes below for quick actions:
📝 WalkthroughWalkthroughAdds PrepareMode and StepPrepared, Extension.OnPrepare and TaskManager.SwitchTaskStepAfterPrepare, scheduler changes to run OnPrepare and atomically persist StepPrepared and prepared metadata, storage change to merge extra_params, and tests/mocks validating prepare-mode and runtime-slot behavior. ChangesPrepare-Mode Task Lifecycle
Sequence Diagram(s): sequenceDiagram
participant Scheduler
participant Extension
participant TaskManager
participant Storage
Scheduler->>Extension: OnPrepare(ctx, handle, task) -- compute Meta/RequiredSlots/MaxNodeCount
Extension-->>Scheduler: nil / error
Scheduler->>TaskManager: SwitchTaskStepAfterPrepare(ctx, preparedTask)
TaskManager->>Storage: UPDATE mysql.tidb_global_task (CAS Init→Prepared) with meta/concurrency/max_node_count
Storage-->>TaskManager: (switched bool, err)
TaskManager-->>Scheduler: (switched bool, err)
alt switched == true
Scheduler->>Scheduler: set task.Step = StepPrepared and continue to GetNextStep
end
Estimated code review effort🎯 4 (Complex) | ⏱️ ~45 minutes Suggested labels
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
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. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 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 `@pkg/dxf/framework/scheduler/scheduler_nokit_test.go`:
- Around line 521-523: The local variable taskPreparedByOtherOwner is assigned
but never used in scheduler_nokit_test.go, causing a compile error; remove the
declaration and its assignments (the three lines creating
taskPreparedByOtherOwner, setting Step and Meta) or, if the intent was to assert
behavior for a prepared task, actually use taskPreparedByOtherOwner in the test
(e.g., pass it to the function under test or add assertions) so it is
referenced. Ensure no other code depends on that variable before removal.
In `@pkg/dxf/framework/storage/task_table.go`:
- Around line 875-883: The CAS update that moves a task to proto.StepPrepared
must also persist extra_params so OnPrepare edits aren't lost; update the SQL
that updates mysql.tidb_global_task to include extra_params = %? in the SET list
and add task.ExtraParams to the parameter list (alongside proto.StepPrepared,
task.Meta, task.RequiredSlots, task.MaxNodeCount) before the task.ID/expected
state args used in the where clause so the prepare-step write saves ExtraParams
as well.
🪄 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: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: bc18dfb9-ee54-4ae1-921e-a8a0ac4fd19e
📒 Files selected for processing (12)
pkg/dxf/framework/integrationtests/framework_test.gopkg/dxf/framework/mock/scheduler_mock.gopkg/dxf/framework/proto/step.gopkg/dxf/framework/proto/step_test.gopkg/dxf/framework/proto/task.gopkg/dxf/framework/proto/task_test.gopkg/dxf/framework/scheduler/interface.gopkg/dxf/framework/scheduler/mock/scheduler_mock.gopkg/dxf/framework/scheduler/scheduler.gopkg/dxf/framework/scheduler/scheduler_nokit_test.gopkg/dxf/framework/storage/table_test.gopkg/dxf/framework/storage/task_table.go
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## master #68581 +/- ##
================================================
- Coverage 76.3752% 75.3136% -1.0617%
================================================
Files 2039 2021 -18
Lines 562884 566944 +4060
================================================
- Hits 429904 426986 -2918
- Misses 132062 139956 +7894
+ Partials 918 2 -916
Flags with carried forward coverage won't be shown. Click here to find out more.
🚀 New features to boost your workflow:
|
|
🔍 Starting code review for this PR... |
ingress-bot
left a comment
ingress-bot
left a comment
There was a problem hiding this comment.
This review was generated by AI and should be verified by a human reviewer.
Manual follow-up is recommended before merge.
Summary
- Total findings: 13
- Inline comments: 12
- Summary-only findings (no inline anchor): 0
Findings (highest risk first)
⚠️ [Major] (8)
FrameworkStep2BusinessStepname and doc overstate what the helper actually returns (pkg/dxf/framework/proto/step.go:199, pkg/dxf/framework/proto/step_test.go:60)- Scheduler-internal step translation is exported from the shared proto package (pkg/dxf/framework/proto/step.go:203, pkg/dxf/framework/scheduler/scheduler.go:499)
- New prepare-mode flag is dropped by max-runtime-slots updates (pkg/dxf/framework/proto/task.go:116, pkg/server/handler/tikvhandler/dxf.go:432, pkg/dxf/framework/storage/task_table.go:1187)
- OnPrepare contract omits reentry/idempotency semantics for CAS race path (pkg/dxf/framework/scheduler/interface.go:156, pkg/dxf/framework/scheduler/scheduler.go:393, pkg/dxf/framework/storage/task_table.go:866)
- Prepare-phase mutation contract is comment-only and split across interface and SQL (pkg/dxf/framework/scheduler/interface.go:163, pkg/dxf/framework/storage/task_table.go:875, pkg/dxf/framework/scheduler/scheduler.go:405)
- Non-retryable OnPrepare failures are retried forever instead of triggering task failure handling (pkg/dxf/framework/scheduler/scheduler.go:394, pkg/dxf/framework/scheduler/scheduler.go:285)
- Prepare-mode retries re-execute
OnPrepareand break replay parity (pkg/dxf/framework/scheduler/scheduler.go:389, pkg/dxf/framework/scheduler/scheduler.go:285, pkg/dxf/framework/storage/task_table.go:868) - Persisting
StepPreparedintroduces an unguarded downgrade incompatibility (pkg/dxf/framework/storage/task_table.go:876, pkg/dxf/framework/scheduler/scheduler.go:499, pkg/dxf/example/scheduler.go:94, pkg/dxf/importinto/scheduler.go:517)
🟡 [Minor] (5)
- Step lifecycle comments are stale after introducing StepPrepared (pkg/dxf/framework/proto/step.go:27, pkg/dxf/framework/storage/task_table.go:866, pkg/dxf/framework/scheduler/scheduler.go:499)
- Framework-only StepPrepared is now accepted as an external target step (pkg/dxf/framework/proto/step.go:48, pkg/dxf/framework/proto/step.go:63, pkg/server/handler/tikvhandler/dxf.go:426)
SwitchTaskStepAfterPreparecontract text is broader than what persistence actually does (pkg/dxf/framework/scheduler/interface.go:86, pkg/dxf/framework/storage/task_table.go:875)- Prepare success path adds a fixed extra scheduler interval before subtask planning (pkg/dxf/framework/scheduler/scheduler.go:406, pkg/dxf/framework/scheduler/scheduler.go:55)
- Prepare transition duplicates task-state SQL transition logic instead of reusing existing transition helper (pkg/dxf/framework/storage/task_table.go:868, pkg/dxf/framework/storage/task_table.go:893)
|
|
||
| // SwitchTaskStepAfterPrepare atomically persists prepare completion from | ||
| // pending+init to pending+prepared. | ||
| func (mgr *TaskManager) SwitchTaskStepAfterPrepare(ctx context.Context, task *proto.Task) (bool, error) { |
There was a problem hiding this comment.
🟡 [Minor] Prepare transition duplicates task-state SQL transition logic instead of reusing existing transition helper
Impact
This PR introduces a second hand-written CAS update path for mysql.tidb_global_task transitions. Keeping transition SQL in parallel implementations increases drift pressure on transition invariants and forces future transition changes to be duplicated across two code paths.
Scope
pkg/dxf/framework/storage/task_table.go:868—TaskManager.SwitchTaskStepAfterPreparepkg/dxf/framework/storage/task_table.go:893—updateTaskStateStep
Evidence
SwitchTaskStepAfterPrepare adds a separate update mysql.tidb_global_task ... where id/state/step statement, while updateTaskStateStep already owns the core task state-step CAS transition update. Both functions now maintain overlapping transition responsibilities with different hard-coded field sets.
Change request
Refactor transition persistence so the prepare path reuses a shared transition primitive with explicit per-transition field options, rather than introducing another standalone SQL transition body. Keep one authoritative transition helper for CAS predicates and shared update semantics.
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
pkg/dxf/framework/scheduler/scheduler.go (1)
398-406:⚠️ Potential issue | 🟠 Major | ⚡ Quick winReload the task after a successful prepare CAS.
This path only patches
task.Steplocally. IfSwitchTaskStepAfterPreparealso normalizes other persisted fields during the transition, this scheduler keeps a stale in-memory task becauserefreshTaskIfNeeded()only refetches whenstate/stepdiverge. The next tick can then plan from outdated prepared data.Suggested fix
if !switched { return nil } - task.Step = proto.StepPrepared - s.task.Store(task) + updatedTask, err := s.taskMgr.GetTaskByID(s.ctx, task.ID) + if err != nil { + return errors.Trace(err) + } + s.task.Store(updatedTask) return nil🤖 Prompt for 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. In `@pkg/dxf/framework/scheduler/scheduler.go` around lines 398 - 406, After a successful CAS in SwitchTaskStepAfterPrepare, the code only mutates task.Step in-memory which can leave other normalized/persisted fields stale; update the scheduler to reload the authoritative task state and replace the in-memory task after the switch succeeds. Concretely: in the block that calls s.taskMgr.SwitchTaskStepAfterPrepare(s.ctx, task), when switched == true, re-fetch the task from the task manager or call the existing refreshTaskIfNeeded/refresh logic (or a direct Read/Get method on taskMgr) and then call s.task.Store(freshTask) instead of only setting task.Step so the in-memory copy matches any other persisted changes.
🤖 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.
Outside diff comments:
In `@pkg/dxf/framework/scheduler/scheduler.go`:
- Around line 398-406: After a successful CAS in SwitchTaskStepAfterPrepare, the
code only mutates task.Step in-memory which can leave other normalized/persisted
fields stale; update the scheduler to reload the authoritative task state and
replace the in-memory task after the switch succeeds. Concretely: in the block
that calls s.taskMgr.SwitchTaskStepAfterPrepare(s.ctx, task), when switched ==
true, re-fetch the task from the task manager or call the existing
refreshTaskIfNeeded/refresh logic (or a direct Read/Get method on taskMgr) and
then call s.task.Store(freshTask) instead of only setting task.Step so the
in-memory copy matches any other persisted changes.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 83ca7229-ed2a-4079-a65c-e2994b5ef455
📒 Files selected for processing (1)
pkg/dxf/framework/scheduler/scheduler.go
ti-chi-bot
Bot
added
size/XXL
There was a problem hiding this comment.
🧹 Nitpick comments (1)
pkg/dxf/framework/scheduler/scheduler_nokit_test.go (1)
543-565: 💤 Low valueTest verifies in-memory task state unchanged when CAS returns false, but the mutated fields from OnPrepare are lost.
When
SwitchTaskStepAfterPreparereturnsfalse, the test correctly asserts thatscheduler.GetTask()equalstaskWithPrepare(the original state before OnPrepare mutations). However, this means the prepared field updates (Meta,RequiredSlots,MaxNodeCount) computed byOnPrepareat lines 548-553 are discarded silently.This is the expected behavior per the current implementation (another owner won the CAS), but it's worth adding a brief comment in the test clarifying that mutations are intentionally not persisted when
switched=falsesince another scheduler instance already transitioned the task.🤖 Prompt for 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. In `@pkg/dxf/framework/scheduler/scheduler_nokit_test.go` around lines 543 - 565, Add a short clarifying comment in the test near the OnPrepare / SwitchTaskStepAfterPrepare expectation block (around the calls to schExt.OnPrepare, taskMgr.SwitchTaskStepAfterPrepare, and the assertion of scheduler.GetTask()) explaining that when SwitchTaskStepAfterPrepare returns false (CAS lost / another owner advanced the task), any mutations performed by OnPrepare (Meta, RequiredSlots, MaxNodeCount) are intentionally not persisted to the in-memory scheduler task; this comment should reference the functions scheduler.onPending, schExt.OnPrepare, taskMgr.SwitchTaskStepAfterPrepare and scheduler.GetTask so future readers understand the behavior being asserted.
🤖 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.
Nitpick comments:
In `@pkg/dxf/framework/scheduler/scheduler_nokit_test.go`:
- Around line 543-565: Add a short clarifying comment in the test near the
OnPrepare / SwitchTaskStepAfterPrepare expectation block (around the calls to
schExt.OnPrepare, taskMgr.SwitchTaskStepAfterPrepare, and the assertion of
scheduler.GetTask()) explaining that when SwitchTaskStepAfterPrepare returns
false (CAS lost / another owner advanced the task), any mutations performed by
OnPrepare (Meta, RequiredSlots, MaxNodeCount) are intentionally not persisted to
the in-memory scheduler task; this comment should reference the functions
scheduler.onPending, schExt.OnPrepare, taskMgr.SwitchTaskStepAfterPrepare and
scheduler.GetTask so future readers understand the behavior being asserted.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Repository UI
Review profile: CHILL
Plan: Pro
Run ID: 2584b6b8-c8f1-439b-a4e6-087ac79dff17
📒 Files selected for processing (8)
pkg/dxf/framework/proto/step.gopkg/dxf/framework/proto/step_test.gopkg/dxf/framework/scheduler/interface.gopkg/dxf/framework/scheduler/scheduler.gopkg/dxf/framework/scheduler/scheduler_nokit_test.gopkg/dxf/framework/storage/task_state_test.gopkg/dxf/framework/storage/task_table.gopkg/server/handler/tikvhandler/dxf.go
✅ Files skipped from review due to trivial changes (1)
- pkg/server/handler/tikvhandler/dxf.go
D3Hunter
left a comment
D3Hunter
left a comment
There was a problem hiding this comment.
Summary
- Total findings: 5
- Inline comments: 5
- Summary-only findings (no inline anchor): 0
Findings (highest risk first)
⚠️ [Major] (2)
- Prepare-mode can silently skip business steps for schedulers that do not handle
StepPrepared(pkg/dxf/framework/scheduler/scheduler.go:394; pkg/dxf/framework/storage/task_table.go:875; pkg/dxf/importinto/scheduler.go:534; pkg/ddl/backfilling_dist_scheduler.go:238) target_stepscannot be cleared after switchingextra_paramsupdates to JSON merge patch (pkg/dxf/framework/storage/task_table.go:1198; pkg/server/handler/tikvhandler/dxf.go:390; pkg/dxf/framework/proto/task.go:109)
🟡 [Minor] (2)
handlePlanErrname understates scope after prepare-mode integration (pkg/dxf/framework/scheduler/scheduler.go:395; pkg/dxf/framework/scheduler/scheduler.go:628)- Integration test entrypoint mixes runtime-slot behavior with prepare-mode lifecycle coverage (
pkg/dxf/framework/integrationtests/framework_test.go:344)
ℹ️ [Info] (1)
- Temporary rollout comments lack explicit tracking or removal criteria (
pkg/dxf/framework/scheduler/scheduler.go:155; pkg/dxf/framework/scheduler/scheduler.go:409)
|
/retest |
|
@D3Hunter: PRs from untrusted users cannot be marked as trusted with DetailsIn response to this:
Instructions for interacting with me using PR comments are available here. If you have questions or suggestions related to my behavior, please file an issue against the kubernetes-sigs/prow repository. |
|
[APPROVALNOTIFIER] This PR is APPROVED This pull-request has been approved by: GMHDBJD, joechenrh The full list of commands accepted by this bot can be found here. The pull request process is described here DetailsNeeds approval from an approver in each of these files:
Approvers can indicate their approval by writing |
ti-chi-bot
Bot
added
approved
lgtm
and removed
needs-1-more-lgtm
4 participants
What problem does this PR solve?
Issue Number: ref #68580
Problem Summary:
IMPORT INTOnextgen prepare flow needs a framework-level marker-step transition (init -> prepared) so prepare work can be completed and persisted before normal business-step scheduling.What changed and how does it work?
StepPreparedand map it back to businessStepInitin next-step resolution.prepare_mode(disabled/required) toproto.ExtraParamswith backward-compatible JSON behavior.prepare_mode=requiredand step isinit, callOnPrepare.preparedon successful CAS transition.SwitchTaskStepAfterPreparestorage API to persist prepare outputs atomically forpending + init -> pending + prepared.StepPreparedmapping, and prepare transition persistence.Check List
Tests
Unit test commands:
Side effects
Documentation
Release note
Summary by CodeRabbit
New Features
Behavior Changes
Tests