From 514b91d90c701b549d1f7e8bf1e4b450a77f4d11 Mon Sep 17 00:00:00 2001 From: Evan Zhou Date: Wed, 17 Jun 2020 11:14:24 +0800 Subject: [PATCH] test: fix data race in distsql/select_result (#18059) --- distsql/select_result.go | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/distsql/select_result.go b/distsql/select_result.go index ab21ec00fc..154c966826 100644 --- a/distsql/select_result.go +++ b/distsql/select_result.go @@ -66,7 +66,7 @@ type selectResult struct { ctx sessionctx.Context selectResp *tipb.SelectResponse - selectRespSize int // record the selectResp.Size() when it is initialized. + selectRespSize int64 // record the selectResp.Size() when it is initialized. respChkIdx int respChunkDecoder *chunk.Decoder @@ -99,10 +99,11 @@ func (r *selectResult) fetchResp(ctx context.Context) error { return errors.Trace(err) } if r.selectResp != nil { - r.memConsume(-int64(r.selectRespSize)) + r.memConsume(-atomic.LoadInt64(&r.selectRespSize)) } if resultSubset == nil { r.selectResp = nil + atomic.StoreInt64(&r.selectRespSize, 0) if !r.durationReported { // final round of fetch // TODO: Add a label to distinguish between success or failure. @@ -117,8 +118,9 @@ func (r *selectResult) fetchResp(ctx context.Context) error { if err != nil { return errors.Trace(err) } - r.selectRespSize = r.selectResp.Size() - r.memConsume(int64(r.selectRespSize)) + respSize := int64(r.selectResp.Size()) + atomic.StoreInt64(&r.selectRespSize, respSize) + r.memConsume(respSize) if err := r.selectResp.Error; err != nil { return terror.ClassTiKV.Synthesize(terror.ErrCode(err.Code), err.Msg) } @@ -284,8 +286,9 @@ func (r *selectResult) Close() error { metrics.DistSQLScanKeysHistogram.Observe(float64(r.feedback.Actual())) } metrics.DistSQLPartialCountHistogram.Observe(float64(r.partialCount)) - if r.selectResp != nil { - r.memConsume(-int64(r.selectRespSize)) + respSize := atomic.SwapInt64(&r.selectRespSize, 0) + if respSize > 0 { + r.memConsume(-respSize) } return r.resp.Close() }