logpuller: extract region request scheduler from subscription client - #5665
logpuller: extract region request scheduler from subscription client#5665lidezhu wants to merge 18 commits into
Conversation
|
Skipping CI for Draft Pull Request. |
ti-chi-bot
Bot
added
do-not-merge/work-in-progress
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
📝 WalkthroughWalkthroughThe log puller extracts region request scheduling into dedicated scheduler, store, and worker components. It rewires subscription lifecycle and failure recovery around shared upstream dependencies, updates tests and mocks, and removes obsolete client-local request management. ChangesLog puller scheduling
Estimated code review effort: 4 (Complex) | ~60 minutes Sequence Diagram(s)sequenceDiagram
participant SubscriptionClient
participant RegionRequestScheduler
participant RegionRequestStore
participant RegionRequestWorker
participant RegionFailureHandler
SubscriptionClient->>RegionRequestScheduler: submit region or range task
RegionRequestScheduler->>RegionRequestStore: route task by store address
RegionRequestStore->>RegionRequestWorker: admit and send request
RegionRequestWorker->>SubscriptionClient: push region or timestamp event
RegionRequestWorker->>RegionFailureHandler: report stream or request failure
RegionFailureHandler->>RegionRequestScheduler: schedule recovery task
Possibly related PRs
Suggested reviewers: Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✨ Finishing Touches 💡 1📝 Generate docstrings 💡
🧪 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 |
ti-chi-bot
Bot
added
the
size/XL
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request refactors the region request scheduling logic in the logpuller package by introducing a dedicated regionRequestScheduler and requestedStore to manage workers connected to TiKV stores, decoupling this responsibility from subscriptionClient. The review feedback suggests adding defensive nil checks for worker.admission in the submit, close, and inflightCount methods of requestedStore to prevent potential nil pointer dereferences during testing or partial initialization.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
ti-chi-bot
Bot
added
size/XXL
f6cc71c to
4d3a6b4
Compare
33e46fc to
7e3df1d
Compare
|
/gemini review |
|
@coderabbitai review |
✅ Action performedReview finished.
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (2)
logservice/logpuller/mock_upstream_test.go (1)
71-76: 📐 Maintainability & Code Quality | 🔵 Trivial | ⚡ Quick winDo not call
requirefrom the serving goroutine.
require.Nilcallst.FailNow, which the testing package allows only from the goroutine that runs the test. From this goroutine the failure can be lost, and the test does not stop where expected.grpc.Server.Servealso returnsgrpc.ErrServerStoppedwhenStopruns beforeServestarts, so this assertion can fail during normal teardown.Log the error instead, or ignore
grpc.ErrServerStoppedand record the error for the test goroutine to assert.♻️ Proposed change
wg.Add(1) go func() { defer wg.Done() - err := grpcServer.Serve(lis) - require.Nil(t, err) + if err := grpcServer.Serve(lis); err != nil && !errors.Is(err, grpc.ErrServerStopped) { + log.Warn("mock grpc server exited with error", zap.Error(err)) + } }()Add the
errorsimport for this change.As per coding guidelines: "Prefer focused deterministic tests".
🤖 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 `@logservice/logpuller/mock_upstream_test.go` around lines 71 - 76, Update the grpcServer.Serve goroutine to remove require.Nil, since test failures must be reported from the test goroutine. Handle the expected grpc.ErrServerStopped teardown result, and log or record any other Serve error for the main test goroutine to assert; add the errors import if needed.Source: Coding guidelines
logservice/logpuller/region_request_scheduler_test.go (1)
93-122: 📐 Maintainability & Code Quality | 🔵 Trivial | 💤 Low valueExtract the shared scheduler fixture.
Both tests repeat the same setup: mock TiKV cluster,
mockPDClient, region cache, backoffer, raw span,subscribedSpan, and theLockRangeassertion. Extract one helper that returns the region cache, the located region, and the lockedsubscribedSpan. The tests then differ only in the stopped flag and the store wiring.Also applies to: 173-202
🤖 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 `@logservice/logpuller/region_request_scheduler_test.go` around lines 93 - 122, The repeated mock-cluster and locked-span setup in the affected scheduler tests should be extracted into a shared test helper. Add a helper that creates the mock TiKV cluster and mockPDClient, initializes the region cache and backoffer, builds the raw span and subscribedSpan, performs and asserts the successful LockRange call, and returns the region cache, located region, and locked subscribedSpan; update both tests to use it so only stopped-state and store-wiring setup remain test-specific.
🤖 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 `@logservice/logpuller/subscription_client_test.go`:
- Around line 426-441: Update the two t.Fatal failure messages in
TestRegionEventSinkPushUnblocksOnClientClose to reference regionEventSink.Push
instead of the removed pushRegionEventToDS method, while leaving the test
behavior unchanged.
---
Nitpick comments:
In `@logservice/logpuller/mock_upstream_test.go`:
- Around line 71-76: Update the grpcServer.Serve goroutine to remove
require.Nil, since test failures must be reported from the test goroutine.
Handle the expected grpc.ErrServerStopped teardown result, and log or record any
other Serve error for the main test goroutine to assert; add the errors import
if needed.
In `@logservice/logpuller/region_request_scheduler_test.go`:
- Around line 93-122: The repeated mock-cluster and locked-span setup in the
affected scheduler tests should be extracted into a shared test helper. Add a
helper that creates the mock TiKV cluster and mockPDClient, initializes the
region cache and backoffer, builds the raw span and subscribedSpan, performs and
asserts the successful LockRange call, and returns the region cache, located
region, and locked subscribedSpan; update both tests to use it so only
stopped-state and store-wiring setup remain test-specific.
🪄 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 Plus
Run ID: 4386cd10-de97-4565-bf5a-ff652d3b174a
📒 Files selected for processing (14)
logservice/logpuller/mock_upstream_test.gologservice/logpuller/region_failure_handler.gologservice/logpuller/region_failure_handler_test.gologservice/logpuller/region_request_scheduler.gologservice/logpuller/region_request_scheduler_test.gologservice/logpuller/region_request_store.gologservice/logpuller/region_request_store_test.gologservice/logpuller/region_request_worker.gologservice/logpuller/region_request_worker_test.gologservice/logpuller/scan_priority_test.gologservice/logpuller/subscription_client.gologservice/logpuller/subscription_client_test.gopkg/metrics/log_puller.goserver/server.go
| func TestRegionEventSinkPushUnblocksOnClientClose(t *testing.T) { | ||
| sink := ®ionEventSink{ | ||
| ds: &mockDynamicStream{}, | ||
| } | ||
| return rawSpan, span | ||
| } | ||
|
|
||
| func newScanPriorityTestRegion(span *subscribedSpan) regionInfo { | ||
| return newRegionInfo(tikv.NewRegionVerID(1, 1, 1), span.span, nil, span, false) | ||
| } | ||
|
|
||
| func newTestScanPriorityPolicy() scanPriorityPolicy { | ||
| return newScanPriorityPolicy(pdutil.NewClock4Test(), 30*time.Minute) | ||
| } | ||
|
|
||
| func TestPushRegionEventToDSUnblocksOnClose(t *testing.T) { | ||
| sink := newTestRegionEventSink(&mockDynamicStream{}) | ||
| client := &subscriptionClient{ | ||
| eventSink: sink, | ||
| regionTaskQueue: priorityqueue.New[*regionPriorityTask](), | ||
| sink.cond = sync.NewCond(&sink.mu) | ||
| client := &subscriptionClient{eventSink: sink} | ||
| client.regionScheduler = ®ionRequestScheduler{ | ||
| taskQueue: priorityqueue.New[*regionPriorityTask](), | ||
| } | ||
| client.ctx, client.cancel = context.WithCancel(context.Background()) | ||
|
|
||
| sink.paused.Store(true) | ||
|
|
||
| done := make(chan struct{}) | ||
| go func() { | ||
| client.pushRegionEventToDS(SubscriptionID(1), regionEvent{}) | ||
| sink.Push(SubscriptionID(1), regionEvent{}) |
There was a problem hiding this comment.
📐 Maintainability & Code Quality | 🟡 Minor | ⚡ Quick win
Update the stale failure messages in this test.
The test now exercises regionEventSink.Push. The two t.Fatal messages at Lines 447 and 456 still name the removed pushRegionEventToDS method. Rename them so a failure points to the current API.
💚 Proposed fix
select {
case <-done:
- t.Fatal("pushRegionEventToDS should block when paused")
+ t.Fatal("regionEventSink.Push should block when paused")
case <-time.After(100 * time.Millisecond):
}
require.NoError(t, client.Close(context.Background()))
select {
case <-done:
case <-time.After(5 * time.Second):
- t.Fatal("pushRegionEventToDS should be unblocked by Close")
+ t.Fatal("regionEventSink.Push should be unblocked by Close")
}📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| func TestRegionEventSinkPushUnblocksOnClientClose(t *testing.T) { | |
| sink := ®ionEventSink{ | |
| ds: &mockDynamicStream{}, | |
| } | |
| return rawSpan, span | |
| } | |
| func newScanPriorityTestRegion(span *subscribedSpan) regionInfo { | |
| return newRegionInfo(tikv.NewRegionVerID(1, 1, 1), span.span, nil, span, false) | |
| } | |
| func newTestScanPriorityPolicy() scanPriorityPolicy { | |
| return newScanPriorityPolicy(pdutil.NewClock4Test(), 30*time.Minute) | |
| } | |
| func TestPushRegionEventToDSUnblocksOnClose(t *testing.T) { | |
| sink := newTestRegionEventSink(&mockDynamicStream{}) | |
| client := &subscriptionClient{ | |
| eventSink: sink, | |
| regionTaskQueue: priorityqueue.New[*regionPriorityTask](), | |
| sink.cond = sync.NewCond(&sink.mu) | |
| client := &subscriptionClient{eventSink: sink} | |
| client.regionScheduler = ®ionRequestScheduler{ | |
| taskQueue: priorityqueue.New[*regionPriorityTask](), | |
| } | |
| client.ctx, client.cancel = context.WithCancel(context.Background()) | |
| sink.paused.Store(true) | |
| done := make(chan struct{}) | |
| go func() { | |
| client.pushRegionEventToDS(SubscriptionID(1), regionEvent{}) | |
| sink.Push(SubscriptionID(1), regionEvent{}) | |
| select { | |
| case <-done: | |
| t.Fatal("regionEventSink.Push should block when paused") | |
| case <-time.After(100 * time.Millisecond): | |
| } | |
| require.NoError(t, client.Close(context.Background())) | |
| select { | |
| case <-done: | |
| case <-time.After(5 * time.Second): | |
| t.Fatal("regionEventSink.Push should be unblocked by Close") | |
| } |
🤖 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 `@logservice/logpuller/subscription_client_test.go` around lines 426 - 441,
Update the two t.Fatal failure messages in
TestRegionEventSinkPushUnblocksOnClientClose to reference regionEventSink.Push
instead of the removed pushRegionEventToDS method, while leaving the test
behavior unchanged.
1 participant
What problem does this PR solve?
Issue Number: close #5871
What is changed and how it works?
Check List
Tests
Questions
Will it cause performance regression or break compatibility?
Do you need to update user documentation, design documentation or monitoring documentation?
Release note
Summary by CodeRabbit
Performance and Reliability
Tests