VStream API: allow cells to be specified for picking source tablets t…

…o stream from (vitessio#10294)

* Add parameter Cells in vstream api flags, for tablet picker to look for candidates in

Signed-off-by: Rohit Nayak <[email protected]>

* Add e2e test for the new Cells vstream flag

Signed-off-by: Rohit Nayak <[email protected]>

* Self-review

Signed-off-by: Rohit Nayak <[email protected]>

* Address review comments

Fix gofmt

Signed-off-by: Rohit Nayak <[email protected]>
Signed-off-by: Vilius Okockis <[email protected]>

go/test/endtoend/vreplication/helper.go

@@ -5,11 +5,14 @@ import (
 	"fmt"
 	"io/ioutil"
 	"net/http"
+	"os/exec"
 	"regexp"
 	"strconv"
 	"strings"
 	"testing"
 
+	"vitess.io/vitess/go/vt/log"
+
 	"github.com/buger/jsonparser"
 	"github.com/stretchr/testify/require"
 
@@ -98,24 +101,29 @@ func validateThatQueryExecutesOnTablet(t *testing.T, conn *mysql.Conn, tablet *c
 	return newCount == count+1
 }
 
-func getQueryCount(url string, query string) int {
-	var headings, row []string
-	var rows [][]string
+func getHTTPBody(url string) string {
 	resp, err := http.Get(url)
 	if err != nil {
-		fmt.Printf("http Get returns %+v\n", err)
-		return 0
+		log.Infof("http Get returns %+v", err)
+		return ""
 	}
 	if resp.StatusCode != 200 {
-		fmt.Printf("http Get returns status %d\n", resp.StatusCode)
-		return 0
+		log.Infof("http Get returns status %d", resp.StatusCode)
+		return ""
 	}
 	respByte, _ := ioutil.ReadAll(resp.Body)
 	defer resp.Body.Close()
 	body := string(respByte)
+	return body
+}
+
+func getQueryCount(url string, query string) int {
+	var headings, row []string
+	var rows [][]string
+	body := getHTTPBody(url)
 	doc, err := goquery.NewDocumentFromReader(strings.NewReader(body))
 	if err != nil {
-		fmt.Printf("goquery parsing returns %+v\n", err)
+		log.Infof("goquery parsing returns %+v\n", err)
 		return 0
 	}
 
@@ -143,7 +151,7 @@ func getQueryCount(url string, query string) int {
 		})
 	})
 	if queryIndex == -1 || countIndex == -1 {
-		fmt.Println("Queryz response is incorrect")
+		log.Infof("Queryz response is incorrect")
 		return 0
 	}
 	for _, row := range rows {
@@ -267,3 +275,19 @@ func printRoutingRules(t *testing.T, vc *VitessCluster, msg string) error {
 	fmt.Printf("Routing Rules::%s:\n%s\n", msg, output)
 	return nil
 }
+
+func osExec(t *testing.T, command string, args []string) (string, error) {
+	cmd := exec.Command(command, args...)
+	output, err := cmd.CombinedOutput()
+	return string(output), err
+}
+
+func getDebugVar(t *testing.T, port int, varPath []string) (string, error) {
+	var val []byte
+	var err error
+	url := fmt.Sprintf("http://localhost:%d/debug/vars", port)
+	body := getHTTPBody(url)
+	val, _, _, err = jsonparser.Get([]byte(body), varPath...)
+	require.NoError(t, err)
+	return string(val), nil
+}

go/test/endtoend/vreplication/vreplication_test.go

@@ -17,15 +17,22 @@ limitations under the License.
 package vreplication
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
+	"io"
 	"io/ioutil"
 	"net/http"
 	"strings"
+	"sync"
 	"testing"
 	"time"
 
 	"vitess.io/vitess/go/vt/log"
+	binlogdatapb "vitess.io/vitess/go/vt/proto/binlogdata"
+	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
+	vtgatepb "vitess.io/vitess/go/vt/proto/vtgate"
+	"vitess.io/vitess/go/vt/vtgate/vtgateconn"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
@@ -142,7 +149,7 @@ func TestBasicVreplicationWorkflow(t *testing.T) {
 
 func TestMultiCellVreplicationWorkflow(t *testing.T) {
 	cells := []string{"zone1", "zone2"}
-	allCellNames = "zone1,zone2"
+	allCellNames = strings.Join(cells, ",")
 
 	vc = NewVitessCluster(t, "TestMultiCellVreplicationWorkflow", cells, mainClusterConfig)
 	require.NotNil(t, vc)
@@ -165,6 +172,86 @@ func TestMultiCellVreplicationWorkflow(t *testing.T) {
 	verifyClusterHealth(t, vc)
 	insertInitialData(t)
 	shardCustomer(t, true, []*Cell{cell1, cell2}, cell2.Name, true)
+
+	// we tag along this test so as not to create the overhead of creating another cluster
+	testVStreamCellFlag(t)
+}
+
+func testVStreamCellFlag(t *testing.T) {
+	vgtid := &binlogdatapb.VGtid{
+		ShardGtids: []*binlogdatapb.ShardGtid{{
+			Keyspace: "product",
+			Shard:    "0",
+			Gtid:     "",
+		}}}
+	filter := &binlogdatapb.Filter{
+		Rules: []*binlogdatapb.Rule{{
+			Match:  "product",
+			Filter: "select * from product",
+		}},
+	}
+	ctx := context.Background()
+
+	type vstreamTestCase struct {
+		cells       string
+		expectError bool
+	}
+	nonExistingCell := "zone7"
+	vstreamTestCases := []vstreamTestCase{
+		{"zone1,zone2", false},
+		{nonExistingCell, true},
+		{"", false},
+	}
+
+	for _, tc := range vstreamTestCases {
+		t.Run("VStreamCellsFlag/"+tc.cells, func(t *testing.T) {
+			conn, err := vtgateconn.Dial(ctx, fmt.Sprintf("localhost:%d", vc.ClusterConfig.vtgateGrpcPort))
+			require.NoError(t, err)
+			defer conn.Close()
+
+			flags := &vtgatepb.VStreamFlags{}
+			if tc.cells != "" {
+				flags.Cells = tc.cells
+			}
+
+			ctx2, cancel := context.WithTimeout(ctx, 30*time.Second)
+			reader, err := conn.VStream(ctx2, topodatapb.TabletType_REPLICA, vgtid, filter, flags)
+			require.NoError(t, err)
+
+			rowsReceived := false
+			var wg sync.WaitGroup
+			wg.Add(1)
+			go func() {
+				defer wg.Done()
+				defer cancel()
+
+				events, err := reader.Recv()
+				switch err {
+				case nil:
+					if len(events) > 0 {
+						log.Infof("received %d events", len(events))
+						rowsReceived = true
+					}
+				case io.EOF:
+					log.Infof("stream ended without data")
+				default:
+					log.Infof("%s:: remote error: %v", time.Now(), err)
+				}
+			}()
+			wg.Wait()
+
+			if tc.expectError {
+				require.False(t, rowsReceived)
+
+				// if no tablet was found the tablet picker adds a key which includes the cell name to the vtgate TabletPickerNoTabletFoundErrorCount stat
+				pickerErrorStat, err := getDebugVar(t, vc.ClusterConfig.vtgatePort, []string{"TabletPickerNoTabletFoundErrorCount"})
+				require.NoError(t, err)
+				require.Contains(t, pickerErrorStat, nonExistingCell)
+			} else {
+				require.True(t, rowsReceived)
+			}
+		})
+	}
 }
 
 func TestCellAliasVreplicationWorkflow(t *testing.T) {

go/vt/discovery/healthcheck.go

@@ -164,7 +164,7 @@ type TabletRecorder interface {
 type keyspaceShardTabletType string
 type tabletAliasString string
 
-//HealthCheck declares what the TabletGateway needs from the HealthCheck
+// HealthCheck declares what the TabletGateway needs from the HealthCheck
 type HealthCheck interface {
 	// CacheStatus returns a displayable version of the health check cache.
 	CacheStatus() TabletsCacheStatusList
@@ -245,18 +245,27 @@ type HealthCheckImpl struct {
 // NewHealthCheck creates a new HealthCheck object.
 // Parameters:
 // retryDelay.
-//   The duration to wait before retrying to connect (e.g. after a failed connection
-//   attempt).
+//
+//	The duration to wait before retrying to connect (e.g. after a failed connection
+//	attempt).
+//
 // healthCheckTimeout.
-//   The duration for which we consider a health check response to be 'fresh'. If we don't get
-//   a health check response from a tablet for more than this duration, we consider the tablet
-//   not healthy.
+//
+//	The duration for which we consider a health check response to be 'fresh'. If we don't get
+//	a health check response from a tablet for more than this duration, we consider the tablet
+//	not healthy.
+//
 // topoServer.
-//   The topology server that this healthcheck object can use to retrieve cell or tablet information
+//
+//	The topology server that this healthcheck object can use to retrieve cell or tablet information
+//
 // localCell.
-//   The localCell for this healthcheck
+//
+//	The localCell for this healthcheck
+//
 // callback.
-//   A function to call when there is a master change. Used to notify vtgate's buffer to stop buffering.
+//
+//	A function to call when there is a master change. Used to notify vtgate's buffer to stop buffering.
 func NewHealthCheck(ctx context.Context, retryDelay, healthCheckTimeout time.Duration, topoServer *topo.Server, localCell, cellsToWatch string) *HealthCheckImpl {
 	log.Infof("loading tablets for cells: %v", cellsToWatch)
 

go/vt/discovery/legacy_healthcheck.go

@@ -378,13 +378,13 @@ func NewLegacyDefaultHealthCheck() LegacyHealthCheck {
 
 // NewLegacyHealthCheck creates a new LegacyHealthCheck object.
 // Parameters:
-// retryDelay.
-//   The duration to wait before retrying to connect (e.g. after a failed connection
-//   attempt).
-// healthCheckTimeout.
-//   The duration for which we consider a health check response to be 'fresh'. If we don't get
-//   a health check response from a tablet for more than this duration, we consider the tablet
-//   not healthy.
+//   - retryDelay
+//     The duration to wait before retrying to connect (e.g. after a failed connection
+//     attempt).
+//   - healthCheckTimeout
+//     The duration for which we consider a health check response to be 'fresh'. If we don't get
+//     a health check response from a tablet for more than this duration, we consider the tablet
+//     not healthy.
 func NewLegacyHealthCheck(retryDelay, healthCheckTimeout time.Duration) LegacyHealthCheck {
 	hc := &LegacyHealthCheckImpl{
 		addrToHealth:       make(map[string]*legacyTabletHealth),

go/vt/discovery/legacy_replicationlag.go

@@ -54,10 +54,10 @@ func LegacyIsReplicationLagVeryHigh(tabletStats *LegacyTabletStats) bool {
 // lags of (30m, 35m, 40m, 45m) return all.
 //
 // One thing to know about this code: vttablet also has a couple flags that impact the logic here:
-// * unhealthy_threshold: if replication lag is higher than this, a tablet will be reported as unhealthy.
-//   The default for this is 2h, same as the discovery_high_replication_lag_minimum_serving here.
-// * degraded_threshold: this is only used by vttablet for display. It should match
-//   discovery_low_replication_lag here, so the vttablet status display matches what vtgate will do of it.
+//   - unhealthy_threshold: if replication lag is higher than this, a tablet will be reported as unhealthy.
+//     The default for this is 2h, same as the discovery_high_replication_lag_minimum_serving here.
+//   - degraded_threshold: this is only used by vttablet for display. It should match
+//     discovery_low_replication_lag here, so the vttablet status display matches what vtgate will do of it.
 func FilterLegacyStatsByReplicationLag(tabletStatsList []*LegacyTabletStats) []*LegacyTabletStats {
 	if !*legacyReplicationLagAlgorithm {
 		return filterLegacyStatsByLag(tabletStatsList)

go/vt/discovery/replicationlag.go

@@ -94,10 +94,10 @@ func IsReplicationLagVeryHigh(tabletHealth *TabletHealth) bool {
 // lags of (30m, 35m, 40m, 45m) return all.
 //
 // One thing to know about this code: vttablet also has a couple flags that impact the logic here:
-// * unhealthy_threshold: if replication lag is higher than this, a tablet will be reported as unhealthy.
-//   The default for this is 2h, same as the discovery_high_replication_lag_minimum_serving here.
-// * degraded_threshold: this is only used by vttablet for display. It should match
-//   discovery_low_replication_lag here, so the vttablet status display matches what vtgate will do of it.
+//   - unhealthy_threshold: if replication lag is higher than this, a tablet will be reported as unhealthy.
+//     The default for this is 2h, same as the discovery_high_replication_lag_minimum_serving here.
+//   - degraded_threshold: this is only used by vttablet for display. It should match
+//     discovery_low_replication_lag here, so the vttablet status display matches what vtgate will do of it.
 func FilterStatsByReplicationLag(tabletHealthList []*TabletHealth) []*TabletHealth {
 	if !*legacyReplicationLagAlgorithm {
 		return filterStatsByLag(tabletHealthList)

go/vt/discovery/tablet_picker.go

@@ -196,6 +196,8 @@ func (tp *TabletPicker) GetMatchingTablets(ctx context.Context) []*topo.TabletIn
 				// if we get an error, either cellAlias doesn't exist or it isn't a cell alias at all. Ignore and continue
 				if err == nil {
 					actualCells = append(actualCells, alias.Cells...)
+				} else {
+					log.Infof("Unable to resolve cell %s, ignoring", cell)
 				}
 			} else {
 				// valid cell, add it to our list