Skip to content

dxf: derive current keyspace when acquiring task runtime #69434

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
ti-chi-bot merged 2 commits into from
Jun 24, 2026
Merged

dxf: derive current keyspace when acquiring task runtime #69434

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
7f796e4
dxf: derive current keyspace when acquiring task runtime
D3Hunter Jun 23, 2026
29c0599
dxf: address runtime review comments
D3Hunter Jun 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 8 additions & 1 deletion pkg/dxf/framework/dxfutil/util.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -29,15 +29,17 @@ type sessionProvider interface {
}

// AcquireTaskRuntime returns a runtime view for the task keyspace and a release function.
// The sessionProvider must supply sessions whose store keyspace is the current node's keyspace;
// this is used to detect whether the task belongs to a different keyspace.
// Callers must call the release function when the returned runtime is no longer used.
func AcquireTaskRuntime(
sessionProvider sessionProvider,
currentKS string,
taskKS string,
holderID string,
) (sqlsvrapi.Runtime, func(), error) {
var taskRuntime sqlsvrapi.Runtime
if err := sessionProvider.WithNewSession(func(se sessionctx.Context) error {
currentKS := se.GetStore().GetKeyspace()
Comment thread
D3Hunter marked this conversation as resolved.
sqlServer := se.GetSQLServer()
if taskKS != currentKS {
var err2 error
Expand Down Expand Up @@ -83,3 +85,8 @@ func CheckTaskRuntime(runtime sqlsvrapi.Runtime, taskKS string) error {
}
return nil
}

// GenHolderID generates a holder ID for the given DXF component and task ID.
func GenHolderID(component string, taskID int64) string {
Comment thread
D3Hunter marked this conversation as resolved.
return fmt.Sprintf("DXF/%s/%d", component, taskID)
}
13 changes: 5 additions & 8 deletions pkg/dxf/framework/dxfutil/util_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,8 +74,9 @@ func (p *taskSessionProvider) WithNewSession(fn func(se sessionctx.Context) erro
return fn(p.se)
}

func newTaskSessionProvider(server *sqlsvrapimock.MockServer) *taskSessionProvider {
func newTaskSessionProvider(server *sqlsvrapimock.MockServer, currentKS string) *taskSessionProvider {
se := utilmock.NewContext()
se.Store = &storeWithKeyspace{keyspace: currentKS}
se.BindDomainAndSchValidator(server, nil)
return &taskSessionProvider{se: se}
}
Expand All @@ -91,8 +92,7 @@ func TestAcquireTaskRuntime(t *testing.T) {
server.EXPECT().GetRuntime().Return(runtime)

gotRuntime, releaseRuntime, err := AcquireTaskRuntime(
newTaskSessionProvider(server),
"task_ks",
newTaskSessionProvider(server, "task_ks"),
"task_ks",
"holder",
)
Expand All @@ -112,8 +112,7 @@ func TestAcquireTaskRuntime(t *testing.T) {
server.EXPECT().AcquireKSRuntime("task_ks", "holder").Return(runtimeHandle, nil)

gotRuntime, releaseRuntime, err := AcquireTaskRuntime(
newTaskSessionProvider(server),
"current_ks",
newTaskSessionProvider(server, "current_ks"),
"task_ks",
"holder",
)
Expand All @@ -132,8 +131,7 @@ func TestAcquireTaskRuntime(t *testing.T) {
server.EXPECT().AcquireKSRuntime("task_ks", "holder").Return(nil, runtimeErr)

gotRuntime, releaseRuntime, err := AcquireTaskRuntime(
newTaskSessionProvider(server),
"current_ks",
newTaskSessionProvider(server, "current_ks"),
"task_ks",
"holder",
)
Expand All @@ -147,7 +145,6 @@ func TestAcquireTaskRuntime(t *testing.T) {

gotRuntime, releaseRuntime, err := AcquireTaskRuntime(
&taskSessionProvider{err: sessionErr},
"current_ks",
"task_ks",
"holder",
)
Expand Down
1 change: 1 addition & 0 deletions pkg/dxf/framework/scheduler/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -74,6 +74,7 @@ go_test(
"//pkg/domain/sqlsvrapi",
"//pkg/domain/sqlsvrapi/mock",
"//pkg/dxf/framework/dxfmetric",
"//pkg/dxf/framework/dxfutil",
"//pkg/dxf/framework/handle",
"//pkg/dxf/framework/mock",
"//pkg/dxf/framework/proto",
Expand Down
5 changes: 2 additions & 3 deletions pkg/dxf/framework/scheduler/scheduler_manager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ package scheduler

import (
"context"
"fmt"
"slices"
"time"

Expand Down Expand Up @@ -352,8 +351,8 @@ func (sm *Manager) startScheduler(basicTask *proto.TaskBase, allocateSlots bool,
return
}

holderID := fmt.Sprintf("DXF/scheduler/%d", task.ID)
taskRuntime, releaseFn, err := dxfutil.AcquireTaskRuntime(sm.taskMgr, sm.store.GetKeyspace(), task.Keyspace, holderID)
holderID := dxfutil.GenHolderID("scheduler", task.ID)
taskRuntime, releaseFn, err := dxfutil.AcquireTaskRuntime(sm.taskMgr, task.Keyspace, holderID)
if err != nil {
sm.logger.Warn("acquire task runtime failed", zap.Int64("task-id", basicTask.ID),
zap.String("task-key", basicTask.Key), zap.Error(err))
Expand Down
13 changes: 9 additions & 4 deletions pkg/dxf/framework/scheduler/scheduler_manager_nokit_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ package scheduler

import (
"context"
"fmt"
"sync/atomic"
"testing"
"time"
Expand All @@ -25,6 +24,7 @@ import (
"github.com/pingcap/failpoint"
"github.com/pingcap/tidb/pkg/domain/sqlsvrapi"
sqlsvrapimock "github.com/pingcap/tidb/pkg/domain/sqlsvrapi/mock"
"github.com/pingcap/tidb/pkg/dxf/framework/dxfutil"
"github.com/pingcap/tidb/pkg/dxf/framework/mock"
"github.com/pingcap/tidb/pkg/dxf/framework/proto"
mockScheduler "github.com/pingcap/tidb/pkg/dxf/framework/scheduler/mock"
Expand Down Expand Up @@ -72,8 +72,11 @@ func expectRuntimeFromNewSession(ctrl *gomock.Controller, taskMgr *mock.MockTask
server := sqlsvrapimock.NewMockServer(ctrl)
server.EXPECT().GetRuntime().Return(runtime).AnyTimes()
taskMgr.EXPECT().WithNewSession(gomock.Any()).DoAndReturn(func(fn func(sessionctx.Context) error) error {
se := utilmock.NewContext()
// Match the mock session store to the runtime so AcquireTaskRuntime takes the local-runtime path.
se.Store = runtime.Store()

@ingress-bot ingress-bot Jun 24, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🧹 [Nit] se.Store assignment in expectRuntimeFromNewSession lacks a comment explaining its role in keyspace routing

The added se.Store = runtime.Store() line controls which branch of AcquireTaskRuntime executes: it makes the mock session's store carry the same keyspace as the runtime, so the same-keyspace path is taken and GetRuntime() is called. Without this context, the assignment reads as boilerplate rather than a contract-critical setup.

return fn(&sessionWithSQLServer{
Context: utilmock.NewContext(),
Context: se,
server: server,
})
}).AnyTimes()
Expand Down Expand Up @@ -265,7 +268,9 @@ func newCrossKeyspaceStartCase(t *testing.T, taskID int64, taskKey string) *cros
server := sqlsvrapimock.NewMockServer(ctrl)
taskMgr.EXPECT().GetTaskByID(gomock.Any(), task.ID).Return(task, nil)
taskMgr.EXPECT().WithNewSession(gomock.Any()).DoAndReturn(func(fn func(sessionctx.Context) error) error {
return fn(&sessionWithSQLServer{Context: utilmock.NewContext(), server: server})
se := utilmock.NewContext()
se.Store = manager.store
return fn(&sessionWithSQLServer{Context: se, server: server})
})

return &crossKeyspaceStartCase{
Expand All @@ -286,7 +291,7 @@ func (tc *crossKeyspaceStartCase) expectRuntimeAcquiredAndReleased() *sqlsvrapim
}

func (tc *crossKeyspaceStartCase) holderID() string {
return fmt.Sprintf("DXF/scheduler/%d", tc.task.ID)
return dxfutil.GenHolderID("scheduler", tc.task.ID)
}

func TestStartSchedulerCrossKeyspaceRuntime(t *testing.T) {
Expand Down
1 change: 1 addition & 0 deletions pkg/dxf/framework/taskexecutor/BUILD.bazel
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -61,6 +61,7 @@ go_test(
deps = [
"//pkg/domain/sqlsvrapi",
"//pkg/domain/sqlsvrapi/mock",
"//pkg/dxf/framework/dxfutil",
"//pkg/dxf/framework/mock",
"//pkg/dxf/framework/mock/execute",
"//pkg/dxf/framework/proto",
Expand Down
5 changes: 2 additions & 3 deletions pkg/dxf/framework/taskexecutor/manager.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,6 @@ package taskexecutor

import (
"context"
"fmt"
"sync"
"time"

Expand Down Expand Up @@ -331,8 +330,8 @@ func (m *Manager) startTaskExecutor(taskBase *proto.TaskBase) (executorStarted b
}
}()

holderID := fmt.Sprintf("DXF/executor/%d", task.ID)
taskRuntime, releaseFn, err := dxfutil.AcquireTaskRuntime(m.taskTable, m.store.GetKeyspace(), task.Keyspace, holderID)
holderID := dxfutil.GenHolderID("executor", task.ID)
taskRuntime, releaseFn, err := dxfutil.AcquireTaskRuntime(m.taskTable, task.Keyspace, holderID)
if err != nil {
m.logger.Warn("acquire task runtime failed", zap.Int64("task-id", taskBase.ID),
zap.String("task-key", taskBase.Key), zap.Error(err))
Expand Down
24 changes: 16 additions & 8 deletions pkg/dxf/framework/taskexecutor/manager_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,12 +17,12 @@ package taskexecutor
import (
"context"
"errors"
"fmt"
"testing"
"time"

"github.com/pingcap/tidb/pkg/domain/sqlsvrapi"
sqlsvrapimock "github.com/pingcap/tidb/pkg/domain/sqlsvrapi/mock"
"github.com/pingcap/tidb/pkg/dxf/framework/dxfutil"
"github.com/pingcap/tidb/pkg/dxf/framework/mock"
"github.com/pingcap/tidb/pkg/dxf/framework/proto"
"github.com/pingcap/tidb/pkg/dxf/framework/storage"
Expand Down Expand Up @@ -68,8 +68,10 @@ func expectRuntimeFromNewSession(ctrl *gomock.Controller, taskTable *mock.MockTa
server := sqlsvrapimock.NewMockServer(ctrl)
server.EXPECT().GetRuntime().Return(runtime).AnyTimes()
taskTable.EXPECT().WithNewSession(gomock.Any()).DoAndReturn(func(fn func(sessionctx.Context) error) error {
se := utilmock.NewContext()
se.Store = runtime.Store()
return fn(&sessionWithSQLServer{
Context: utilmock.NewContext(),
Context: se,
server: server,
})
}).AnyTimes()
Expand Down Expand Up @@ -224,7 +226,9 @@ func newCrossKeyspaceStartCase(t *testing.T, taskID int64, taskKey string) *cros
server := sqlsvrapimock.NewMockServer(ctrl)
taskTable.EXPECT().GetTaskByID(gomock.Any(), task.ID).Return(task, nil)
taskTable.EXPECT().WithNewSession(gomock.Any()).DoAndReturn(func(fn func(sessionctx.Context) error) error {
return fn(&sessionWithSQLServer{Context: utilmock.NewContext(), server: server})
se := utilmock.NewContext()
se.Store = m.store
return fn(&sessionWithSQLServer{Context: se, server: server})
})

return &crossKeyspaceStartCase{
Expand All @@ -245,7 +249,7 @@ func (tc *crossKeyspaceStartCase) expectRuntimeAcquiredAndReleased() *sqlsvrapim
}

func (tc *crossKeyspaceStartCase) holderID() string {
return fmt.Sprintf("DXF/executor/%d", tc.task.ID)
return dxfutil.GenHolderID("executor", tc.task.ID)
}

func TestStartTaskExecutorCrossKeyspaceRuntime(t *testing.T) {
Expand Down Expand Up @@ -675,7 +679,7 @@ func TestStartTaskExecutorResolveTaskRuntimeFromTaskKeyspace(t *testing.T) {
runtimeHandle := newRuntimeHandle(ctrl, taskStore)
runtimeHandle.EXPECT().Release()
server := sqlsvrapimock.NewMockServer(ctrl)
server.EXPECT().AcquireKSRuntime(taskKS, "DXF/executor/1").Return(runtimeHandle, nil)
server.EXPECT().AcquireKSRuntime(taskKS, dxfutil.GenHolderID("executor", task.ID)).Return(runtimeHandle, nil)

mockExecutor := mock.NewMockTaskExecutor(ctrl)
var gotStore kv.Storage
Expand All @@ -685,7 +689,9 @@ func TestStartTaskExecutorResolveTaskRuntimeFromTaskKeyspace(t *testing.T) {
})
mockTaskTable.EXPECT().GetTaskByID(gomock.Any(), task.ID).Return(task, nil)
mockTaskTable.EXPECT().WithNewSession(gomock.Any()).DoAndReturn(func(fn func(sessionctx.Context) error) error {
return fn(&sessionWithSQLServer{Context: utilmock.NewContext(), server: server})
se := utilmock.NewContext()
se.Store = m.store
return fn(&sessionWithSQLServer{Context: se, server: server})
})
mockExecutor.EXPECT().Init(gomock.Any()).Return(nil)
runCh := make(chan struct{})
Expand Down Expand Up @@ -727,7 +733,7 @@ func TestStartTaskExecutorResolveTaskRuntimeError(t *testing.T) {

runtimeErr := errors.New("ks runtime not found")
server := sqlsvrapimock.NewMockServer(ctrl)
server.EXPECT().AcquireKSRuntime(taskKS, "DXF/executor/2").Return(nil, runtimeErr)
server.EXPECT().AcquireKSRuntime(taskKS, dxfutil.GenHolderID("executor", task.ID)).Return(nil, runtimeErr)
factoryCalled := false
RegisterTaskType(task.Type, func(context.Context, *proto.Task, Param) TaskExecutor {
factoryCalled = true
Expand All @@ -736,7 +742,9 @@ func TestStartTaskExecutorResolveTaskRuntimeError(t *testing.T) {

mockTaskTable.EXPECT().GetTaskByID(gomock.Any(), task.ID).Return(task, nil)
mockTaskTable.EXPECT().WithNewSession(gomock.Any()).DoAndReturn(func(fn func(sessionctx.Context) error) error {
return fn(&sessionWithSQLServer{Context: utilmock.NewContext(), server: server})
se := utilmock.NewContext()
se.Store = m.store
return fn(&sessionWithSQLServer{Context: se, server: server})
})

require.False(t, m.startTaskExecutor(&task.TaskBase))
Expand Down
Loading