From 3d2b00dac5d9faf401459408b3e395378b77dec3 Mon Sep 17 00:00:00 2001 From: wshwsh12 <793703860@qq.com> Date: Thu, 2 Feb 2023 17:29:53 +0800 Subject: [PATCH] cherry-pick 40901 --- distsql/distsql.go | 5 +++-- executor/builder.go | 2 ++ executor/mpp_gather.go | 5 ++++- executor/tiflashtest/tiflash_test.go | 23 +++++++++++++++++++++++ kv/mpp.go | 3 ++- store/copr/mpp.go | 23 ++++++++++++++++++++++- 6 files changed, 56 insertions(+), 5 deletions(-) diff --git a/distsql/distsql.go b/distsql/distsql.go index 1f18d084eb0b2..db80d626356d7 100644 --- a/distsql/distsql.go +++ b/distsql/distsql.go @@ -30,6 +30,7 @@ import ( "github.com/pingcap/tidb/statistics" "github.com/pingcap/tidb/types" "github.com/pingcap/tidb/util/logutil" + "github.com/pingcap/tidb/util/memory" "github.com/pingcap/tidb/util/trxevents" "github.com/pingcap/tipb/go-tipb" "github.com/tikv/client-go/v2/tikvrpc/interceptor" @@ -38,11 +39,11 @@ import ( ) // DispatchMPPTasks dispatches all tasks and returns an iterator. -func DispatchMPPTasks(ctx context.Context, sctx sessionctx.Context, tasks []*kv.MPPDispatchRequest, fieldTypes []*types.FieldType, planIDs []int, rootID int, startTs uint64) (SelectResult, error) { +func DispatchMPPTasks(ctx context.Context, sctx sessionctx.Context, tasks []*kv.MPPDispatchRequest, fieldTypes []*types.FieldType, planIDs []int, rootID int, startTs uint64, memTracker *memory.Tracker) (SelectResult, error) { ctx = WithSQLKvExecCounterInterceptor(ctx, sctx.GetSessionVars().StmtCtx) _, allowTiFlashFallback := sctx.GetSessionVars().AllowFallbackToTiKV[kv.TiFlash] ctx = SetTiFlashMaxThreadsInContext(ctx, sctx) - resp := sctx.GetMPPClient().DispatchMPPTasks(ctx, sctx.GetSessionVars().KVVars, tasks, allowTiFlashFallback, startTs) + resp := sctx.GetMPPClient().DispatchMPPTasks(ctx, sctx.GetSessionVars().KVVars, tasks, allowTiFlashFallback, startTs, memTracker) if resp == nil { return nil, errors.New("client returns nil response") } diff --git a/executor/builder.go b/executor/builder.go index d4270397eecd0..ab0ad87efc891 100644 --- a/executor/builder.go +++ b/executor/builder.go @@ -3402,7 +3402,9 @@ func (b *executorBuilder) buildMPPGather(v *plannercore.PhysicalTableReader) Exe is: b.is, originalPlan: v.GetTablePlan(), startTS: startTs, + memTracker: memory.NewTracker(v.ID(), -1), } + gather.memTracker.AttachTo(b.ctx.GetSessionVars().StmtCtx.MemTracker) return gather } diff --git a/executor/mpp_gather.go b/executor/mpp_gather.go index 42526774dbdd5..38108adf4f59b 100644 --- a/executor/mpp_gather.go +++ b/executor/mpp_gather.go @@ -26,6 +26,7 @@ import ( "github.com/pingcap/tidb/sessionctx" "github.com/pingcap/tidb/util/chunk" "github.com/pingcap/tidb/util/logutil" + "github.com/pingcap/tidb/util/memory" "github.com/pingcap/tipb/go-tipb" "go.uber.org/zap" ) @@ -49,6 +50,8 @@ type MPPGather struct { mppReqs []*kv.MPPDispatchRequest respIter distsql.SelectResult + + memTracker *memory.Tracker } func (e *MPPGather) appendMPPDispatchReq(pf *plannercore.Fragment) error { @@ -124,7 +127,7 @@ func (e *MPPGather) Open(ctx context.Context) (err error) { failpoint.Return(errors.Errorf("The number of tasks is not right, expect %d tasks but actually there are %d tasks", val.(int), len(e.mppReqs))) } }) - e.respIter, err = distsql.DispatchMPPTasks(ctx, e.ctx, e.mppReqs, e.retFieldTypes, planIDs, e.id, e.startTS) + e.respIter, err = distsql.DispatchMPPTasks(ctx, e.ctx, e.mppReqs, e.retFieldTypes, planIDs, e.id, e.startTS, e.memTracker) if err != nil { return errors.Trace(err) } diff --git a/executor/tiflashtest/tiflash_test.go b/executor/tiflashtest/tiflash_test.go index 2198182a8950e..4346c7dd8277e 100644 --- a/executor/tiflashtest/tiflash_test.go +++ b/executor/tiflashtest/tiflash_test.go @@ -1256,3 +1256,26 @@ func TestTiflashEmptyDynamicPruneResult(t *testing.T) { tk.MustQuery("select /*+ read_from_storage(tiflash[t2]) */ * from IDT_RP24833 partition(p2) t2 where t2. col1 <= -8448770111093677011;").Check(testkit.Rows()) tk.MustQuery("select /*+ read_from_storage(tiflash[t1, t2]) */ * from IDT_RP24833 partition(p3, p4) t1 join IDT_RP24833 partition(p2) t2 on t1.col1 = t2.col1 where t1. col1 between -8448770111093677011 and -8448770111093677011 and t2. col1 <= -8448770111093677011;").Check(testkit.Rows()) } + +func TestMPPMemoryTracker(t *testing.T) { + store := testkit.CreateMockStore(t, withMockTiFlash(2)) + tk := testkit.NewTestKit(t, store) + tk.MustExec("set tidb_mem_quota_query = 1 << 30") + tk.MustExec("set global tidb_mem_oom_action = 'CANCEL'") + tk.MustExec("use test") + tk.MustExec("create table t(a int);") + tk.MustExec("insert into t values (1);") + tk.MustExec("alter table t set tiflash replica 1") + tb := external.GetTableByName(t, tk, "test", "t") + err := domain.GetDomain(tk.Session()).DDL().UpdateTableReplicaInfo(tk.Session(), tb.Meta().ID, true) + require.NoError(t, err) + tk.MustExec("set tidb_enforce_mpp = on;") + tk.MustQuery("select * from t").Check(testkit.Rows("1")) + require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/store/copr/testMPPOOMPanic", `return(true)`)) + defer func() { + require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/store/copr/testMPPOOMPanic")) + }() + err = tk.QueryToErr("select * from t") + require.NotNil(t, err) + require.True(t, strings.Contains(err.Error(), "Out Of Memory Quota!")) +} diff --git a/kv/mpp.go b/kv/mpp.go index 2e398af595650..3acf89c4ef529 100644 --- a/kv/mpp.go +++ b/kv/mpp.go @@ -19,6 +19,7 @@ import ( "time" "github.com/pingcap/kvproto/pkg/mpp" + "github.com/pingcap/tidb/util/memory" ) // MPPTaskMeta means the meta info such as location of a mpp task. @@ -83,7 +84,7 @@ type MPPClient interface { ConstructMPPTasks(context.Context, *MPPBuildTasksRequest, time.Duration) ([]MPPTaskMeta, error) // DispatchMPPTasks dispatches ALL mpp requests at once, and returns an iterator that transfers the data. - DispatchMPPTasks(ctx context.Context, vars interface{}, reqs []*MPPDispatchRequest, needTriggerFallback bool, startTs uint64) Response + DispatchMPPTasks(ctx context.Context, vars interface{}, reqs []*MPPDispatchRequest, needTriggerFallback bool, startTs uint64, memTracker *memory.Tracker) Response } // MPPBuildTasksRequest request the stores allocation for a mpp plan fragment. diff --git a/store/copr/mpp.go b/store/copr/mpp.go index b7464ae8ac66c..ca1976885a003 100644 --- a/store/copr/mpp.go +++ b/store/copr/mpp.go @@ -16,6 +16,7 @@ package copr import ( "context" + "fmt" "io" "strconv" "sync" @@ -35,6 +36,7 @@ import ( "github.com/pingcap/tidb/util" "github.com/pingcap/tidb/util/logutil" "github.com/pingcap/tidb/util/mathutil" + "github.com/pingcap/tidb/util/memory" "github.com/tikv/client-go/v2/tikv" "github.com/tikv/client-go/v2/tikvrpc" "go.uber.org/zap" @@ -160,6 +162,8 @@ type mppIterator struct { mu sync.Mutex enableCollectExecutionInfo bool + + memTracker *memory.Tracker } func (m *mppIterator) run(ctx context.Context) { @@ -197,6 +201,22 @@ func (m *mppIterator) sendError(err error) { } func (m *mppIterator) sendToRespCh(resp *mppResponse) (exit bool) { + defer func() { + if r := recover(); r != nil { + logutil.BgLogger().Error("mppIterator panic", zap.Stack("stack"), zap.Any("recover", r)) + m.sendError(errors.New(fmt.Sprint(r))) + } + }() + if m.memTracker != nil { + respSize := resp.MemSize() + failpoint.Inject("testMPPOOMPanic", func(val failpoint.Value) { + if val.(bool) && respSize != 0 { + respSize = 1 << 30 + } + }) + m.memTracker.Consume(respSize) + defer m.memTracker.Consume(-respSize) + } select { case m.respChan <- resp: case <-m.finishCh: @@ -508,7 +528,7 @@ func (m *mppIterator) Next(ctx context.Context) (kv.ResultSubset, error) { } // DispatchMPPTasks dispatches all the mpp task and waits for the responses. -func (c *MPPClient) DispatchMPPTasks(ctx context.Context, variables interface{}, dispatchReqs []*kv.MPPDispatchRequest, needTriggerFallback bool, startTs uint64) kv.Response { +func (c *MPPClient) DispatchMPPTasks(ctx context.Context, variables interface{}, dispatchReqs []*kv.MPPDispatchRequest, needTriggerFallback bool, startTs uint64, memTracker *memory.Tracker) kv.Response { vars := variables.(*tikv.Variables) ctxChild, cancelFunc := context.WithCancel(ctx) iter := &mppIterator{ @@ -521,6 +541,7 @@ func (c *MPPClient) DispatchMPPTasks(ctx context.Context, variables interface{}, vars: vars, needTriggerFallback: needTriggerFallback, enableCollectExecutionInfo: config.GetGlobalConfig().Instance.EnableCollectExecutionInfo.Load(), + memTracker: memTracker, } go iter.run(ctxChild) return iter