From 6391d58d25124b1506c237b0a9bf44fd3bb67fa3 Mon Sep 17 00:00:00 2001 From: Feng Liyuan Date: Mon, 1 Jul 2019 21:45:09 +0800 Subject: [PATCH] =?UTF-8?q?distsql:=20clean=20the=20memory=20usage=20of=20?= =?UTF-8?q?MemTracker=20when=20a=20query=20end=E2=80=A6=20(#10970)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- distsql/distsql_test.go | 17 +++++++++++++ distsql/request_builder_test.go | 5 ++++ distsql/select_result.go | 45 +++++++++++++++++++++++---------- 3 files changed, 53 insertions(+), 14 deletions(-) diff --git a/distsql/distsql_test.go b/distsql/distsql_test.go index 478881e47720f..6fda7832eb474 100644 --- a/distsql/distsql_test.go +++ b/distsql/distsql_test.go @@ -42,6 +42,7 @@ func (s *testSuite) createSelectNormal(batch, totalRows int, c *C, planIDs []str SetDesc(false). SetKeepOrder(false). SetFromSessionVars(variable.NewSessionVars()). + SetMemTracker(s.sctx, stringutil.StringerStr("testSuite.createSelectNormal")). Build() c.Assert(err, IsNil) @@ -106,6 +107,21 @@ func (s *testSuite) TestSelectNormal(c *C) { c.Assert(numAllRows, Equals, 2) err := response.Close() c.Assert(err, IsNil) + c.Assert(response.memTracker.BytesConsumed(), Equals, int64(0)) +} + +func (s *testSuite) TestSelectMemTracker(c *C) { + response, colTypes := s.createSelectNormal(2, 6, c, nil) + response.Fetch(context.TODO()) + + // Test Next. + chk := chunk.New(colTypes, 3, 3) + err := response.Next(context.TODO(), chk) + c.Assert(err, IsNil) + c.Assert(chk.IsFull(), Equals, true) + err = response.Close() + c.Assert(err, IsNil) + c.Assert(response.memTracker.BytesConsumed(), Equals, int64(0)) } func (s *testSuite) TestSelectNormalChunkSize(c *C) { @@ -113,6 +129,7 @@ func (s *testSuite) TestSelectNormalChunkSize(c *C) { response.Fetch(context.TODO()) s.testChunkSize(response, colTypes, c) c.Assert(response.Close(), IsNil) + c.Assert(response.memTracker.BytesConsumed(), Equals, int64(0)) } func (s *testSuite) TestSelectWithRuntimeStats(c *C) { diff --git a/distsql/request_builder_test.go b/distsql/request_builder_test.go index b64dd63218892..a2b472b5ad833 100644 --- a/distsql/request_builder_test.go +++ b/distsql/request_builder_test.go @@ -26,8 +26,10 @@ import ( "github.com/pingcap/tidb/types" "github.com/pingcap/tidb/util/codec" "github.com/pingcap/tidb/util/logutil" + "github.com/pingcap/tidb/util/memory" "github.com/pingcap/tidb/util/mock" "github.com/pingcap/tidb/util/ranger" + "github.com/pingcap/tidb/util/stringutil" "github.com/pingcap/tidb/util/testleak" "github.com/pingcap/tipb/go-tipb" ) @@ -49,6 +51,9 @@ type testSuite struct { func (s *testSuite) SetUpSuite(c *C) { ctx := mock.NewContext() + ctx.GetSessionVars().StmtCtx = &stmtctx.StatementContext{ + MemTracker: memory.NewTracker(stringutil.StringerStr("testSuite"), variable.DefTiDBMemQuotaDistSQL), + } ctx.Store = &mock.Store{ Client: &mock.Client{ MockResponse: &mockResponse{ diff --git a/distsql/select_result.go b/distsql/select_result.go index 72a283d1b2f8f..3935df52ad450 100644 --- a/distsql/select_result.go +++ b/distsql/select_result.go @@ -66,8 +66,9 @@ type selectResult struct { fieldTypes []*types.FieldType ctx sessionctx.Context - selectResp *tipb.SelectResponse - respChkIdx int + selectResp *tipb.SelectResponse + selectRespSize int // record the selectResp.Size() when it is initialized. + respChkIdx int feedback *statistics.QueryFeedback partialCount int64 // number of partial results. @@ -103,20 +104,25 @@ func (r *selectResult) fetch(ctx context.Context) { if err != nil { result.err = err } else if resultSubset == nil { + // If the result is drained, the resultSubset would be nil return } else { result.result = resultSubset - if r.memTracker != nil { - r.memTracker.Consume(int64(resultSubset.MemSize())) - } + r.memConsume(int64(resultSubset.MemSize())) } select { case r.results <- result: case <-r.closed: // If selectResult called Close() already, make fetch goroutine exit. + if resultSubset != nil { + r.memConsume(-int64(resultSubset.MemSize())) + } return case <-ctx.Done(): + if resultSubset != nil { + r.memConsume(-int64(resultSubset.MemSize())) + } return } } @@ -161,24 +167,21 @@ func (r *selectResult) getSelectResp() error { if re.err != nil { return errors.Trace(re.err) } - if r.memTracker != nil && r.selectResp != nil { - r.memTracker.Consume(-int64(r.selectResp.Size())) + if r.selectResp != nil { + r.memConsume(-int64(r.selectRespSize)) } if re.result == nil { r.selectResp = nil return nil } - if r.memTracker != nil { - r.memTracker.Consume(-int64(re.result.MemSize())) - } + r.memConsume(-int64(re.result.MemSize())) r.selectResp = new(tipb.SelectResponse) err := r.selectResp.Unmarshal(re.result.GetData()) if err != nil { return errors.Trace(err) } - if r.memTracker != nil && r.selectResp != nil { - r.memTracker.Consume(int64(r.selectResp.Size())) - } + r.selectRespSize = r.selectResp.Size() + r.memConsume(int64(r.selectRespSize)) if err := r.selectResp.Error; err != nil { return terror.ClassTiKV.New(terror.ErrCode(err.Code), err.Msg) } @@ -234,13 +237,27 @@ func (r *selectResult) readRowsData(chk *chunk.Chunk) (err error) { return nil } +func (r *selectResult) memConsume(bytes int64) { + if r.memTracker != nil { + r.memTracker.Consume(bytes) + } +} + // Close closes selectResult. func (r *selectResult) Close() error { - // Close this channel tell fetch goroutine to exit. if r.feedback.Actual() >= 0 { metrics.DistSQLScanKeysHistogram.Observe(float64(r.feedback.Actual())) } metrics.DistSQLPartialCountHistogram.Observe(float64(r.partialCount)) + // Close this channel to tell the fetch goroutine to exit. close(r.closed) + for re := range r.results { + if re.result != nil { + r.memConsume(-int64(re.result.MemSize())) + } + } + if r.selectResp != nil { + r.memConsume(-int64(r.selectRespSize)) + } return r.resp.Close() }