tablet_picker: keep trying to find a tablet until context expires (vi…

…tessio#6546) * tablet_picker: keep trying to find a tablet until context expires Signed-off-by: deepthi <[email protected]> * tablet_picker: keyspace/shard/cells must be provided Signed-off-by: deepthi <[email protected]> * tablet picker: return ks/cell/shard in test result for TestShardMigrateMainflow Signed-off-by: Rohit Nayak <[email protected]> * tablet picker: fix wrangler tests by choosing a valid source per-test Signed-off-by: Rohit Nayak <[email protected]> * tablet picker: fix worker and testlib tests Signed-off-by: Rohit Nayak <[email protected]> * tablet_picker: gofmted Signed-off-by: Rohit Nayak <[email protected]> * tablet picker: fix test races Signed-off-by: Rohit Nayak <[email protected]> * tablet picker: gofmt Signed-off-by: Rohit Nayak <[email protected]> * tablet picker: gofmt Signed-off-by: Rohit Nayak <[email protected]> * healthcheck: cleanup test code where we were setting the same flag twice Signed-off-by: deepthi <[email protected]> Co-authored-by: Rohit Nayak <[email protected]>
tinyspeck · Aug 14, 2020 · e16a6c5 · e16a6c5
1 parent 1d74017
commit e16a6c5
Show file tree

Hide file tree

Showing 24 changed files with 608 additions and 148 deletions.
diff --git a/go/test/endtoend/vreplication/cluster.go b/go/test/endtoend/vreplication/cluster.go
@@ -124,8 +124,8 @@ func InitCluster(t *testing.T, cellNames []string) *VitessCluster {
 		globalConfig.topoPort, globalConfig.hostname, globalConfig.tmpDir)
 	vc.Vtctld = vtctld
 	assert.NotNil(t, vc.Vtctld)
-	// use first cell as `-cell` and all cells as `-cells_to_watch`
-	vc.Vtctld.Setup(cellNames[0], "-cells_to_watch", strings.Join(cellNames, ","))
+	// use first cell as `-cell`
+	vc.Vtctld.Setup(cellNames[0])
 
 	vc.Vtctl = cluster.VtctlProcessInstance(globalConfig.topoPort, globalConfig.hostname)
 	assert.NotNil(t, vc.Vtctl)

diff --git a/go/vt/discovery/healthcheck_test.go b/go/vt/discovery/healthcheck_test.go
@@ -48,13 +48,17 @@ import (
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
 )
 
+var connMap map[string]*fakeConn
+var connMapMu sync.Mutex
+
 func init() {
 	tabletconn.RegisterDialer("fake_gateway", tabletDialer)
 
 	//log error
 	if err := flag.Set("tablet_protocol", "fake_gateway"); err != nil {
 		log.Errorf("failed to set flag \"tablet_protocol\" to \"fake_gateway\":%v", err)
 	}
+	connMap = make(map[string]*fakeConn)
 }
 
 func TestHealthCheck(t *testing.T) {
@@ -85,9 +89,10 @@ func TestHealthCheck(t *testing.T) {
 	mustMatch(t, want, result, "Wrong TabletHealth data")
 
 	shr := &querypb.StreamHealthResponse{
-		TabletAlias:                         tablet.Alias,
-		Target:                              &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
-		Serving:                             true,
+		TabletAlias: tablet.Alias,
+		Target:      &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
+		Serving:     true,
+
 		TabletExternallyReparentedTimestamp: 0,
 		RealtimeStats:                       &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.5},
 	}
@@ -323,10 +328,11 @@ func TestHealthCheckCloseWaitsForGoRoutines(t *testing.T) {
 		RealtimeStats:                       &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.2},
 	}
 	want = &TabletHealth{
-		Tablet:              tablet,
-		Target:              &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
-		Serving:             true,
-		Stats:               &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.2},
+		Tablet:  tablet,
+		Target:  &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
+		Serving: true,
+		Stats:   &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.2},
+
 		MasterTermStartTime: 0,
 	}
 	input <- shr
@@ -549,11 +555,13 @@ func TestGetHealthyTablets(t *testing.T) {
 
 	// second tablet turns into a master
 	shr2 = &querypb.StreamHealthResponse{
-		TabletAlias:                         tablet2.Alias,
-		Target:                              &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_MASTER},
-		Serving:                             true,
+		TabletAlias: tablet2.Alias,
+		Target:      &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_MASTER},
+		Serving:     true,
+
 		TabletExternallyReparentedTimestamp: 10,
-		RealtimeStats:                       &querypb.RealtimeStats{SecondsBehindMaster: 0, CpuUsage: 0.2},
+
+		RealtimeStats: &querypb.RealtimeStats{SecondsBehindMaster: 0, CpuUsage: 0.2},
 	}
 	input2 <- shr2
 	// wait for result
