diff --git a/Makefile b/Makefile
index 6ce62a320ebc4..96c92efe3295e 100644
--- a/Makefile
+++ b/Makefile
@@ -272,7 +272,7 @@ tools/bin/gotestsum:
 
 # mockgen@v0.2.0 is imcompatible with v0.3.0, so install it always.
 mockgen:
-	GOBIN=$(shell pwd)/tools/bin $(GO) install github.com/lance6716/mock/mockgen@proposal_embed_option
+	GOBIN=$(shell pwd)/tools/bin $(GO) install github.com/lance6716/mock/mockgen@v0.4.0-patch
 
 # Usage:
 #
@@ -392,7 +392,7 @@ mock_lightning: mockgen
 gen_mock: mockgen
 	tools/bin/mockgen -package mock github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor TaskTable,Pool,TaskExecutor,Extension > pkg/disttask/framework/mock/task_executor_mock.go
 	tools/bin/mockgen -package mock github.com/pingcap/tidb/pkg/disttask/framework/scheduler Scheduler,CleanUpRoutine,TaskManager > pkg/disttask/framework/mock/scheduler_mock.go
-	tools/bin/mockgen -package mock github.com/pingcap/tidb/pkg/disttask/framework/scheduler Extension > pkg/disttask/framework/scheduler/mock/scheduler_mock.go
+	tools/bin/mockgen -destination pkg/disttask/framework/scheduler/mock/scheduler_mock.go -package mock github.com/pingcap/tidb/pkg/disttask/framework/scheduler Extension
 	tools/bin/mockgen -embed -package mockexecute github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/execute StepExecutor > pkg/disttask/framework/mock/execute/execute_mock.go
 	tools/bin/mockgen -package mock github.com/pingcap/tidb/pkg/disttask/importinto MiniTaskExecutor > pkg/disttask/importinto/mock/import_mock.go
 	tools/bin/mockgen -package mock github.com/pingcap/tidb/pkg/disttask/framework/planner LogicalPlan,PipelineSpec > pkg/disttask/framework/mock/plan_mock.go
diff --git a/br/pkg/mock/storage/storage.go b/br/pkg/mock/storage/storage.go
index 6ca014bb92b93..5e547f4c18436 100644
--- a/br/pkg/mock/storage/storage.go
+++ b/br/pkg/mock/storage/storage.go
@@ -5,6 +5,7 @@
 //
 //	mockgen -package mockstorage github.com/pingcap/tidb/br/pkg/storage ExternalStorage
 //
+
 // Package mockstorage is a generated GoMock package.
 package mockstorage
 
