From 1d0873966822bc43b9bb80fc15e891e325510be8 Mon Sep 17 00:00:00 2001 From: Matt Lord Date: Tue, 27 Aug 2024 11:18:17 -0400 Subject: [PATCH] Improve target_shards handling Signed-off-by: Matt Lord --- .../command/vreplication/vdiff/vdiff.go | 4 ++-- go/vt/vtctl/workflow/server.go | 12 ++++------ go/vt/vtctl/workflow/utils.go | 22 +++++++++++++++++++ 3 files changed, 28 insertions(+), 10 deletions(-) diff --git a/go/cmd/vtctldclient/command/vreplication/vdiff/vdiff.go b/go/cmd/vtctldclient/command/vreplication/vdiff/vdiff.go index 6ad5abd818f..b2aefeee03b 100644 --- a/go/cmd/vtctldclient/command/vreplication/vdiff/vdiff.go +++ b/go/cmd/vtctldclient/command/vreplication/vdiff/vdiff.go @@ -198,7 +198,7 @@ vtctldclient --server localhost:15999 vdiff --workflow commerce2customer --targe for _, shard := range resumeOptions.TargetShards { if !key.IsValidKeyRange(shard) { - return fmt.Errorf("invalid target shard provided: %s", shard) + return fmt.Errorf("invalid target shard provided: %q", shard) } } @@ -250,7 +250,7 @@ vtctldclient --server localhost:15999 vdiff --workflow commerce2customer --targe for _, shard := range stopOptions.TargetShards { if !key.IsValidKeyRange(shard) { - return fmt.Errorf("invalid target shard provided: %s", shard) + return fmt.Errorf("invalid target shard provided: %q", shard) } } diff --git a/go/vt/vtctl/workflow/server.go b/go/vt/vtctl/workflow/server.go index dc45310494e..afffc0dba63 100644 --- a/go/vt/vtctl/workflow/server.go +++ b/go/vt/vtctl/workflow/server.go @@ -1927,10 +1927,8 @@ func (s *Server) VDiffResume(ctx context.Context, req *vtctldatapb.VDiffResumeRe } if len(targetShards) > 0 { - for key, target := range ts.targets { - if !slices.Contains(targetShards, target.GetShard().ShardName()) { - delete(ts.targets, key) - } + if err := applyTargetShards(ts, targetShards); err != nil { + return nil, err } } @@ -2012,10 +2010,8 @@ func (s *Server) VDiffStop(ctx context.Context, req *vtctldatapb.VDiffStopReques } if len(targetShards) > 0 { - for key, target := range ts.targets { - if !slices.Contains(targetShards, target.GetShard().ShardName()) { - delete(ts.targets, key) - } + if err := applyTargetShards(ts, targetShards); err != nil { + return nil, err } } diff --git a/go/vt/vtctl/workflow/utils.go b/go/vt/vtctl/workflow/utils.go index 374d96396f2..c3694cfa3e6 100644 --- a/go/vt/vtctl/workflow/utils.go +++ b/go/vt/vtctl/workflow/utils.go @@ -23,11 +23,13 @@ import ( "fmt" "hash/fnv" "math" + "slices" "sort" "strconv" "strings" "sync" + "golang.org/x/exp/maps" "google.golang.org/protobuf/encoding/prototext" "vitess.io/vitess/go/mysql/sqlerror" @@ -967,3 +969,23 @@ func defaultErrorHandler(logger logutil.Logger, message string, err error) (*[]s logger.Error(werr) return nil, werr } + +// applyTargetShards applies the targetShards, coming from a command, to the trafficSwitcher. +// It will return an error if the targetShards list contains a shard that does not exist in +// the target keyspace. +// It will remove any target shards in the trafficSwitcher that are not in the targetShards list. +func applyTargetShards(ts *trafficSwitcher, targetShards []string) error { + shards := maps.Keys(ts.targets) + for _, targetShard := range targetShards { + if !slices.Contains(shards, targetShard) { + return vterrors.Errorf(vtrpcpb.Code_FAILED_PRECONDITION, "specified target shard %s not found in target keyspace %s", + targetShard, ts.targetKeyspace) + } + } + for key, target := range ts.targets { + if !slices.Contains(targetShards, target.GetShard().ShardName()) { + delete(ts.targets, key) + } + } + return nil +}