Golang tabletconn.RegisterDialer函数代码示例

func init() {
	sandboxMap = make(map[string]*sandbox)
	createSandbox(TEST_SHARDED)
	createSandbox(TEST_UNSHARDED)
	tabletconn.RegisterDialer("sandbox", sandboxDialer)
	flag.Set("tablet_protocol", "sandbox")
}

func init() {
	sandboxMap = make(map[string]*sandbox)
	createSandbox(KsTestSharded)
	createSandbox(KsTestUnsharded)
	tabletconn.RegisterDialer("sandbox", sandboxDialer)
	flag.Set("tablet_protocol", "sandbox")
}

// createSourceTablet is a helper method to create the source tablet
// in the given keyspace/shard.
func createSourceTablet(t *testing.T, name string, ts topo.Server, keyspace, shard string) {
	vshard, kr, err := topo.ValidateShardName(shard)
	if err != nil {
		t.Fatalf("ValidateShardName(%v) failed: %v", shard, err)
	}

	ctx := context.Background()
	tablet := &topodatapb.Tablet{
		Alias: &topodatapb.TabletAlias{
			Cell: "cell1",
			Uid:  100,
		},
		Keyspace: keyspace,
		Shard:    vshard,
		Type:     topodatapb.TabletType_REPLICA,
		KeyRange: kr,
		PortMap: map[string]int32{
			"vt": 80,
		},
	}
	if err := ts.CreateTablet(ctx, tablet); err != nil {
		t.Fatalf("CreateTablet failed: %v", err)
	}

	// register a tablet conn dialer that will return the instance
	// we want
	tabletconn.RegisterDialer(name, func(tablet *topodatapb.Tablet, timeout time.Duration) (tabletconn.TabletConn, error) {
		return &fakeTabletConn{
			tablet: tablet,
		}, nil
	})
	flag.Set("tablet_protocol", name)
}

// TestNewRestartableResultReader tests the correct error handling e.g.
// if the connection to a tablet fails due to a canceled context.
func TestNewRestartableResultReader(t *testing.T) {
	wantErr := errors.New("restartable_result_reader_test.go: context canceled")

	tabletconn.RegisterDialer("fake_dialer", func(tablet *topodatapb.Tablet, timeout time.Duration) (tabletconn.TabletConn, error) {
		return nil, wantErr
	})
	protocol := flag.CommandLine.Lookup("tablet_protocol").Value.String()
	flag.Set("tablet_protocol", "fake_dialer")
	// Restore the previous flag value after the test.
	defer flag.Set("tablet_protocol", protocol)

	// Create dependencies e.g. a "singleTabletProvider" instance.
	ts := zktestserver.New(t, []string{"cell1"})
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient())
	alias := &topodatapb.TabletAlias{
		Cell: "cell1",
		Uid:  1,
	}
	tablet := &topodatapb.Tablet{
		Keyspace: "ks1",
		Shard:    "-80",
		Alias:    alias,
	}
	ctx := context.Background()
	if err := ts.CreateTablet(ctx, tablet); err != nil {
		t.Fatalf("CreateTablet failed: %v", err)
	}
	tp := newSingleTabletProvider(ctx, ts, alias)

	_, err := NewRestartableResultReader(ctx, wr.Logger(), tp, nil /* td */, chunk{}, false)
	if err == nil || !strings.Contains(err.Error(), wantErr.Error()) {
		t.Fatalf("NewRestartableResultReader() should have failed because the context is canceled: %v", err)
	}
}

// TestSuite executes a set of tests on the provided gateway. The provided
// gateway needs to be configured with one established connection for
// tabletconntest.TestTarget.{Keyspace, Shard, TabletType} to the
// provided tabletconntest.FakeQueryService.
func TestSuite(t *testing.T, name string, g gateway.Gateway, f *tabletconntest.FakeQueryService) {

	protocolName := "gateway-test-" + name

	tabletconn.RegisterDialer(protocolName, func(tablet *topodatapb.Tablet, timeout time.Duration) (tabletconn.TabletConn, error) {
		return &gatewayAdapter{Gateway: g}, nil
	})

	tabletconntest.TestSuite(t, protocolName, &topodatapb.Tablet{
		Keyspace: tabletconntest.TestTarget.Keyspace,
		Shard:    tabletconntest.TestTarget.Shard,
		Type:     tabletconntest.TestTarget.TabletType,
	}, f)
}