@@ -766,6 +774,9 @@ func TestDebugURLFormatting(t *testing.T) {
 }
 
 func tabletDialer(tablet *topodatapb.Tablet, _ grpcclient.FailFast) (queryservice.QueryService, error) {
+	connMapMu.Lock()
+	defer connMapMu.Unlock()
+
 	key := TabletToMapKey(tablet)
 	if qs, ok := connMap[key]; ok {
 		return qs, nil
@@ -794,6 +805,8 @@ type fakeConn struct {
 }
 
 func createFakeConn(tablet *topodatapb.Tablet, c chan *querypb.StreamHealthResponse) *fakeConn {
+	connMapMu.Lock()
+	defer connMapMu.Unlock()
 	key := TabletToMapKey(tablet)
 	conn := &fakeConn{
 		QueryService: fakes.ErrorQueryService,
@@ -866,6 +879,8 @@ func createFixedHealthConn(tablet *topodatapb.Tablet, fixedResult *querypb.Strea
 		tablet:       tablet,
 		fixedResult:  fixedResult,
 	}
+	connMapMu.Lock()
+	defer connMapMu.Unlock()
 	connMap[key] = conn
 	return conn
 }

diff --git a/go/vt/discovery/legacy_healthcheck_flaky_test.go b/go/vt/discovery/legacy_healthcheck_flaky_test.go
@@ -26,31 +26,13 @@ import (
 	"testing"
 	"time"
 
-	"vitess.io/vitess/go/vt/log"
-
 	"golang.org/x/net/context"
-	"vitess.io/vitess/go/vt/grpcclient"
-	"vitess.io/vitess/go/vt/status"
-	"vitess.io/vitess/go/vt/topo"
-	"vitess.io/vitess/go/vt/vttablet/queryservice"
-	"vitess.io/vitess/go/vt/vttablet/tabletconn"
-
 	querypb "vitess.io/vitess/go/vt/proto/query"
 	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	"vitess.io/vitess/go/vt/status"
+	"vitess.io/vitess/go/vt/topo"
 )
 
-var connMap map[string]*fakeConn
-
-func init() {
-	tabletconn.RegisterDialer("fake_discovery", discoveryDialer)
-
-	//log error
-	if err := flag.Set("tablet_protocol", "fake_discovery"); err != nil {
-		log.Errorf("flag.Set(\"tablet_protocol\", \"fake_discovery\") failed : %v", err)
-	}
-	connMap = make(map[string]*fakeConn)
-}
-
 func testChecksum(t *testing.T, want, got int64) {
 	t.Helper()
 	if want != got {
@@ -88,18 +70,20 @@ func TestLegacyHealthCheck(t *testing.T) {
 
 	// one tablet after receiving a StreamHealthResponse
 	shr := &querypb.StreamHealthResponse{
-		Target:                              &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_MASTER},
-		Serving:                             true,
+		Target:  &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_MASTER},
+		Serving: true,
+
 		TabletExternallyReparentedTimestamp: 10,
 		RealtimeStats:                       &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.2},
 	}
 	want = &LegacyTabletStats{
-		Key:                                 "a,vt:1",
-		Tablet:                              tablet,
-		Target:                              &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_MASTER},
-		Up:                                  true,
-		Serving:                             true,
-		Stats:                               &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.2},
+		Key:     "a,vt:1",
+		Tablet:  tablet,
+		Target:  &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_MASTER},
+		Up:      true,
+		Serving: true,
+		Stats:   &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.2},
+
 		TabletExternallyReparentedTimestamp: 10,
 	}
 	input <- shr
@@ -119,12 +103,13 @@ func TestLegacyHealthCheck(t *testing.T) {
 		Cell:   "cell",
 		Target: &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_MASTER},
 		TabletsStats: LegacyTabletStatsList{{
-			Key:                                 "a,vt:1",
-			Tablet:                              tablet,
-			Target:                              &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_MASTER},
-			Up:                                  true,
-			Serving:                             true,
-			Stats:                               &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.2},
+			Key:     "a,vt:1",
+			Tablet:  tablet,
+			Target:  &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_MASTER},
+			Up:      true,
+			Serving: true,
+			Stats:   &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.2},
+
 			TabletExternallyReparentedTimestamp: 10,
 		}},
 	}}
@@ -135,10 +120,12 @@ func TestLegacyHealthCheck(t *testing.T) {
 
 	// TabletType changed, should get both old and new event
 	shr = &querypb.StreamHealthResponse{
-		Target:                              &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
-		Serving:                             true,
+		Target:  &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
+		Serving: true,
+
 		TabletExternallyReparentedTimestamp: 0,
-		RealtimeStats:                       &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.5},
+
+		RealtimeStats: &querypb.RealtimeStats{SecondsBehindMaster: 1, CpuUsage: 0.5},
 	}
 	input <- shr
 	t.Logf(`input <- {{Keyspace: "k", Shard: "s", TabletType: REPLICA}, Serving: true, TabletExternallyReparentedTimestamp: 0, {SecondsBehindMaster: 1, CpuUsage: 0.5}}`)
@@ -658,11 +645,3 @@ func newListener() *listener {
 func (l *listener) StatsUpdate(ts *LegacyTabletStats) {
 	l.output <- ts
 }
-
-func discoveryDialer(tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
-	key := TabletToMapKey(tablet)
-	if qs, ok := connMap[key]; ok {
-		return qs, nil
-	}
-	return nil, fmt.Errorf("tablet %v not found", key)
-}