From cca066ded4c8234251dc63a235122bb0a3925d04 Mon Sep 17 00:00:00 2001
From: Sugu Sougoumarane <ssougou@gmail.com>
Date: Sat, 30 Nov 2019 16:07:45 -0800
Subject: [PATCH] vdiff: fix data race in test

Fixes #5489

Signed-off-by: Sugu Sougoumarane <ssougou@gmail.com>
---
 go/vt/wrangler/vdiff_env_test.go | 19 ++++++++++++-------
 1 file changed, 12 insertions(+), 7 deletions(-)

diff --git a/go/vt/wrangler/vdiff_env_test.go b/go/vt/wrangler/vdiff_env_test.go
index f41db6315e8..44f4f9b85ee 100644
--- a/go/vt/wrangler/vdiff_env_test.go
+++ b/go/vt/wrangler/vdiff_env_test.go
@@ -19,6 +19,7 @@ package wrangler
 import (
 	"flag"
 	"fmt"
+	"sync"
 
 	"golang.org/x/net/context"
 	"vitess.io/vitess/go/sqltypes"
@@ -51,11 +52,13 @@ const (
 type testVDiffEnv struct {
 	wr         *Wrangler
 	workflow   string
-	tablets    map[int]*testVDiffTablet
 	topoServ   *topo.Server
 	cell       string
 	tabletType topodatapb.TabletType
 	tmc        *testVDiffTMClient
+
+	mu      sync.Mutex
+	tablets map[int]*testVDiffTablet
 }
 
 // vdiffEnv has to be a global for RegisterDialer to work.
@@ -63,6 +66,8 @@ var vdiffEnv *testVDiffEnv
 
 func init() {
 	tabletconn.RegisterDialer("VDiffTest", func(tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
+		vdiffEnv.mu.Lock()
+		defer vdiffEnv.mu.Unlock()
 		return vdiffEnv.tablets[int(tablet.Alias.Uid)], nil
 	})
 }
@@ -163,12 +168,17 @@ func newTestVDiffEnv(sourceShards, targetShards []string, query string, position
 }
 
 func (env *testVDiffEnv) close() {
+	env.mu.Lock()
+	defer env.mu.Unlock()
 	for _, t := range env.tablets {
-		env.deleteTablet(t.tablet)
+		env.topoServ.DeleteTablet(context.Background(), t.tablet.Alias)
 	}
+	env.tablets = nil
 }
 
 func (env *testVDiffEnv) addTablet(id int, keyspace, shard string, tabletType topodatapb.TabletType) *testVDiffTablet {
+	env.mu.Lock()
+	defer env.mu.Unlock()
 	tablet := &topodatapb.Tablet{
 		Alias: &topodatapb.TabletAlias{
 			Cell: env.cell,
@@ -198,11 +208,6 @@ func (env *testVDiffEnv) addTablet(id int, keyspace, shard string, tabletType to
 	return env.tablets[id]
 }
 
-func (env *testVDiffEnv) deleteTablet(tablet *topodatapb.Tablet) {
-	env.topoServ.DeleteTablet(context.Background(), tablet.Alias)
-	delete(env.tablets, int(tablet.Alias.Uid))
-}
-
 //----------------------------------------------
 // testVDiffTablet