Modify Wrangler to take in a TabletManagerClient at creation

go/cmd/vtctl/vtctl.go

@@ -19,6 +19,7 @@ import (
 	"github.com/youtube/vitess/go/exit"
 	"github.com/youtube/vitess/go/vt/logutil"
 	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"
+	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
 	"github.com/youtube/vitess/go/vt/topo"
 	"github.com/youtube/vitess/go/vt/vtctl"
 	"github.com/youtube/vitess/go/vt/wrangler"
@@ -77,7 +78,7 @@ func main() {
 	defer topo.CloseServers()
 
 	ctx, cancel := context.WithTimeout(context.Background(), *waitTime)
-	wr := wrangler.New(logutil.NewConsoleLogger(), topoServer, *lockWaitTimeout)
+	wr := wrangler.New(logutil.NewConsoleLogger(), topoServer, tmclient.NewTabletManagerClient(), *lockWaitTimeout)
 	installSignalHandlers(cancel)
 
 	err := vtctl.RunCommand(ctx, wr, args)

go/cmd/vtctld/actions.go

@@ -10,6 +10,7 @@ import (
 	"github.com/youtube/vitess/go/acl"
 	"github.com/youtube/vitess/go/vt/logutil"
 	"github.com/youtube/vitess/go/vt/tabletmanager/actionnode"
+	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
 	"github.com/youtube/vitess/go/vt/topo"
 	"github.com/youtube/vitess/go/vt/wrangler"
 	"golang.org/x/net/context"
@@ -98,7 +99,7 @@ func (ar *ActionRepository) ApplyKeyspaceAction(actionName, keyspace string, r *
 
 	// FIXME(alainjobart) copy web context info
 	ctx, cancel := context.WithTimeout(context.TODO(), *actionTimeout)
-	wr := wrangler.New(logutil.NewConsoleLogger(), ar.ts, *lockTimeout)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ar.ts, tmclient.NewTabletManagerClient(), *lockTimeout)
 	output, err := action(ctx, wr, keyspace, r)
 	cancel()
 	if err != nil {
@@ -126,7 +127,7 @@ func (ar *ActionRepository) ApplyShardAction(actionName, keyspace, shard string,
 
 	// FIXME(alainjobart) copy web context info
 	ctx, cancel := context.WithTimeout(context.TODO(), *actionTimeout)
-	wr := wrangler.New(logutil.NewConsoleLogger(), ar.ts, *lockTimeout)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ar.ts, tmclient.NewTabletManagerClient(), *lockTimeout)
 	output, err := action(ctx, wr, keyspace, shard, r)
 	cancel()
 	if err != nil {
@@ -158,7 +159,7 @@ func (ar *ActionRepository) ApplyTabletAction(actionName string, tabletAlias top
 	// run the action
 	// FIXME(alainjobart) copy web context info
 	ctx, cancel := context.WithTimeout(context.TODO(), *actionTimeout)
-	wr := wrangler.New(logutil.NewConsoleLogger(), ar.ts, *lockTimeout)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ar.ts, tmclient.NewTabletManagerClient(), *lockTimeout)
 	output, err := action.method(ctx, wr, tabletAlias, r)
 	cancel()
 	if err != nil {

go/cmd/vtctld/tablet_data_test.go

@@ -16,7 +16,7 @@ import (
 
 func TestTabletData(t *testing.T) {
 	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
-	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Second)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)
 
 	tablet1 := testlib.NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER, testlib.TabletKeyspaceShard(t, "ks", "-80"))
 	tablet1.StartActionLoop(t, wr)

go/cmd/vtjanitor/vtjanitor.go

@@ -13,6 +13,7 @@ import (
 	"github.com/youtube/vitess/go/vt/logutil"
 	"github.com/youtube/vitess/go/vt/servenv"
 	"github.com/youtube/vitess/go/vt/tabletmanager/actionnode"
+	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
 	"github.com/youtube/vitess/go/vt/topo"
 	"github.com/youtube/vitess/go/vt/wrangler"
 )
@@ -52,7 +53,7 @@ func main() {
 
 	ts := topo.GetServer()
 
-	scheduler, err := janitor.New(*keyspace, *shard, ts, wrangler.New(logutil.NewConsoleLogger(), ts, *lockTimeout), *sleepTime)
+	scheduler, err := janitor.New(*keyspace, *shard, ts, wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), *lockTimeout), *sleepTime)
 	if err != nil {
 		log.Errorf("janitor.New: %v", err)
 		exit.Return(1)

go/cmd/vtworker/vtworker.go

@@ -25,6 +25,7 @@ import (
 	"github.com/youtube/vitess/go/exit"
 	"github.com/youtube/vitess/go/vt/logutil"
 	"github.com/youtube/vitess/go/vt/servenv"
+	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
 	"github.com/youtube/vitess/go/vt/topo"
 	"github.com/youtube/vitess/go/vt/worker"
 	"github.com/youtube/vitess/go/vt/wrangler"
@@ -100,7 +101,7 @@ func main() {
 	defer topo.CloseServers()
 
 	// The logger will be replaced when we start a job.
-	wr = wrangler.New(logutil.NewConsoleLogger(), ts, 30*time.Second)
+	wr = wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), 30*time.Second)
 	if len(args) == 0 {
 		// In interactive mode, initialize the web UI to choose a command.
 		initInteractiveMode()

go/vt/tabletmanager/faketmclient/fake_client.go

@@ -21,10 +21,8 @@ import (
 
 type timeoutError error
 
-func init() {
-	tmclient.RegisterTabletManagerClientFactory("fake", func() tmclient.TabletManagerClient {
-		return &FakeTabletManagerClient{}
-	})
+func NewFakeTabletManagerClient() tmclient.TabletManagerClient {
+	return &FakeTabletManagerClient{}
 }
 
 // FakeTabletManagerClient implements tmclient.TabletManagerClient

go/vt/topo/test/faketopo/faketopo.go

@@ -10,6 +10,7 @@ import (
 	"github.com/youtube/vitess/go/vt/key"
 	"github.com/youtube/vitess/go/vt/logutil"
 	"github.com/youtube/vitess/go/vt/mysqlctl"
+	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
 	"github.com/youtube/vitess/go/vt/topo"
 	"github.com/youtube/vitess/go/vt/wrangler"
 	"golang.org/x/net/context"
@@ -51,7 +52,7 @@ type Fixture struct {
 
 // New creates a topology fixture.
 func New(t *testing.T, logger logutil.Logger, ts topo.Server, cells []string) *Fixture {
-	wr := wrangler.New(logger, ts, 1*time.Second)
+	wr := wrangler.New(logger, ts, tmclient.NewTabletManagerClient(), 1*time.Second)
 
 	return &Fixture{
 		T:        t,

go/vt/vtctl/gorpcvtctlserver/server.go

@@ -13,6 +13,7 @@ import (
 
 	"github.com/youtube/vitess/go/vt/logutil"
 	"github.com/youtube/vitess/go/vt/servenv"
+	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
 	"github.com/youtube/vitess/go/vt/topo"
 	"github.com/youtube/vitess/go/vt/vtctl"
 	"github.com/youtube/vitess/go/vt/vtctl/gorpcproto"
@@ -47,7 +48,7 @@ func (s *VtctlServer) ExecuteVtctlCommand(ctx context.Context, query *gorpcproto
 	}()
 
 	// create the wrangler
-	wr := wrangler.New(logger, s.ts, query.LockTimeout)
+	wr := wrangler.New(logger, s.ts, tmclient.NewTabletManagerClient(), query.LockTimeout)
 	// FIXME(alainjobart) use a single context, copy the source info from it
 	ctx, cancel := context.WithTimeout(context.TODO(), query.ActionTimeout)
 

go/vt/worker/split_clone_test.go

@@ -19,6 +19,7 @@ import (
 	"github.com/youtube/vitess/go/vt/logutil"
 	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"
 	_ "github.com/youtube/vitess/go/vt/tabletmanager/gorpctmclient"
+	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
 	_ "github.com/youtube/vitess/go/vt/tabletserver/gorpctabletconn"
 	"github.com/youtube/vitess/go/vt/tabletserver/proto"
 	"github.com/youtube/vitess/go/vt/topo"
@@ -235,7 +236,7 @@ func TestSplitClonePopulateBlpCheckpoint(t *testing.T) {
 
 func testSplitClone(t *testing.T, strategy string) {
 	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
-	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Second)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)
 
 	sourceMaster := testlib.NewFakeTablet(t, wr, "cell1", 0,
 		topo.TYPE_MASTER, testlib.TabletKeyspaceShard(t, "ks", "-80"))

go/vt/worker/vertical_split_clone_test.go

@@ -18,6 +18,7 @@ import (
 	"github.com/youtube/vitess/go/vt/logutil"
 	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"
 	_ "github.com/youtube/vitess/go/vt/tabletmanager/gorpctmclient"
+	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
 	_ "github.com/youtube/vitess/go/vt/tabletserver/gorpctabletconn"
 	"github.com/youtube/vitess/go/vt/tabletserver/proto"
 	"github.com/youtube/vitess/go/vt/topo"
@@ -219,7 +220,7 @@ func TestVerticalSplitClonePopulateBlpCheckpoint(t *testing.T) {
 
 func testVerticalSplitClone(t *testing.T, strategy string) {
 	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
-	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Second)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)
 
 	sourceMaster := testlib.NewFakeTablet(t, wr, "cell1", 0,
 		topo.TYPE_MASTER, testlib.TabletKeyspaceShard(t, "source_ks", "0"))

go/vt/worker/vertical_split_diff_test.go

@@ -17,6 +17,7 @@ import (
 	"github.com/youtube/vitess/go/vt/dbconnpool"
 	"github.com/youtube/vitess/go/vt/logutil"
 	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"
+	"github.com/youtube/vitess/go/vt/tabletmanager/faketmclient"
 	_ "github.com/youtube/vitess/go/vt/tabletmanager/gorpctmclient"
 	_ "github.com/youtube/vitess/go/vt/tabletserver/gorpctabletconn"
 	"github.com/youtube/vitess/go/vt/tabletserver/proto"
@@ -220,7 +221,7 @@ func TestVerticalSplitDiff(t *testing.T) {
 
 func testVerticalSplitDiff(t *testing.T, strategy string) {
 	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
-	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Second)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ts, faketmclient.NewFakeTabletManagerClient(), time.Second)
 	ctx := context.Background()
 
 	sourceMaster := testlib.NewFakeTablet(t, wr, "cell1", 0,

go/vt/wrangler/testlib/copy_schema_shard_test.go

@@ -15,6 +15,7 @@ import (
 	"github.com/youtube/vitess/go/vt/logutil"
 	myproto "github.com/youtube/vitess/go/vt/mysqlctl/proto"
 	_ "github.com/youtube/vitess/go/vt/tabletmanager/gorpctmclient"
+	"github.com/youtube/vitess/go/vt/tabletmanager/tmclient"
 	_ "github.com/youtube/vitess/go/vt/tabletserver/gorpctabletconn"
 	"github.com/youtube/vitess/go/vt/topo"
 	"github.com/youtube/vitess/go/vt/wrangler"
@@ -124,7 +125,7 @@ func DestinationsFactory(t *testing.T) func() (dbconnpool.PoolConnection, error)
 
 func TestCopySchemaShard(t *testing.T) {
 	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
-	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Second)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)
 
 	sourceMaster := NewFakeTablet(t, wr, "cell1", 0,
 		topo.TYPE_MASTER, TabletKeyspaceShard(t, "ks", "-80"))

go/vt/wrangler/testlib/reparent_external_test.go

@@ -36,7 +36,7 @@ func testTabletExternallyReparented(t *testing.T, fast bool) {
 
 	ctx := context.Background()
 	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
-	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Second)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)
 
 	// Create an old master, a new master, two good slaves, one bad slave
 	oldMaster := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)
@@ -187,7 +187,7 @@ func testTabletExternallyReparentedWithDifferentMysqlPort(t *testing.T, fast boo
 	tabletmanager.SetReparentFlags(fast, time.Minute /* finalizeTimeout */)
 
 	ts := zktopo.NewTestServer(t, []string{"cell1"})
-	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Second)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)
 
 	// Create an old master, a new master, two good slaves, one bad slave
 	oldMaster := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)
@@ -245,7 +245,7 @@ func testTabletExternallyReparentedContinueOnUnexpectedMaster(t *testing.T, fast
 	tabletmanager.SetReparentFlags(fast, time.Minute /* finalizeTimeout */)
 
 	ts := zktopo.NewTestServer(t, []string{"cell1"})
-	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Second)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)
 
 	// Create an old master, a new master, two good slaves, one bad slave
 	oldMaster := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)
@@ -297,7 +297,7 @@ func testTabletExternallyReparentedFailedOldMaster(t *testing.T, fast bool) {
 	tabletmanager.SetReparentFlags(fast, time.Minute /* finalizeTimeout */)
 
 	ts := zktopo.NewTestServer(t, []string{"cell1", "cell2"})
-	wr := wrangler.New(logutil.NewConsoleLogger(), ts, time.Second)
+	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient(), time.Second)
 
 	// Create an old master, a new master, and a good slave.
 	oldMaster := NewFakeTablet(t, wr, "cell1", 0, topo.TYPE_MASTER)

go/vt/wrangler/wrangler.go

@@ -45,11 +45,11 @@ type Wrangler struct {
 // of the time, we want to immediately know that our action will
 // fail. However, automated action will need some time to arbitrate
 // the locks.
-func New(logger logutil.Logger, ts topo.Server, lockTimeout time.Duration) *Wrangler {
+func New(logger logutil.Logger, ts topo.Server, tmc tmclient.TabletManagerClient, lockTimeout time.Duration) *Wrangler {
 	return &Wrangler{
 		logger:      logger,
 		ts:          ts,
-		tmc:         tmclient.NewTabletManagerClient(),
+		tmc:         tmc,
 		lockTimeout: lockTimeout,
 	}
 }