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 }