Merge pull request vitessio#2089 from michael-berlin/vtctl_migrate_se…

…rved_types_sleep

vtctl: MigrateServedTypes: Wait for several seconds for the old table…

go/vt/wrangler/keyspace.go

@@ -5,6 +5,7 @@
 package wrangler
 
 import (
+	"flag"
 	"fmt"
 	"strings"
 	"sync"
@@ -23,6 +24,12 @@ import (
 	topodatapb "github.com/youtube/vitess/go/vt/proto/topodata"
 )
 
+// TODO(b/26388813): Remove these flags once vtctl WaitForDrain is integrated in the vtctl MigrateServed* commands.
+var (
+	waitForDrainSleepRdonly  = flag.Duration("wait_for_drain_sleep_rdonly", 5*time.Second, "time to wait before shutting the query service on old RDONLY tablets during MigrateServedTypes")
+	waitForDrainSleepReplica = flag.Duration("wait_for_drain_sleep_replica", 15*time.Second, "time to wait before shutting the query service on old REPLICA tablets during MigrateServedTypes")
+)
+
 // keyspace related methods for Wrangler
 
 // SetKeyspaceShardingInfo locks a keyspace and sets its ShardingColumnName
@@ -151,6 +158,22 @@ func (wr *Wrangler) MigrateServedTypes(ctx context.Context, keyspace, shard stri
 			// For a forwards migration, we just disabled query service on the source shards
 			refreshShards = sourceShards
 		}
+
+		// TODO(b/26388813): Integrate vtctl WaitForDrain here instead of just sleeping.
+		var waitForDrainSleep time.Duration
+		switch servedType {
+		case topodatapb.TabletType_RDONLY:
+			waitForDrainSleep = *waitForDrainSleepRdonly
+		case topodatapb.TabletType_REPLICA:
+			waitForDrainSleep = *waitForDrainSleepReplica
+		default:
+			wr.Logger().Warningf("invalid TabletType: %v for MigrateServedTypes command", servedType)
+		}
+
+		wr.Logger().Infof("WaitForDrain: Sleeping for %.0f seconds before shutting down query service on old tablets...", waitForDrainSleep.Seconds())
+		time.Sleep(waitForDrainSleep)
+		wr.Logger().Infof("WaitForDrain: Sleeping finished. Shutting down queryservice on old tablets now.")
+
 		for _, si := range refreshShards {
 			rec.RecordError(wr.RefreshTabletsByShard(ctx, si, servedType, cells))
 		}

go/vt/wrangler/testlib/migrate_served_types_test.go

@@ -5,6 +5,7 @@
 package testlib
 
 import (
+	"flag"
 	"testing"
 
 	"github.com/youtube/vitess/go/sqltypes"
@@ -43,6 +44,10 @@ func checkShardSourceShards(t *testing.T, ts topo.Server, shard string, expected
 }
 
 func TestMigrateServedTypes(t *testing.T) {
+	// TODO(b/26388813): Remove the next two lines once vtctl WaitForDrain is integrated in the vtctl MigrateServed* commands.
+	flag.Set("wait_for_drain_sleep_rdonly", "0s")
+	flag.Set("wait_for_drain_sleep_replica", "0s")
+
 	db := fakesqldb.Register()
 	ts := zktestserver.New(t, []string{"cell1", "cell2"})
 	wr := wrangler.New(logutil.NewConsoleLogger(), ts, tmclient.NewTabletManagerClient())

test/utils.py

@@ -868,6 +868,10 @@ def run_vtctl_vtctl(clargs, auto_log=False, expect_fail=False,
   args.extend(['-throttler_client_protocol',
                protocols_flavor().throttler_client_protocol()])
   args.extend(['-vtgate_protocol', protocols_flavor().vtgate_protocol()])
+  # TODO(b/26388813): Remove the next two lines once vtctl WaitForDrain is
+  #                   integrated in the vtctl MigrateServed* commands.
+  args.extend(['--wait_for_drain_sleep_rdonly', '0s'])
+  args.extend(['--wait_for_drain_sleep_replica', '0s'])
 
   if auto_log:
     args.append('--stderrthreshold=%s' % get_log_level())
@@ -1235,6 +1239,10 @@ def start(self, enable_schema_change_dir=False):
         protocols_flavor().throttler_client_protocol(),
         '-vtgate_protocol', protocols_flavor().vtgate_protocol(),
     ] + environment.topo_server().flags()
+    # TODO(b/26388813): Remove the next two lines once vtctl WaitForDrain is
+    #                   integrated in the vtctl MigrateServed* commands.
+    args.extend(['--wait_for_drain_sleep_rdonly', '0s'])
+    args.extend(['--wait_for_drain_sleep_replica', '0s'])
     if enable_schema_change_dir:
       args += [
           '--schema_change_dir', self.schema_change_dir,