// createSourceTablet is a helper method to create the source tablet
// in the given keyspace/shard.
func createSourceTablet(t *testing.T, name string, ts topo.Server, keyspace, shard string) {
	vshard, kr, err := topo.ValidateShardName(shard)
	if err != nil {
		t.Fatalf("ValidateShardName(%v) failed: %v", shard, err)
	}

	ctx := context.Background()
	tablet := &pbt.Tablet{
		Alias: &pbt.TabletAlias{
			Cell: "cell1",
			Uid:  100,
		},
		Type:     pbt.TabletType_REPLICA,
		KeyRange: kr,
		Keyspace: keyspace,
		Shard:    vshard,
		PortMap: map[string]int32{
			"vt": 80,
		},
	}
	if err := ts.CreateTablet(ctx, tablet); err != nil {
		t.Fatalf("CreateTablet failed: %v", err)
	}
	if err := topotools.UpdateTabletEndpoints(ctx, ts, tablet); err != nil {
		t.Fatalf("topotools.UpdateTabletEndpoints failed: %v", err)
	}

	// register a tablet conn dialer that will return the instance
	// we want
	tabletconn.RegisterDialer(name, func(ctx context.Context, endPoint *pbt.EndPoint, k, s string, tabletType pbt.TabletType, timeout time.Duration) (tabletconn.TabletConn, error) {
		return &fakeTabletConn{
			endPoint:   endPoint,
			keyspace:   keyspace,
			shard:      vshard,
			tabletType: pbt.TabletType_REPLICA,
		}, nil
	})
	flag.Lookup("tablet_protocol").Value.Set(name)
}

//.........这里部分代码省略.........
			sct := topodatapb.KeyspaceIdType_UNSET
			if shard != "0" {
				var err error
				sct, err = key.ParseKeyspaceIDType(*shardingColumnType)
				if err != nil {
					log.Fatalf("parseKeyspaceIDType(%v) failed: %v", *shardingColumnType, err)
				}
				scn = *shardingColumnName
			}

			if err := ts.CreateKeyspace(ctx, keyspace, &topodatapb.Keyspace{
				ShardingColumnName: scn,
				ShardingColumnType: sct,
			}); err != nil {
				log.Fatalf("CreateKeyspace(%v) failed: %v", keyspace, err)
			}
			keyspaceMap[keyspace] = true
		}

		// create the master
		alias := &topodatapb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating master tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		masterController := tabletserver.NewServer()
		masterAgent := tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), masterController, dbcfgs, mysqld, keyspace, shard, dbname, "replica")
		if err := masterAgent.TabletExternallyReparented(ctx, ""); err != nil {
			log.Fatalf("TabletExternallyReparented failed on master: %v", err)
		}
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: topodatapb.TabletType_MASTER,
			dbname:     dbname,

			qsc:   masterController,
			agent: masterAgent,
		}
		uid++

		// create a replica slave
		alias = &topodatapb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating replica tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		replicaController := tabletserver.NewServer()
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: topodatapb.TabletType_REPLICA,
			dbname:     dbname,

			qsc:   replicaController,
			agent: tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), replicaController, dbcfgs, mysqld, keyspace, shard, dbname, "replica"),
		}
		uid++

		// create a rdonly slave
		alias = &topodatapb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating rdonly tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		rdonlyController := tabletserver.NewServer()
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: topodatapb.TabletType_RDONLY,
			dbname:     dbname,

			qsc:   rdonlyController,
			agent: tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), rdonlyController, dbcfgs, mysqld, keyspace, shard, dbname, "rdonly"),
		}
		uid++
	}

	// Rebuild the SrvKeyspace objects, we we can support range-based
	// sharding queries.
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, nil)
	for keyspace := range keyspaceMap {
		if err := wr.RebuildKeyspaceGraph(ctx, keyspace, nil, true); err != nil {
			log.Fatalf("cannot rebuild %v: %v", keyspace, err)
		}
	}

	// Register the tablet dialer for tablet server
	tabletconn.RegisterDialer("internal", dialer)
	*tabletconn.TabletProtocol = "internal"

	// Register the tablet manager client factory for tablet manager
	tmclient.RegisterTabletManagerClientFactory("internal", func() tmclient.TabletManagerClient {
		return &internalTabletManagerClient{}
	})
	*tmclient.TabletManagerProtocol = "internal"
}

