disttask: add StepExecFrameworkInfo as an example of Impl access Base

Makefile

@@ -272,7 +272,7 @@ tools/bin/gotestsum:
 
 # [email protected] is imcompatible with v0.3.0, so install it always.
 mockgen:
-	GOBIN=$(shell pwd)/tools/bin $(GO) install go.uber.org/mock/mockgen@v0.3.0
+	GOBIN=$(shell pwd)/tools/bin $(GO) install github.com/lance6716/mock/mockgen@v0.4.0-patch
 
 # Usage:
 #
@@ -392,8 +392,8 @@ 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 -package execute github.com/pingcap/tidb/pkg/disttask/framework/taskexecutor/execute StepExecutor > pkg/disttask/framework/mock/execute/execute_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
 	tools/bin/mockgen -package mock github.com/pingcap/tidb/pkg/util/sqlexec RestrictedSQLExecutor > pkg/util/sqlexec/mock/restricted_sql_executor_mock.go

pkg/ddl/backfilling_dist_executor.go

@@ -191,7 +191,7 @@ func (s *backfillDistExecutor) Init(ctx context.Context) error {
 	return nil
 }
 
-func (s *backfillDistExecutor) GetStepExecutor(task *proto.Task, _ *proto.StepResource) (execute.StepExecutor, error) {
+func (s *backfillDistExecutor) GetStepExecutor(task *proto.Task) (execute.StepExecutor, error) {
 	switch task.Step {
 	case proto.BackfillStepReadIndex, proto.BackfillStepMergeSort, proto.BackfillStepWriteAndIngest:
 		return s.newBackfillSubtaskExecutor(task.Step)

pkg/ddl/backfilling_read_index.go

@@ -38,6 +38,7 @@ import (
 )
 
 type readIndexExecutor struct {
+	execute.StepExecFrameworkInfo
 	d       *ddl
 	job     *model.Job
 	indexes []*model.IndexInfo

pkg/disttask/framework/taskexecutor/BUILD.bazel

@@ -51,13 +51,14 @@ go_test(
     ],
     embed = [":taskexecutor"],
     flaky = True,
-    shard_count = 15,
+    shard_count = 16,
     deps = [
         "//pkg/disttask/framework/mock",
         "//pkg/disttask/framework/mock/execute",
         "//pkg/disttask/framework/proto",
         "//pkg/disttask/framework/scheduler",
         "//pkg/disttask/framework/storage",
+        "//pkg/disttask/framework/taskexecutor/execute",
         "//pkg/disttask/framework/testutil",
         "//pkg/kv",
         "//pkg/testkit",

pkg/disttask/framework/taskexecutor/execute/interface.go

@@ -16,6 +16,7 @@ package execute
 
 import (
 	"context"
+	"reflect"
 
 	"github.com/pingcap/tidb/pkg/disttask/framework/proto"
 )
@@ -29,6 +30,8 @@ import (
 //		else OnFinished
 //	Cleanup
 type StepExecutor interface {
+	StepExecFrameworkInfo
+
 	// Init is used to initialize the environment.
 	// if failed, task executor will retry later.
 	Init(context.Context) error
@@ -49,3 +52,49 @@ type StepExecutor interface {
 type SubtaskSummary struct {
 	RowCount int64
 }
+
+// StepExecFrameworkInfo is an interface that should be embedded into the
+// implementation of StepExecutor. It's set by the framework automatically and
+// the implementation can use it to access necessary information. The framework
+// will init it before `StepExecutor.Init`, before that you cannot call methods
+// in this interface.
+type StepExecFrameworkInfo interface {
+	// restricted is a private method to prevent other package mistakenly implements
+	// StepExecFrameworkInfo. So when StepExecFrameworkInfo is composed with other
+	// interfaces, the implementation of other interface must embed
+	// StepExecFrameworkInfo.
+	restricted()
+	// GetResource returns the expected resource of this step executor.
+	GetResource() *proto.StepResource
+}
+
+var stepExecFrameworkInfoName = reflect.TypeOf((*StepExecFrameworkInfo)(nil)).Elem().Name()
+
+type frameworkInfo struct {
+	resource *proto.StepResource
+}
+
+func (*frameworkInfo) restricted() {}
+
+func (f *frameworkInfo) GetResource() *proto.StepResource {
+	return f.resource
+}
+
+// SetFrameworkInfo sets the framework info for the StepExecutor.
+func SetFrameworkInfo(exec StepExecutor, resource *proto.StepResource) {
+	if exec == nil {
+		return
+	}
+	toInject := &frameworkInfo{resource: resource}
+	// use reflection to set the framework info
+	e := reflect.ValueOf(exec)
+	if e.Kind() == reflect.Ptr || e.Kind() == reflect.Interface {
+		e = e.Elem()
+	}
+	info := e.FieldByName(stepExecFrameworkInfoName)
+	// if `exec` embeds StepExecutor rather than StepExecFrameworkInfo, the field
+	// will not be found. This is happened in mock generated code.
+	if info.IsValid() && info.CanSet() {
+		info.Set(reflect.ValueOf(toInject))
+	}
+}