vitess.io/vitess@v0.16.2/go/vt/vtgate/sandbox_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  	"sync"
    23  
    24  	"vitess.io/vitess/go/json2"
    25  	"vitess.io/vitess/go/vt/grpcclient"
    26  	"vitess.io/vitess/go/vt/key"
    27  	"vitess.io/vitess/go/vt/topo"
    28  	"vitess.io/vitess/go/vt/topo/memorytopo"
    29  	"vitess.io/vitess/go/vt/vterrors"
    30  	"vitess.io/vitess/go/vt/vttablet/queryservice"
    31  	"vitess.io/vitess/go/vt/vttablet/sandboxconn"
    32  	"vitess.io/vitess/go/vt/vttablet/tabletconn"
    33  	"vitess.io/vitess/go/vt/vttablet/tabletconntest"
    34  
    35  	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
    36  	vschemapb "vitess.io/vitess/go/vt/proto/vschema"
    37  )
    38  
    39  // sandbox_test.go provides a sandbox for unit testing VTGate.
    40  
    41  const (
    42  	KsTestSharded             = "TestExecutor"
    43  	KsTestUnsharded           = "TestUnsharded"
    44  	KsTestUnshardedServedFrom = "TestUnshardedServedFrom"
    45  	KsTestBadVSchema          = "TestXBadVSchema"
    46  )
    47  
    48  func init() {
    49  	ksToSandbox = make(map[string]*sandbox)
    50  	createSandbox(KsTestSharded)
    51  	createSandbox(KsTestUnsharded)
    52  	createSandbox(KsTestBadVSchema)
    53  	tabletconn.RegisterDialer("sandbox", sandboxDialer)
    54  	tabletconntest.SetProtocol("go.vt.vtgate.sandbox_test", "sandbox")
    55  }
    56  
    57  var sandboxMu sync.Mutex
    58  var ksToSandbox map[string]*sandbox
    59  
    60  func createSandbox(keyspace string) *sandbox {
    61  	sandboxMu.Lock()
    62  	defer sandboxMu.Unlock()
    63  	s := &sandbox{VSchema: "{}"}
    64  	s.Reset()
    65  	ksToSandbox[keyspace] = s
    66  	return s
    67  }
    68  
    69  func getSandbox(keyspace string) *sandbox {
    70  	sandboxMu.Lock()
    71  	defer sandboxMu.Unlock()
    72  	return ksToSandbox[keyspace]
    73  }
    74  
    75  func getSandboxSrvVSchema() *vschemapb.SrvVSchema {
    76  	result := &vschemapb.SrvVSchema{
    77  		Keyspaces: map[string]*vschemapb.Keyspace{},
    78  	}
    79  	sandboxMu.Lock()
    80  	defer sandboxMu.Unlock()
    81  	for keyspace, sandbox := range ksToSandbox {
    82  		var vs vschemapb.Keyspace
    83  		if err := json2.Unmarshal([]byte(sandbox.VSchema), &vs); err != nil {
    84  			panic(err)
    85  		}
    86  		result.Keyspaces[keyspace] = &vs
    87  	}
    88  	return result
    89  }
    90  
    91  type sandbox struct {
    92  	// Use sandmu to access the variables below
    93  	sandmu sync.Mutex
    94  
    95  	// SrvKeyspaceCounter tracks how often GetSrvKeyspace was called
    96  	SrvKeyspaceCounter int
    97  
    98  	// SrvKeyspaceMustFail specifies how often GetSrvKeyspace must fail before succeeding
    99  	SrvKeyspaceMustFail int
   100  
   101  	// DialCounter tracks how often sandboxDialer was called
   102  	DialCounter int
   103  
   104  	// DialMustFail specifies how often sandboxDialer must fail before succeeding
   105  	DialMustFail int
   106  
   107  	// KeyspaceServedFrom specifies the served-from keyspace for vertical resharding
   108  	KeyspaceServedFrom string
   109  
   110  	// ShardSpec specifies the sharded keyranges
   111  	ShardSpec string
   112  
   113  	// SrvKeyspaceCallback specifies the callback function in GetSrvKeyspace
   114  	SrvKeyspaceCallback func()
   115  
   116  	// VSchema specifies the vschema in JSON format.
   117  	VSchema string
   118  }
   119  
   120  // Reset cleans up sandbox internal state.
   121  func (s *sandbox) Reset() {
   122  	s.sandmu.Lock()
   123  	defer s.sandmu.Unlock()
   124  	s.SrvKeyspaceCounter = 0
   125  	s.SrvKeyspaceMustFail = 0
   126  	s.DialCounter = 0
   127  	s.DialMustFail = 0
   128  	s.KeyspaceServedFrom = ""
   129  	s.ShardSpec = DefaultShardSpec
   130  	s.SrvKeyspaceCallback = nil
   131  }
   132  
   133  // DefaultShardSpec is the default sharding scheme for testing.
   134  var DefaultShardSpec = "-20-40-60-80-a0-c0-e0-"
   135  
   136  func getAllShards(shardSpec string) ([]*topodatapb.KeyRange, error) {
   137  	shardedKrArray, err := key.ParseShardingSpec(shardSpec)
   138  	if err != nil {
   139  		return nil, err
   140  	}
   141  	return shardedKrArray, nil
   142  }
   143  
   144  func createShardedSrvKeyspace(shardSpec, servedFromKeyspace string) (*topodatapb.SrvKeyspace, error) {
   145  	shardKrArray, err := getAllShards(shardSpec)
   146  	if err != nil {
   147  		return nil, err
   148  	}
   149  	shards := make([]*topodatapb.ShardReference, 0, len(shardKrArray))
   150  	for i := 0; i < len(shardKrArray); i++ {
   151  		shard := &topodatapb.ShardReference{
   152  			Name:     key.KeyRangeString(shardKrArray[i]),
   153  			KeyRange: shardKrArray[i],
   154  		}
   155  		shards = append(shards, shard)
   156  	}
   157  	shardedSrvKeyspace := &topodatapb.SrvKeyspace{
   158  		Partitions: []*topodatapb.SrvKeyspace_KeyspacePartition{
   159  			{
   160  				ServedType:      topodatapb.TabletType_PRIMARY,
   161  				ShardReferences: shards,
   162  			},
   163  			{
   164  				ServedType:      topodatapb.TabletType_REPLICA,
   165  				ShardReferences: shards,
   166  			},
   167  			{
   168  				ServedType:      topodatapb.TabletType_RDONLY,
   169  				ShardReferences: shards,
   170  			},
   171  		},
   172  	}
   173  	if servedFromKeyspace != "" {
   174  		shardedSrvKeyspace.ServedFrom = []*topodatapb.SrvKeyspace_ServedFrom{
   175  			{
   176  				TabletType: topodatapb.TabletType_RDONLY,
   177  				Keyspace:   servedFromKeyspace,
   178  			},
   179  			{
   180  				TabletType: topodatapb.TabletType_PRIMARY,
   181  				Keyspace:   servedFromKeyspace,
   182  			},
   183  		}
   184  	}
   185  	return shardedSrvKeyspace, nil
   186  }
   187  
   188  func createUnshardedKeyspace() (*topodatapb.SrvKeyspace, error) {
   189  	shard := &topodatapb.ShardReference{
   190  		Name: "0",
   191  	}
   192  
   193  	unshardedSrvKeyspace := &topodatapb.SrvKeyspace{
   194  		Partitions: []*topodatapb.SrvKeyspace_KeyspacePartition{
   195  			{
   196  				ServedType:      topodatapb.TabletType_PRIMARY,
   197  				ShardReferences: []*topodatapb.ShardReference{shard},
   198  			},
   199  			{
   200  				ServedType:      topodatapb.TabletType_REPLICA,
   201  				ShardReferences: []*topodatapb.ShardReference{shard},
   202  			},
   203  			{
   204  				ServedType:      topodatapb.TabletType_RDONLY,
   205  				ShardReferences: []*topodatapb.ShardReference{shard},
   206  			},
   207  		},
   208  	}
   209  	return unshardedSrvKeyspace, nil
   210  }
   211  
   212  // sandboxTopo satisfies the srvtopo.Server interface
   213  type sandboxTopo struct {
   214  	topoServer *topo.Server
   215  }
   216  
   217  // newSandboxForCells creates a new topo with a backing memory topo for
   218  // the given cells.
   219  //
   220  // when this version is used, WatchSrvVSchema can properly simulate watches
   221  func newSandboxForCells(cells []string) *sandboxTopo {
   222  	return &sandboxTopo{
   223  		topoServer: memorytopo.NewServer(cells...),
   224  	}
   225  }
   226  
   227  // GetTopoServer is part of the srvtopo.Server interface
   228  func (sct *sandboxTopo) GetTopoServer() (*topo.Server, error) {
   229  	return sct.topoServer, nil
   230  }
   231  
   232  // GetSrvKeyspaceNames is part of the srvtopo.Server interface.
   233  func (sct *sandboxTopo) GetSrvKeyspaceNames(ctx context.Context, cell string, staleOK bool) ([]string, error) {
   234  	sandboxMu.Lock()
   235  	defer sandboxMu.Unlock()
   236  	keyspaces := make([]string, 0, 1)
   237  	for k := range ksToSandbox {
   238  		keyspaces = append(keyspaces, k)
   239  	}
   240  	return keyspaces, nil
   241  }
   242  
   243  // GetSrvKeyspace is part of the srvtopo.Server interface.
   244  func (sct *sandboxTopo) GetSrvKeyspace(ctx context.Context, cell, keyspace string) (*topodatapb.SrvKeyspace, error) {
   245  	sand := getSandbox(keyspace)
   246  	if sand == nil {
   247  		return nil, fmt.Errorf("topo error GetSrvKeyspace")
   248  	}
   249  	sand.sandmu.Lock()
   250  	defer sand.sandmu.Unlock()
   251  	if sand.SrvKeyspaceCallback != nil {
   252  		sand.SrvKeyspaceCallback()
   253  	}
   254  	sand.SrvKeyspaceCounter++
   255  	if sand.SrvKeyspaceMustFail > 0 {
   256  		sand.SrvKeyspaceMustFail--
   257  		return nil, fmt.Errorf("topo error GetSrvKeyspace")
   258  	}
   259  	switch keyspace {
   260  	case KsTestUnshardedServedFrom:
   261  		servedFromKeyspace, err := createUnshardedKeyspace()
   262  		if err != nil {
   263  			return nil, err
   264  		}
   265  		servedFromKeyspace.ServedFrom = []*topodatapb.SrvKeyspace_ServedFrom{
   266  			{
   267  				TabletType: topodatapb.TabletType_RDONLY,
   268  				Keyspace:   KsTestUnsharded,
   269  			},
   270  			{
   271  				TabletType: topodatapb.TabletType_PRIMARY,
   272  				Keyspace:   KsTestUnsharded,
   273  			},
   274  		}
   275  		return servedFromKeyspace, nil
   276  	case KsTestUnsharded:
   277  		return createUnshardedKeyspace()
   278  	}
   279  
   280  	return createShardedSrvKeyspace(sand.ShardSpec, sand.KeyspaceServedFrom)
   281  }
   282  
   283  func (sct *sandboxTopo) WatchSrvKeyspace(ctx context.Context, cell, keyspace string, callback func(*topodatapb.SrvKeyspace, error) bool) {
   284  	// panic("not supported: WatchSrvKeyspace")
   285  }
   286  
   287  // WatchSrvVSchema is part of the srvtopo.Server interface.
   288  //
   289  // If the sandbox was created with a backing topo service, piggy back on it
   290  // to properly simulate watches, otherwise just immediately call back the
   291  // caller.
   292  func (sct *sandboxTopo) WatchSrvVSchema(ctx context.Context, cell string, callback func(*vschemapb.SrvVSchema, error) bool) {
   293  	srvVSchema := getSandboxSrvVSchema()
   294  
   295  	if sct.topoServer == nil {
   296  		callback(srvVSchema, nil)
   297  		return
   298  	}
   299  
   300  	sct.topoServer.UpdateSrvVSchema(ctx, cell, srvVSchema)
   301  	current, updateChan, _ := sct.topoServer.WatchSrvVSchema(ctx, cell)
   302  	if !callback(current.Value, nil) {
   303  		panic("sandboxTopo callback returned false")
   304  	}
   305  	go func() {
   306  		for {
   307  			update := <-updateChan
   308  			if !callback(update.Value, update.Err) {
   309  				panic("sandboxTopo callback returned false")
   310  			}
   311  		}
   312  	}()
   313  }
   314  
   315  func sandboxDialer(tablet *topodatapb.Tablet, failFast grpcclient.FailFast) (queryservice.QueryService, error) {
   316  	sand := getSandbox(tablet.Keyspace)
   317  	sand.sandmu.Lock()
   318  	defer sand.sandmu.Unlock()
   319  	sand.DialCounter++
   320  	if sand.DialMustFail > 0 {
   321  		sand.DialMustFail--
   322  		return nil, vterrors.VT14001()
   323  	}
   324  	sbc := sandboxconn.NewSandboxConn(tablet)
   325  	return sbc, nil
   326  }