From 51b75c16257dc4113098340114693c8a62c6bb55 Mon Sep 17 00:00:00 2001
From: Matt Lord <mattalord@gmail.com>
Date: Tue, 21 May 2024 10:04:13 -0400
Subject: [PATCH] Add timeouts to vx.CallbackContext usage

Signed-off-by: Matt Lord <mattalord@gmail.com>
---
 go/vt/vtctl/workflow/server.go      | 10 +++++++---
 go/vt/vtctl/workflow/server_test.go | 19 +++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/go/vt/vtctl/workflow/server.go b/go/vt/vtctl/workflow/server.go
index 7af3a3982ea..00c36b2d3dc 100644
--- 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
index b66e0283203..a6a38fd3040 100644
--- 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)