executor: fix the issue that pb memory cannot be released quickly

executor/memory_test.go

@@ -0,0 +1,109 @@
+// Copyright 2019 PingCAP, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package executor_test
+
+import (
+ "context"
+ "fmt"
+ "runtime"
+
+ . "github.com/pingcap/check"
+ "github.com/pingcap/tidb/domain"
+ "github.com/pingcap/tidb/kv"
+ "github.com/pingcap/tidb/session"
+ "github.com/pingcap/tidb/store/mockstore"
+)
+
+var _ = SerialSuites(&testMemoryLeak{})
+
+type testMemoryLeak struct {
+ store kv.Storage
+ domain *domain.Domain
+}
+
+func (s *testMemoryLeak) SetUpSuite(c *C) {
+ var err error
+ s.store, err = mockstore.NewMockTikvStore()
+ c.Assert(err, IsNil)
+ s.domain, err = session.BootstrapSession(s.store)
+ c.Assert(err, IsNil)
+}
+
+func (s *testMemoryLeak) TestPBMemoryLeak(c *C) {
+ se, err := session.CreateSession4Test(s.store)
+ c.Assert(err, IsNil)
+ _, err = se.Execute(context.Background(), "create database test_mem")
+ c.Assert(err, IsNil)
+ _, err = se.Execute(context.Background(), "use test_mem")
+ c.Assert(err, IsNil)
+
+ // prepare data
+ totalSize := uint64(256 << 20) // 256MB
+ blockSize := uint64(8 << 10) // 8KB
+ delta := totalSize / 5
+ numRows := totalSize / blockSize
+ _, err = se.Execute(context.Background(), fmt.Sprintf("create table t (c varchar(%v))", blockSize))
+ c.Assert(err, IsNil)
+ defer func() {
+ _, err = se.Execute(context.Background(), "drop table t")
+ c.Assert(err, IsNil)
+ }()
+ sql := fmt.Sprintf("insert into t values (space(%v))", blockSize)
+ for i := uint64(0); i < numRows; i++ {
+ _, err = se.Execute(context.Background(), sql)
+ c.Assert(err, IsNil)
+ }
+
+ // read data
+ runtime.GC()
+ allocatedBegin, inUseBegin := s.readMem()
+ records, err := se.Execute(context.Background(), "select * from t")
+ c.Assert(err, IsNil)
+ record := records[0]
+ rowCnt := 0
+ chk := record.NewRecordBatch()
+ for {
+ c.Assert(record.Next(context.Background(), chk), IsNil)
+ rowCnt += chk.NumRows()
+ if chk.NumRows() == 0 {
+ break
+ }
+ }
+ c.Assert(rowCnt, Equals, int(numRows))
+
+ // check memory before close
+ runtime.GC()
+ allocatedAfter, inUseAfter := s.readMem()
+ c.Assert(allocatedAfter-allocatedBegin, GreaterEqual, totalSize)
+ c.Assert(s.memDiff(inUseAfter, inUseBegin), Less, delta)
+
+ se.Close()
+ runtime.GC()
+ allocatedFinal, inUseFinal := s.readMem()
+ c.Assert(allocatedFinal-allocatedAfter, Less, delta)
+ c.Assert(s.memDiff(inUseFinal, inUseAfter), Less, delta)
+}
+
+func (s *testMemoryLeak) readMem() (allocated, heapInUse uint64) {
+ var stat runtime.MemStats
+ runtime.ReadMemStats(&stat)
+ return stat.TotalAlloc, stat.HeapInuse
+}
+
+func (s *testMemoryLeak) memDiff(m1, m2 uint64) uint64 {
+ if m1 > m2 {
+ return m1 - m2
+ }
+ return m2 - m1
+}

store/tikv/coprocessor.go

@@ -367,8 +367,8 @@ type copIteratorTaskSender struct {
 }
 
 type copResponse struct {
- pbResp *coprocessor.Response
- execdetails.ExecDetails
+ pbResp  *coprocessor.Response
+ detail *execdetails.ExecDetails
  startKey kv.Key
  err error
  respSize int64
@@ -390,7 +390,7 @@ func (rs *copResponse) GetStartKey() kv.Key {
 }
 
 func (rs *copResponse) GetExecDetails() *execdetails.ExecDetails {
- return &rs.ExecDetails
+ return rs.detail
 }
 
 // MemSize returns how many bytes of memory this response use
@@ -401,9 +401,11 @@ func (rs *copResponse) MemSize() int64 {
 
  // ignore rs.err
  rs.respSize += int64(cap(rs.startKey))
- rs.respSize += int64(sizeofExecDetails)
- if rs.CommitDetail != nil {
- rs.respSize += int64(sizeofCommitDetails)
+ if rs.detail != nil {
+ rs.respSize += int64(sizeofExecDetails)
+ if rs.detail.CommitDetail != nil {
+ rs.respSize += int64(sizeofCommitDetails)
+ }
  }
  if rs.pbResp != nil {
  // Using a approximate size since it's hard to get a accurate value.
@@ -768,19 +770,22 @@ func (worker *copIteratorWorker) handleCopResponse(bo *Backoffer, rpcCtx *RPCCon
  } else {
  resp.startKey = task.ranges.at(0).StartKey
  }
- resp.BackoffTime = time.Duration(bo.totalSleep) * time.Millisecond
+ if resp.detail == nil {
+ resp.detail = new(execdetails.ExecDetails)
+ }
+ resp.detail.BackoffTime = time.Duration(bo.totalSleep) * time.Millisecond
  if rpcCtx != nil {
- resp.CalleeAddress = rpcCtx.Addr
+ resp.detail.CalleeAddress = rpcCtx.Addr
  }
  if pbDetails := resp.pbResp.ExecDetails; pbDetails != nil {
  if handleTime := pbDetails.HandleTime; handleTime != nil {
- resp.WaitTime = time.Duration(handleTime.WaitMs) * time.Millisecond
- resp.ProcessTime = time.Duration(handleTime.ProcessMs) * time.Millisecond
+ resp.detail.WaitTime = time.Duration(handleTime.WaitMs) * time.Millisecond
+ resp.detail.ProcessTime = time.Duration(handleTime.ProcessMs) * time.Millisecond
  }
  if scanDetail := pbDetails.ScanDetail; scanDetail != nil {
  if scanDetail.Write != nil {
- resp.TotalKeys += scanDetail.Write.Total
- resp.ProcessedKeys += scanDetail.Write.Processed
+ resp.detail.TotalKeys += scanDetail.Write.Total
+ resp.detail.ProcessedKeys += scanDetail.Write.Processed
  }
  }
  }