Add timeouts to vx.CallbackContext usage

Signed-off-by: Matt Lord <mattalord@gmail.com>
vitessio · May 21, 2024 · 51b75c1 · 51b75c1
1 parent 4401a21
commit 51b75c1
Show file tree

Hide file tree

Showing 2 changed files with 26 additions and 3 deletions.
diff --git a/go/vt/vtctl/workflow/server.go b/go/vt/vtctl/workflow/server.go
@@ -2005,7 +2005,9 @@ func (s *Server) WorkflowDelete(ctx context.Context, req *vtctldatapb.WorkflowDe
  s.optimizeCopyStateTable(tablet.Tablet)
  return res.Result, err
  }
- res, err := vx.CallbackContext(ctx, callback)
+ delCtx, delCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout*2)
+ defer delCancel()
+ res, err := vx.CallbackContext(delCtx, callback)
  if err != nil {
  return nil, err
  }
@@ -2015,7 +2017,7 @@ func (s *Server) WorkflowDelete(ctx context.Context, req *vtctldatapb.WorkflowDe
  }
 
  // Cleanup related data and artifacts.
- if _, err := s.DropTargets(ctx, ts, req.GetKeepData(), req.GetKeepRoutingRules(), false); err != nil {
+ if _, err := s.DropTargets(delCtx, ts, req.GetKeepData(), req.GetKeepRoutingRules(), false); err != nil {
  if topo.IsErrType(err, topo.NoNode) {
  return nil, vterrors.Wrapf(err, "%s keyspace does not exist", req.Keyspace)
  }
@@ -2289,7 +2291,9 @@ func (s *Server) WorkflowUpdate(ctx context.Context, req *vtctldatapb.WorkflowUp
  }
  return res.Result, err
  }
- res, err := vx.CallbackContext(ctx, callback)
+ updCtx, updCancel := context.WithTimeout(ctx, topo.RemoteOperationTimeout*2)
+ defer updCancel()
+ res, err := vx.CallbackContext(updCtx, callback)
  if err != nil {
  if topo.IsErrType(err, topo.NoNode) {
  return nil, vterrors.Wrapf(err, "%s keyspace does not exist", req.Keyspace)

diff --git a/go/vt/vtctl/workflow/server_test.go b/go/vt/vtctl/workflow/server_test.go
@@ -211,6 +211,18 @@ func TestWorkflowDelete(t *testing.T) {
  s := NewServer(vtenv.NewTestEnv(), ts, tmc)
  err := s.ts.CreateKeyspace(ctx, ksName, &topodatapb.Keyspace{})
  require.NoError(t, err)
+ err = s.ts.CreateShard(ctx, ksName, "0")
+ require.NoError(t, err)
+ err = s.ts.CreateTablet(ctx, &topodatapb.Tablet{
+ Hostname: "host1",
+ Keyspace: ksName,
+ Shard: "0",
+ Alias: &topodatapb.TabletAlias{
+ Cell: cellName,
+ Uid: 100,
+ },
+ })
+ require.NoError(t, err)
 
  testcases := []struct {
  name string
@@ -228,6 +240,13 @@ func TestWorkflowDelete(t *testing.T) {
  }
  for _, tc := range testcases {
  t.Run(tc.name, func(t *testing.T) {
+ res, err := s.MoveTablesCreate(ctx, &vtctldatapb.MoveTablesCreateRequest{
+ TargetKeyspace: ksName,
+ Workflow: wfName,
+ AllTables: true,
+ })
+ require.NoError(t, err)
+ require.NotNil(t, res)
  got, err := s.WorkflowDelete(ctx, tc.req)
  if (err != nil) != tc.wantErr {
  t.Errorf("Server.WorkflowDelete() error = %v, wantErr %v", err, tc.wantErr)