func init() {
	tabletconn.RegisterDialer("grpc", DialTablet)
}

// initTabletMap creates the action agents and associated data structures
// for all tablets
func initTabletMap(ts topo.Server, topology string, mysqld mysqlctl.MysqlDaemon, dbcfgs *dbconfigs.DBConfigs, mycnf *mysqlctl.Mycnf) {
	tabletMap = make(map[uint32]*tablet)

	ctx := context.Background()

	// disable publishing of stats from query service
	flag.Lookup("queryserver-config-enable-publish-stats").Value.Set("false")

	var uid uint32 = 1
	keyspaceMap := make(map[string]bool)
	for _, entry := range strings.Split(topology, ",") {
		slash := strings.IndexByte(entry, '/')
		column := strings.IndexByte(entry, ':')
		if slash == -1 || column == -1 {
			log.Fatalf("invalid topology entry: %v", entry)
		}

		keyspace := entry[:slash]
		shard := entry[slash+1 : column]
		dbname := entry[column+1:]
		keyspaceMap[keyspace] = true

		localDBConfigs := &(*dbcfgs)
		localDBConfigs.App.DbName = dbname

		// create the master
		alias := &pb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating master tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		masterController := tabletserver.NewServer()
		masterAgent := tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), masterController, localDBConfigs, mysqld, keyspace, shard, dbname, "replica")
		if err := masterAgent.TabletExternallyReparented(ctx, ""); err != nil {
			log.Fatalf("TabletExternallyReparented failed on master: %v", err)
		}
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: pb.TabletType_MASTER,
			dbname:     dbname,

			qsc:   masterController,
			agent: masterAgent,
		}
		uid++

		// create a replica slave
		alias = &pb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating replica tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		replicaController := tabletserver.NewServer()
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: pb.TabletType_REPLICA,
			dbname:     dbname,

			qsc:   replicaController,
			agent: tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), replicaController, localDBConfigs, mysqld, keyspace, shard, dbname, "replica"),
		}
		uid++

		// create a rdonly slave
		alias = &pb.TabletAlias{
			Cell: cell,
			Uid:  uid,
		}
		log.Infof("Creating rdonly tablet %v for %v/%v", topoproto.TabletAliasString(alias), keyspace, shard)
		flag.Lookup("debug-url-prefix").Value.Set(fmt.Sprintf("/debug-%d", uid))
		rdonlyController := tabletserver.NewServer()
		tabletMap[uid] = &tablet{
			keyspace:   keyspace,
			shard:      shard,
			tabletType: pb.TabletType_RDONLY,
			dbname:     dbname,

			qsc:   rdonlyController,
			agent: tabletmanager.NewComboActionAgent(ctx, ts, alias, int32(8000+uid), int32(9000+uid), rdonlyController, localDBConfigs, mysqld, keyspace, shard, dbname, "rdonly"),
		}
		uid++
	}

	// Rebuild the SrvKeyspace objects, we we can support range-based
	// sharding queries.
	wr := wrangler.New(logutil.NewConsoleLogger(), ts, nil, 30*time.Second /*lockTimeout*/)
	for keyspace := range keyspaceMap {
		if err := wr.RebuildKeyspaceGraph(ctx, keyspace, nil, true); err != nil {
			log.Fatalf("cannot rebuild %v: %v", keyspace, err)
		}
	}

	// Register the tablet dialer
	tabletconn.RegisterDialer("internal", dialer)
//.........这里部分代码省略.........

func init() {
	tabletconn.RegisterDialer("fake_discovery", discoveryDialer)
	flag.Set("tablet_protocol", "fake_discovery")
	connMap = make(map[string]*fakeConn)
}

