*: add resource group name when dispatch req to tiflash (#47096)

close #47093
pingcap · Sep 20, 2023 · 3f3b07e · 3f3b07e
1 parent e8d2223
commit 3f3b07e
Show file tree

Hide file tree

Showing 4 changed files with 6 additions and 0 deletions.
diff --git a/executor/internal/mpp/local_mpp_coordinator.go b/executor/internal/mpp/local_mpp_coordinator.go
@@ -229,6 +229,7 @@ func (c *localMppCoordinator) appendMPPDispatchReq(pf *plannercore.Fragment) err
 			CoordinatorAddress:     c.coordinatorAddr,
 			ReportExecutionSummary: c.reportExecutionInfo,
 			State:                  kv.MppTaskReady,
+			ResourceGroupName:      c.sessionCtx.GetSessionVars().ResourceGroupName,
 		}
 		c.reqMap[req.ID] = &mppRequestReport{mppReq: req, receivedReport: false, errMsg: "", executionSummaries: nil}
 		c.mppReqs = append(c.mppReqs, req)

diff --git a/kv/mpp.go b/kv/mpp.go
@@ -154,6 +154,7 @@ type MPPDispatchRequest struct {
 	CoordinatorAddress     string
 	ReportExecutionSummary bool
 	State                  MppTaskStates
+	ResourceGroupName      string
 }
 
 // CancelMPPTasksParam represents parameter for MPPClient's CancelMPPTasks

diff --git a/store/copr/batch_coprocessor.go b/store/copr/batch_coprocessor.go
@@ -1276,6 +1276,9 @@ func (b *batchCopIterator) handleTaskOnce(ctx context.Context, bo *backoff.Backo
 		RecordTimeStat: true,
 		RecordScanStat: true,
 		TaskId:         b.req.TaskID,
+		ResourceControlContext: &kvrpcpb.ResourceControlContext{
+			ResourceGroupName: b.req.ResourceGroupName,
+		},
 	})
 	if b.req.ResourceGroupTagger != nil {
 		b.req.ResourceGroupTagger(req)

diff --git a/store/copr/mpp.go b/store/copr/mpp.go
@@ -108,6 +108,7 @@ func (c *MPPClient) DispatchMPPTask(param kv.DispatchMPPTaskParam) (resp *mpp.Di
 		CoordinatorAddress:     req.CoordinatorAddress,
 		ReportExecutionSummary: req.ReportExecutionSummary,
 		MppVersion:             req.MppVersion.ToInt64(),
+		ResourceGroupName:      req.ResourceGroupName,
 	}
 
 	mppReq := &mpp.DispatchTaskRequest{