@@ -39,6 +40,11 @@ func (m *MockExternalStorage) EXPECT() *MockExternalStorageMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockExternalStorage) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // Close mocks base method.
 func (m *MockExternalStorage) Close() {
 	m.ctrl.T.Helper()
diff --git a/pkg/disttask/framework/mock/execute/execute_mock.go b/pkg/disttask/framework/mock/execute/execute_mock.go
index d8be965a479a9..edd5af519dad1 100644
--- a/pkg/disttask/framework/mock/execute/execute_mock.go
+++ b/pkg/disttask/framework/mock/execute/execute_mock.go
@@ -3,11 +3,11 @@
 //
 // Generated by this command:
 //
-//	mockgen -embed -package execute github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/execute StepExecutor
+//	mockgen -embed -package mockexecute github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/execute StepExecutor
 //
 
-// Package execute is a generated GoMock package.
-package execute
+// Package mockexecute is a generated GoMock package.
+package mockexecute
 
 import (
 	context "context"
diff --git a/pkg/disttask/framework/mock/plan_mock.go b/pkg/disttask/framework/mock/plan_mock.go
index 1192073675f11..b0db8b51dc48d 100644
--- a/pkg/disttask/framework/mock/plan_mock.go
+++ b/pkg/disttask/framework/mock/plan_mock.go
@@ -5,6 +5,7 @@
 //
 //	mockgen -package mock github.com/pingcap/tidb/pkg/disttask/framework/planner LogicalPlan,PipelineSpec
 //
+
 // Package mock is a generated GoMock package.
 package mock
 
@@ -38,6 +39,11 @@ func (m *MockLogicalPlan) EXPECT() *MockLogicalPlanMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockLogicalPlan) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // FromTaskMeta mocks base method.
 func (m *MockLogicalPlan) FromTaskMeta(arg0 []byte) error {
 	m.ctrl.T.Helper()
@@ -105,6 +111,11 @@ func (m *MockPipelineSpec) EXPECT() *MockPipelineSpecMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockPipelineSpec) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // ToSubtaskMeta mocks base method.
 func (m *MockPipelineSpec) ToSubtaskMeta(arg0 planner.PlanCtx) ([]byte, error) {
 	m.ctrl.T.Helper()
diff --git a/pkg/disttask/framework/mock/scheduler_mock.go b/pkg/disttask/framework/mock/scheduler_mock.go
index 5493ab1d2f556..a37a16b894be7 100644
--- a/pkg/disttask/framework/mock/scheduler_mock.go
+++ b/pkg/disttask/framework/mock/scheduler_mock.go
@@ -5,6 +5,7 @@
 //
 //	mockgen -package mock github.com/pingcap/tidb/pkg/disttask/framework/scheduler Scheduler,CleanUpRoutine,TaskManager
 //
+
 // Package mock is a generated GoMock package.
 package mock
 
@@ -41,6 +42,11 @@ func (m *MockScheduler) EXPECT() *MockSchedulerMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockScheduler) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // Close mocks base method.
 func (m *MockScheduler) Close() {
 	m.ctrl.T.Helper()
@@ -200,6 +206,11 @@ func (m *MockCleanUpRoutine) EXPECT() *MockCleanUpRoutineMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockCleanUpRoutine) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // CleanUp mocks base method.
 func (m *MockCleanUpRoutine) CleanUp(arg0 context.Context, arg1 *proto.Task) error {
 	m.ctrl.T.Helper()
@@ -237,6 +248,11 @@ func (m *MockTaskManager) EXPECT() *MockTaskManagerMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockTaskManager) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // CancelTask mocks base method.
 func (m *MockTaskManager) CancelTask(arg0 context.Context, arg1 int64) error {
 	m.ctrl.T.Helper()
diff --git a/pkg/disttask/framework/mock/task_executor_mock.go b/pkg/disttask/framework/mock/task_executor_mock.go
index 67a57b574cc28..2ecd10ab0472a 100644
--- a/pkg/disttask/framework/mock/task_executor_mock.go
+++ b/pkg/disttask/framework/mock/task_executor_mock.go
@@ -5,6 +5,7 @@
 //
 //	mockgen -package mock github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor TaskTable,Pool,TaskExecutor,Extension
 //
+
 // Package mock is a generated GoMock package.
 package mock
 
@@ -41,6 +42,11 @@ func (m *MockTaskTable) EXPECT() *MockTaskTableMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockTaskTable) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // CancelSubtask mocks base method.
 func (m *MockTaskTable) CancelSubtask(arg0 context.Context, arg1 string, arg2 int64) error {
 	m.ctrl.T.Helper()
@@ -315,6 +321,11 @@ func (m *MockPool) EXPECT() *MockPoolMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockPool) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // ReleaseAndWait mocks base method.
 func (m *MockPool) ReleaseAndWait() {
 	m.ctrl.T.Helper()
@@ -378,6 +389,11 @@ func (m *MockTaskExecutor) EXPECT() *MockTaskExecutorMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockTaskExecutor) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // Cancel mocks base method.
 func (m *MockTaskExecutor) Cancel() {
 	m.ctrl.T.Helper()
@@ -491,6 +507,11 @@ func (m *MockExtension) EXPECT() *MockExtensionMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockExtension) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // GetStepExecutor mocks base method.
 func (m *MockExtension) GetStepExecutor(arg0 *proto.Task) (execute.StepExecutor, error) {
 	m.ctrl.T.Helper()
diff --git a/pkg/disttask/framework/scheduler/mock/scheduler_mock.go b/pkg/disttask/framework/scheduler/mock/scheduler_mock.go
index 54e18806e8a62..e820c22016f16 100644
--- a/pkg/disttask/framework/scheduler/mock/scheduler_mock.go
+++ b/pkg/disttask/framework/scheduler/mock/scheduler_mock.go
@@ -3,8 +3,9 @@
 //
 // Generated by this command:
 //
-//	mockgen -package mock github.com/pingcap/tidb/pkg/disttask/framework/scheduler Extension
+//	mockgen -destination pkg/disttask/framework/scheduler/mock/scheduler_mock.go -package mock github.com/pingcap/tidb/pkg/disttask/framework/scheduler Extension
 //
+
 // Package mock is a generated GoMock package.
 package mock
 
@@ -40,6 +41,11 @@ func (m *MockExtension) EXPECT() *MockExtensionMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockExtension) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // GetEligibleInstances mocks base method.
 func (m *MockExtension) GetEligibleInstances(arg0 context.Context, arg1 *proto.Task) ([]string, error) {
 	m.ctrl.T.Helper()
diff --git a/pkg/disttask/framework/taskexecutor/task_executor_test.go b/pkg/disttask/framework/taskexecutor/task_executor_test.go
index f469880fe1bb3..7db71b127048d 100644
--- a/pkg/disttask/framework/taskexecutor/task_executor_test.go
+++ b/pkg/disttask/framework/taskexecutor/task_executor_test.go
@@ -21,7 +21,7 @@ import (
 
 	"github.com/pingcap/errors"
 	"github.com/pingcap/tidb/pkg/disttask/framework/mock"
-	mockexecute "github.com/pingcap/tidb/pkg/disttask/framework/mock/execute"
+	"github.com/pingcap/tidb/pkg/disttask/framework/mock/execute"
 	"github.com/pingcap/tidb/pkg/disttask/framework/proto"
 	"github.com/pingcap/tidb/pkg/disttask/framework/storage"
 	"github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/execute"
diff --git a/pkg/disttask/importinto/mock/import_mock.go b/pkg/disttask/importinto/mock/import_mock.go
index 2aa27848b5bb5..13135320fbd38 100644
--- a/pkg/disttask/importinto/mock/import_mock.go
+++ b/pkg/disttask/importinto/mock/import_mock.go
@@ -5,6 +5,7 @@
 //
 //	mockgen -package mock github.com/pingcap/tidb/pkg/disttask/importinto MiniTaskExecutor
 //
+
 // Package mock is a generated GoMock package.
 package mock
 
@@ -39,6 +40,11 @@ func (m *MockMiniTaskExecutor) EXPECT() *MockMiniTaskExecutorMockRecorder {
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockMiniTaskExecutor) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // Run mocks base method.
 func (m *MockMiniTaskExecutor) Run(arg0 context.Context, arg1, arg2 backend.EngineWriter) error {
 	m.ctrl.T.Helper()
diff --git a/pkg/util/sqlexec/mock/restricted_sql_executor_mock.go b/pkg/util/sqlexec/mock/restricted_sql_executor_mock.go
index b0da3b6039828..7db258957e7ab 100644
--- a/pkg/util/sqlexec/mock/restricted_sql_executor_mock.go
+++ b/pkg/util/sqlexec/mock/restricted_sql_executor_mock.go
@@ -5,6 +5,7 @@
 //
 //	mockgen -package mock github.com/pingcap/tidb/pkg/util/sqlexec RestrictedSQLExecutor
 //
+
 // Package mock is a generated GoMock package.
 package mock
 
@@ -41,6 +42,11 @@ func (m *MockRestrictedSQLExecutor) EXPECT() *MockRestrictedSQLExecutorMockRecor
 	return m.recorder
 }
 
+// ISGOMOCK indicates that this struct is a gomock mock.
+func (m *MockRestrictedSQLExecutor) ISGOMOCK() struct{} {
+	return struct{}{}
+}
+
 // ExecRestrictedSQL mocks base method.
 func (m *MockRestrictedSQLExecutor) ExecRestrictedSQL(arg0 context.Context, arg1 []func(*sqlexec.ExecOption), arg2 string, arg3 ...any) ([]chunk.Row, []*ast.ResultField, error) {
 	m.ctrl.T.Helper()