test: fix data race in distsql/select_result (#18059)

This commit is contained in:
Evan Zhou
2020-06-17 11:14:24 +08:00
committed by GitHub
parent 7f03358895
commit 514b91d90c

View File

@ -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()
}