func init() {
	tabletconn.RegisterDialer("sandbox", sandboxDialer)
	flag.Set("tablet_protocol", "sandbox")
}

//.........这里部分代码省略.........
						Keyspace:   kpb.ServedFrom,
					},
					{
						TabletType: topodatapb.TabletType_RDONLY,
						Keyspace:   kpb.ServedFrom,
					},
				},
			}); err != nil {
				return fmt.Errorf("CreateKeyspace(%v) failed: %v", keyspace, err)
			}
		} else {
			// create a regular keyspace
			if err := ts.CreateKeyspace(ctx, keyspace, &topodatapb.Keyspace{
				ShardingColumnName: kpb.ShardingColumnName,
				ShardingColumnType: sct,
			}); err != nil {
				return fmt.Errorf("CreateKeyspace(%v) failed: %v", keyspace, err)
			}

			// iterate through the shards
			for _, spb := range kpb.Shards {
				shard := spb.Name
				dbname := spb.DbNameOverride
				if dbname == "" {
					dbname = fmt.Sprintf("vt_%v_%v", keyspace, shard)
				}
				dbcfgs.App.DbName = dbname

				// create the master
				if err := createTablet(ctx, ts, cell, uid, keyspace, shard, dbname, topodatapb.TabletType_MASTER, mysqld, dbcfgs); err != nil {
					return err
				}
				uid++

				// create a replica slave
				if err := createTablet(ctx, ts, cell, uid, keyspace, shard, dbname, topodatapb.TabletType_REPLICA, mysqld, dbcfgs); err != nil {
					return err
				}
				uid++

				// create a rdonly slave
				if err := createTablet(ctx, ts, cell, uid, keyspace, shard, dbname, topodatapb.TabletType_RDONLY, mysqld, dbcfgs); err != nil {
					return err
				}
				uid++
			}
		}

		// vschema for the keyspace
		if schemaDir != "" {
			f := path.Join(schemaDir, keyspace, "vschema.json")
			if _, err := os.Stat(f); err == nil {
				// load the vschema
				formal, err := vindexes.LoadFormalKeyspace(f)
				if err != nil {
					return fmt.Errorf("cannot load vschema file %v for keyspace %v: %v", f, keyspace, err)
				}

				if err := ts.SaveVSchema(ctx, keyspace, formal); err != nil {
					return fmt.Errorf("SaveVSchema(%v) failed: %v", keyspace, err)
				}
			} else {
				log.Infof("File %v doesn't exist, skipping vschema for keyspace %v", f, keyspace)
			}
		}

		// Rebuild the SrvKeyspace object, so we can support
		// range-based sharding queries, and export the redirects.
		if err := wr.RebuildKeyspaceGraph(ctx, keyspace, nil); err != nil {
			return fmt.Errorf("cannot rebuild %v: %v", keyspace, err)
		}
	}

	// Rebuild the SrvVSchema object
	if err := topotools.RebuildVSchema(ctx, wr.Logger(), ts, []string{cell}); err != nil {
		return fmt.Errorf("RebuildVSchemaGraph failed: %v", err)
	}

	// Register the tablet dialer for tablet server
	tabletconn.RegisterDialer("internal", dialer)
	*tabletconn.TabletProtocol = "internal"

	// Register the tablet manager client factory for tablet manager
	tmclient.RegisterTabletManagerClientFactory("internal", func() tmclient.TabletManagerClient {
		return &internalTabletManagerClient{}
	})
	*tmclient.TabletManagerProtocol = "internal"

	// run healthcheck on all vttablets
	tmc := tmclient.NewTabletManagerClient()
	for _, tablet := range tabletMap {
		tabletInfo, err := ts.GetTablet(ctx, tablet.agent.TabletAlias)
		if err != nil {
			return fmt.Errorf("cannot find tablet: %+v", tablet.agent.TabletAlias)
		}
		tmc.RunHealthCheck(ctx, tabletInfo.Tablet)
	}

	return nil
}

func init() {
	tabletconn.RegisterDialer(protocolName, DialTablet)
}