vitess.io/vitess@v0.16.2/go/vt/vtgate/tabletgateway_test.go (about)

     1  /*
     2  Copyright 2019 The Vitess Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package vtgate
    18  
    19  import (
    20  	"context"
    21  	"fmt"
    22  	"strings"
    23  	"testing"
    24  
    25  	"github.com/stretchr/testify/assert"
    26  	"github.com/stretchr/testify/require"
    27  
    28  	"vitess.io/vitess/go/sqltypes"
    29  	"vitess.io/vitess/go/vt/discovery"
    30  	querypb "vitess.io/vitess/go/vt/proto/query"
    31  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    32  	vtrpcpb "vitess.io/vitess/go/vt/proto/vtrpc"
    33  	"vitess.io/vitess/go/vt/topo"
    34  	"vitess.io/vitess/go/vt/vterrors"
    35  )
    36  
    37  func TestTabletGatewayExecute(t *testing.T) {
    38  	testTabletGatewayGeneric(t, func(tg *TabletGateway, target *querypb.Target) error {
    39  		_, err := tg.Execute(context.Background(), target, "query", nil, 0, 0, nil)
    40  		return err
    41  	})
    42  	testTabletGatewayTransact(t, func(tg *TabletGateway, target *querypb.Target) error {
    43  		_, err := tg.Execute(context.Background(), target, "query", nil, 1, 0, nil)
    44  		return err
    45  	})
    46  }
    47  
    48  func TestTabletGatewayExecuteStream(t *testing.T) {
    49  	testTabletGatewayGeneric(t, func(tg *TabletGateway, target *querypb.Target) error {
    50  		err := tg.StreamExecute(context.Background(), target, "query", nil, 0, 0, nil, func(qr *sqltypes.Result) error {
    51  			return nil
    52  		})
    53  		return err
    54  	})
    55  }
    56  
    57  func TestTabletGatewayBegin(t *testing.T) {
    58  	testTabletGatewayGeneric(t, func(tg *TabletGateway, target *querypb.Target) error {
    59  		_, err := tg.Begin(context.Background(), target, nil)
    60  		return err
    61  	})
    62  }
    63  
    64  func TestTabletGatewayCommit(t *testing.T) {
    65  	testTabletGatewayTransact(t, func(tg *TabletGateway, target *querypb.Target) error {
    66  		_, err := tg.Commit(context.Background(), target, 1)
    67  		return err
    68  	})
    69  }
    70  
    71  func TestTabletGatewayRollback(t *testing.T) {
    72  	testTabletGatewayTransact(t, func(tg *TabletGateway, target *querypb.Target) error {
    73  		_, err := tg.Rollback(context.Background(), target, 1)
    74  		return err
    75  	})
    76  }
    77  
    78  func TestTabletGatewayBeginExecute(t *testing.T) {
    79  	testTabletGatewayGeneric(t, func(tg *TabletGateway, target *querypb.Target) error {
    80  		_, _, err := tg.BeginExecute(context.Background(), target, nil, "query", nil, 0, nil)
    81  		return err
    82  	})
    83  }
    84  
    85  func TestTabletGatewayShuffleTablets(t *testing.T) {
    86  	hc := discovery.NewFakeHealthCheck(nil)
    87  	tg := NewTabletGateway(context.Background(), hc, nil, "local")
    88  
    89  	ts1 := &discovery.TabletHealth{
    90  		Tablet:  topo.NewTablet(1, "cell1", "host1"),
    91  		Target:  &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
    92  		Serving: true,
    93  		Stats:   &querypb.RealtimeStats{ReplicationLagSeconds: 1, CpuUsage: 0.2},
    94  	}
    95  
    96  	ts2 := &discovery.TabletHealth{
    97  		Tablet:  topo.NewTablet(2, "cell1", "host2"),
    98  		Target:  &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
    99  		Serving: true,
   100  		Stats:   &querypb.RealtimeStats{ReplicationLagSeconds: 1, CpuUsage: 0.2},
   101  	}
   102  
   103  	ts3 := &discovery.TabletHealth{
   104  		Tablet:  topo.NewTablet(3, "cell2", "host3"),
   105  		Target:  &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
   106  		Serving: true,
   107  		Stats:   &querypb.RealtimeStats{ReplicationLagSeconds: 1, CpuUsage: 0.2},
   108  	}
   109  
   110  	ts4 := &discovery.TabletHealth{
   111  		Tablet:  topo.NewTablet(4, "cell2", "host4"),
   112  		Target:  &querypb.Target{Keyspace: "k", Shard: "s", TabletType: topodatapb.TabletType_REPLICA},
   113  		Serving: true,
   114  		Stats:   &querypb.RealtimeStats{ReplicationLagSeconds: 1, CpuUsage: 0.2},
   115  	}
   116  
   117  	sameCellTablets := []*discovery.TabletHealth{ts1, ts2}
   118  	diffCellTablets := []*discovery.TabletHealth{ts3, ts4}
   119  	mixedTablets := []*discovery.TabletHealth{ts1, ts2, ts3, ts4}
   120  	// repeat shuffling 10 times and every time the same cell tablets should be in the front
   121  	for i := 0; i < 10; i++ {
   122  		tg.shuffleTablets("cell1", sameCellTablets)
   123  		assert.Len(t, sameCellTablets, 2, "Wrong number of TabletHealth")
   124  		assert.Equal(t, sameCellTablets[0].Tablet.Alias.Cell, "cell1", "Wrong tablet cell")
   125  		assert.Equal(t, sameCellTablets[1].Tablet.Alias.Cell, "cell1", "Wrong tablet cell")
   126  
   127  		tg.shuffleTablets("cell1", diffCellTablets)
   128  		assert.Len(t, diffCellTablets, 2, "should shuffle in only diff cell tablets")
   129  		assert.Contains(t, diffCellTablets, ts3, "diffCellTablets should contain %v", ts3)
   130  		assert.Contains(t, diffCellTablets, ts4, "diffCellTablets should contain %v", ts4)
   131  
   132  		tg.shuffleTablets("cell1", mixedTablets)
   133  		assert.Len(t, mixedTablets, 4, "should have 4 tablets, got %+v", mixedTablets)
   134  
   135  		assert.Contains(t, mixedTablets[0:2], ts1, "should have same cell tablets in the front, got %+v", mixedTablets)
   136  		assert.Contains(t, mixedTablets[0:2], ts2, "should have same cell tablets in the front, got %+v", mixedTablets)
   137  
   138  		assert.Contains(t, mixedTablets[2:4], ts3, "should have diff cell tablets in the rear, got %+v", mixedTablets)
   139  		assert.Contains(t, mixedTablets[2:4], ts4, "should have diff cell tablets in the rear, got %+v", mixedTablets)
   140  	}
   141  }
   142  
   143  func TestTabletGatewayReplicaTransactionError(t *testing.T) {
   144  	keyspace := "ks"
   145  	shard := "0"
   146  	// transactions on REPLICA are not allowed from tabletgateway
   147  	// they have to be executed directly on tabletserver
   148  	tabletType := topodatapb.TabletType_REPLICA
   149  	host := "1.1.1.1"
   150  	port := int32(1001)
   151  	target := &querypb.Target{
   152  		Keyspace:   keyspace,
   153  		Shard:      shard,
   154  		TabletType: tabletType,
   155  	}
   156  	hc := discovery.NewFakeHealthCheck(nil)
   157  	tg := NewTabletGateway(context.Background(), hc, nil, "cell")
   158  
   159  	_ = hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil)
   160  	_, err := tg.Execute(context.Background(), target, "query", nil, 1, 0, nil)
   161  	verifyContainsError(t, err, "query service can only be used for non-transactional queries on replicas", vtrpcpb.Code_INTERNAL)
   162  }
   163  
   164  func testTabletGatewayGeneric(t *testing.T, f func(tg *TabletGateway, target *querypb.Target) error) {
   165  	t.Helper()
   166  	keyspace := "ks"
   167  	shard := "0"
   168  	tabletType := topodatapb.TabletType_REPLICA
   169  	host := "1.1.1.1"
   170  	port := int32(1001)
   171  	target := &querypb.Target{
   172  		Keyspace:   keyspace,
   173  		Shard:      shard,
   174  		TabletType: tabletType,
   175  	}
   176  	hc := discovery.NewFakeHealthCheck(nil)
   177  	tg := NewTabletGateway(context.Background(), hc, nil, "cell")
   178  
   179  	// no tablet
   180  	want := []string{"target: ks.0.replica", `no healthy tablet available for 'keyspace:"ks" shard:"0" tablet_type:REPLICA`}
   181  	err := f(tg, target)
   182  	verifyShardErrors(t, err, want, vtrpcpb.Code_UNAVAILABLE)
   183  
   184  	// tablet with error
   185  	hc.Reset()
   186  	hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, false, 10, fmt.Errorf("no connection"))
   187  	err = f(tg, target)
   188  	verifyShardErrors(t, err, want, vtrpcpb.Code_UNAVAILABLE)
   189  
   190  	// tablet without connection
   191  	hc.Reset()
   192  	_ = hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, false, 10, nil).Tablet()
   193  	err = f(tg, target)
   194  	verifyShardErrors(t, err, want, vtrpcpb.Code_UNAVAILABLE)
   195  
   196  	// retry error
   197  	hc.Reset()
   198  	sc1 := hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil)
   199  	sc2 := hc.AddTestTablet("cell", host, port+1, keyspace, shard, tabletType, true, 10, nil)
   200  	sc1.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 1
   201  	sc2.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 1
   202  
   203  	err = f(tg, target)
   204  	verifyContainsError(t, err, "target: ks.0.replica", vtrpcpb.Code_FAILED_PRECONDITION)
   205  
   206  	// fatal error
   207  	hc.Reset()
   208  	sc1 = hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil)
   209  	sc2 = hc.AddTestTablet("cell", host, port+1, keyspace, shard, tabletType, true, 10, nil)
   210  	sc1.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 1
   211  	sc2.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 1
   212  	err = f(tg, target)
   213  	verifyContainsError(t, err, "target: ks.0.replica", vtrpcpb.Code_FAILED_PRECONDITION)
   214  
   215  	// server error - no retry
   216  	hc.Reset()
   217  	sc1 = hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil)
   218  	sc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1
   219  	err = f(tg, target)
   220  	assert.Equal(t, vtrpcpb.Code_INVALID_ARGUMENT, vterrors.Code(err))
   221  
   222  	// no failure
   223  	hc.Reset()
   224  	hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil)
   225  	err = f(tg, target)
   226  	assert.NoError(t, err)
   227  }
   228  
   229  func testTabletGatewayTransact(t *testing.T, f func(tg *TabletGateway, target *querypb.Target) error) {
   230  	t.Helper()
   231  	keyspace := "ks"
   232  	shard := "0"
   233  	// test with PRIMARY because replica transactions don't use gateway's queryservice
   234  	// they are executed directly on tabletserver
   235  	tabletType := topodatapb.TabletType_PRIMARY
   236  	host := "1.1.1.1"
   237  	port := int32(1001)
   238  	target := &querypb.Target{
   239  		Keyspace:   keyspace,
   240  		Shard:      shard,
   241  		TabletType: tabletType,
   242  	}
   243  	hc := discovery.NewFakeHealthCheck(nil)
   244  	tg := NewTabletGateway(context.Background(), hc, nil, "cell")
   245  
   246  	// retry error - no retry
   247  	sc1 := hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil)
   248  	sc2 := hc.AddTestTablet("cell", host, port+1, keyspace, shard, tabletType, true, 10, nil)
   249  	sc1.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 1
   250  	sc2.MustFailCodes[vtrpcpb.Code_FAILED_PRECONDITION] = 1
   251  
   252  	err := f(tg, target)
   253  	verifyContainsError(t, err, "target: ks.0.primary", vtrpcpb.Code_FAILED_PRECONDITION)
   254  
   255  	// server error - no retry
   256  	hc.Reset()
   257  	sc1 = hc.AddTestTablet("cell", host, port, keyspace, shard, tabletType, true, 10, nil)
   258  	sc1.MustFailCodes[vtrpcpb.Code_INVALID_ARGUMENT] = 1
   259  	err = f(tg, target)
   260  	verifyContainsError(t, err, "target: ks.0.primary", vtrpcpb.Code_INVALID_ARGUMENT)
   261  }
   262  
   263  func verifyContainsError(t *testing.T, err error, wantErr string, wantCode vtrpcpb.Code) {
   264  	require.Error(t, err)
   265  	if !strings.Contains(err.Error(), wantErr) {
   266  		assert.Failf(t, "", "wanted error: \n%s\n, got error: \n%v\n", wantErr, err)
   267  	}
   268  	if code := vterrors.Code(err); code != wantCode {
   269  		assert.Failf(t, "", "wanted error code: %v, got: %v", wantCode, code)
   270  	}
   271  }
   272  
   273  func verifyShardErrors(t *testing.T, err error, wantErrors []string, wantCode vtrpcpb.Code) {
   274  	require.Error(t, err)
   275  	for _, wantErr := range wantErrors {
   276  		require.Contains(t, err.Error(), wantErr, "wanted error: \n%s\n, got error: \n%v\n", wantErr, err)
   277  	}
   278  	require.Equal(t, vterrors.Code(err), wantCode, "wanted error code: %s, got: %v", wantCode, vterrors.Code(err))
   279  }