Repository: vitessio/vitess
Pull request #10006 (2022-03-30 06:57:19)

Alternate fix for flaky TestReparentDuringWorkerCopy

<!--
  Thank you for your contribution to the Vitess project.
  How to contribute: https://vitess.io/docs/contributing/
  Please first make sure there is an open Issue to discuss the feature/fix suggested in this PR.
  If this is a new feature, please mark the Issue as "RFC".
 -->

<!-- if this PR is Work in Progress please create it as a Draft Pull Request -->

## Description
<!-- A few sentences describing the overall goals of the pull request's commits. -->
<!-- If this is a bug fix and you think the fix should be backported, please write so. -->


As seen in the comment https://github.com/vitessio/vitess/issues/9985#issuecomment-1079666698, the test `TestReparentDuringWorkerCopy` is flaky.
The PR #9986 tried to fix it by adding a 1 second wait between the two PRS calls. That did seem to resolve the problem on local testing but the CI is still flaky. This PR reverses that change and instead goes ahead with the alternate proposed fix https://github.com/vitessio/vitess/pull/9986#discussion_r835756712.

## Related Issue(s)
<!-- List related issues and pull requests. If this PR fixes an issue, please add it using Fixes #????  -->
#9985 

## Checklist
- [x] Should this PR be backported? No
- [x] Tests were added or are not required
- [x] Documentation was added or is not required

## Deployment Notes
<!-- Notes regarding deployment of the contained body of work. These should note any db migrations, etc. -->


Patch:
diff --git a/go/vt/sqlparser/ast_rewriting.go b/go/vt/sqlparser/ast_rewriting.go
index 18a2d93be2..817aa3a887 100644
--- a/go/vt/sqlparser/ast_rewriting.go
+++ b/go/vt/sqlparser/ast_rewriting.go
@@ -280,19 +280,19 @@ func newASTRewriter(keyspace string, selectLimit int, setVarComment string, sysV
 }
 
 const (
-	//LastInsertIDName is a reserved bind var name for last_insert_id()
+	// LastInsertIDName is a reserved bind var name for last_insert_id()
 	LastInsertIDName = "__lastInsertId"
 
-	//DBVarName is a reserved bind var name for database()
+	// DBVarName is a reserved bind var name for database()
 	DBVarName = "__vtdbname"
 
-	//FoundRowsName is a reserved bind var name for found_rows()
+	// FoundRowsName is a reserved bind var name for found_rows()
 	FoundRowsName = "__vtfrows"
 
-	//RowCountName is a reserved bind var name for row_count()
+	// RowCountName is a reserved bind var name for row_count()
 	RowCountName = "__vtrcount"
 
-	//UserDefinedVariableName is what we prepend bind var names for user defined variables
+	// UserDefinedVariableName is what we prepend bind var names for user defined variables
 	UserDefinedVariableName = "__vtudv"
 )
 
@@ -506,6 +506,9 @@ func (er *astRewriter) funcRewrite(cursor *Cursor, node *FuncExpr) {
 }
 
 func (er *astRewriter) unnestSubQueries(cursor *Cursor, subquery *Subquery) {
+	if _, isExists := cursor.Parent().(*ExistsExpr); isExists {
+		return
+	}
 	sel, isSimpleSelect := subquery.Select.(*Select)
 	if !isSimpleSelect {
 		return
diff --git a/go/vt/vtgate/planbuilder/physical/route_planning.go b/go/vt/vtgate/planbuilder/physical/route_planning.go
index 58c02c94d3..e7926c5ef6 100644
--- a/go/vt/vtgate/planbuilder/physical/route_planning.go
+++ b/go/vt/vtgate/planbuilder/physical/route_planning.go
@@ -569,17 +569,22 @@ func leaves(op abstract.Operator) (sources []abstract.Operator) {
 }
 
 func tryMergeReferenceTable(aRoute, bRoute *Route, merger mergeFunc) (*Route, error) {
-	// if either side is a reference table, we can just merge it and use the opcode of the other side
-	var opCode engine.Opcode
-	var selected *VindexOption
+	var (
+		// if either side is a reference table, we can just merge it and use the opcode of the other side
+		opCode engine.Opcode
+		vindex *VindexOption
+		ks     *vindexes.Keyspace
+	)
 
 	switch {
 	case aRoute.RouteOpCode == engine.Reference:
-		selected = bRoute.Selected
+		vindex = bRoute.Selected
 		opCode = bRoute.RouteOpCode
+		ks = bRoute.Keyspace
 	case bRoute.RouteOpCode == engine.Reference:
-		selected = aRoute.Selected
+		vindex = aRoute.Selected
 		opCode = aRoute.RouteOpCode
+		ks = aRoute.Keyspace
 	default:
 		return nil, nil
 	}
@@ -589,7 +594,8 @@ func tryMergeReferenceTable(aRoute, bRoute *Route, merger mergeFunc) (*Route, er
 		return nil, err
 	}
 	r.RouteOpCode = opCode
-	r.Selected = selected
+	r.Selected = vindex
+	r.Keyspace = ks
 	return r, nil
 }
 


Test patch:
diff --git a/go/test/endtoend/worker/worker_test.go b/go/test/endtoend/worker/worker_test.go
index 45cd9b3767..cb4dc6aa79 100644
--- a/go/test/endtoend/worker/worker_test.go
+++ b/go/test/endtoend/worker/worker_test.go
@@ -261,7 +261,7 @@ func verifySuccessfulWorkerCopyWithReparent(t *testing.T, isMysqlDown bool) {
 		// If MySQL is down, we wait until vtworker retried at least once to make
 		// sure it reached the point where a write failed due to MySQL being down.
 		// There should be two retries at least, one for each destination shard.
-		pollForVarsWorkerRetryCount(t, 1)
+		pollForVarsWorkerRetryCount(t, 2)
 
 		// Bring back primaries. Since we test with semi-sync now, we need at least
 		// one replica for the new primary. This test is already quite expensive,
@@ -306,12 +306,6 @@ func verifySuccessfulWorkerCopyWithReparent(t *testing.T, isMysqlDown bool) {
 	// Reparent away from the old primaries.
 	localCluster.VtctlclientProcess.ExecuteCommand("PlannedReparentShard", "--", "--keyspace_shard",
 		"test_keyspace/-80", "--new_primary", shard0Replica.Alias)
-
-	// Sleep here for one second. (as long as executefetch_retry_time)
-	// We want the next PRS call to run after the first one has passed and the retry time for the
-	// failing query has expired. More info on why this is needed https://github.com/vitessio/vitess/issues/9985#issuecomment-1079666698
-	time.Sleep(1 * time.Second)
-
 	localCluster.VtctlclientProcess.ExecuteCommand("PlannedReparentShard", "--", "--keyspace_shard",
 		"test_keyspace/80-", "--new_primary", shard1Replica.Alias)
 
@@ -411,7 +405,7 @@ func pollForVarsWorkerRetryCount(t *testing.T, count int) {
 		}
 		continue
 	}
-	assert.Greater(t, workerRetryCountInt, count)
+	assert.GreaterOrEqual(t, workerRetryCountInt, count)
 }
 
 // vttablet: the Tablet instance to modify.
diff --git a/go/vt/sqlparser/ast_rewriting_test.go b/go/vt/sqlparser/ast_rewriting_test.go
index ad793ce654..d1d4c5f36a 100644
--- a/go/vt/sqlparser/ast_rewriting_test.go
+++ b/go/vt/sqlparser/ast_rewriting_test.go
@@ -163,6 +163,9 @@ func TestRewrites(in *testing.T) {
 	}, {
 		in:       "select (select 42) from dual",
 		expected: "select 42 as `(select 42 from dual)` from dual",
+	}, {
+		in:       "select exists(select 1) from user",
+		expected: "select exists(select 1) from user",
 	}, {
 		in:       "select * from user where col = (select 42)",
 		expected: "select * from user where col = 42",
diff --git a/go/vt/vtgate/planbuilder/plan_test.go b/go/vt/vtgate/planbuilder/plan_test.go
index 001b386f36..593b3164f1 100644
--- a/go/vt/vtgate/planbuilder/plan_test.go
+++ b/go/vt/vtgate/planbuilder/plan_test.go
@@ -258,6 +258,30 @@ func TestOne(t *testing.T) {
 	testFile(t, "onecase.txt", "", vschema)
 }
 
+func TestOneWithMainAsDefault(t *testing.T) {
+	vschema := &vschemaWrapper{
+		v: loadSchema(t, "schema_test.json", true),
+		keyspace: &vindexes.Keyspace{
+			Name:    "main",
+			Sharded: false,
+		},
+	}
+
+	testFile(t, "onecase.txt", "", vschema)
+}
+
+func TestOneWithSecondUserAsDefault(t *testing.T) {
+	vschema := &vschemaWrapper{
+		v: loadSchema(t, "schema_test.json", true),
+		keyspace: &vindexes.Keyspace{
+			Name:    "second_user",
+			Sharded: true,
+		},
+	}
+
+	testFile(t, "onecase.txt", "", vschema)
+}
+
 func TestRubyOnRailsQueries(t *testing.T) {
 	vschemaWrapper := &vschemaWrapper{
 		v:             loadSchema(t, "rails_schema_test.json", true),
@@ -372,6 +396,21 @@ func TestWithDefaultKeyspaceFromFile(t *testing.T) {
 	testFile(t, "call_cases.txt", testOutputTempDir, vschema)
 }
 
+func TestWithDefaultKeyspaceFromFileSharded(t *testing.T) {
+	// We are testing this separately so we can set a default keyspace
+	vschema := &vschemaWrapper{
+		v: loadSchema(t, "schema_test.json", true),
+		keyspace: &vindexes.Keyspace{
+			Name:    "second_user",
+			Sharded: true,
+		},
+		tabletType: topodatapb.TabletType_PRIMARY,
+	}
+
+	testOutputTempDir := makeTestOutput(t)
+	testFile(t, "select_cases_with_default.txt", testOutputTempDir, vschema)
+}
+
 func TestWithSystemSchemaAsDefaultKeyspace(t *testing.T) {
 	// We are testing this separately so we can set a default keyspace
 	vschema := &vschemaWrapper{
diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases.txt b/go/vt/vtgate/planbuilder/testdata/select_cases.txt
index 652ad343b9..00bd87869a 100644
--- a/go/vt/vtgate/planbuilder/testdata/select_cases.txt
+++ b/go/vt/vtgate/planbuilder/testdata/select_cases.txt
@@ -2195,8 +2195,8 @@ Gen4 plan same as above
     "OperatorType": "Route",
     "Variant": "Scatter",
     "Keyspace": {
-      "Name": "main",
-      "Sharded": false
+      "Name": "user",
+      "Sharded": true
     },
     "FieldQuery": "select 42, id from dual, `user` where 1 != 1",
     "Query": "select 42, id from dual, `user`",
@@ -3110,3 +3110,65 @@ Gen4 plan same as above
     ]
   }
 }
+
+# yeah, it does not make sense, but it's valid
+"select exists(select 1) from user where id = 5"
+{
+  "QueryType": "SELECT",
+  "Original": "select exists(select 1) from user where id = 5",
+  "Instructions": {
+    "OperatorType": "Subquery",
+    "Variant": "PulloutExists",
+    "PulloutVars": [
+      "__sq_has_values1",
+      "__sq1"
+    ],
+    "Inputs": [
+      {
+        "OperatorType": "Route",
+        "Variant": "Reference",
+        "Keyspace": {
+          "Name": "main",
+          "Sharded": false
+        },
+        "FieldQuery": "select 1 from dual where 1 != 1",
+        "Query": "select 1 from dual",
+        "Table": "dual"
+      },
+      {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select :__sq_has_values1 from `user` where 1 != 1",
+        "Query": "select :__sq_has_values1 from `user` where id = 5",
+        "Table": "`user`",
+        "Values": [
+          "INT64(5)"
+        ],
+        "Vindex": "user_index"
+      }
+    ]
+  }
+}
+{
+  "QueryType": "SELECT",
+  "Original": "select exists(select 1) from user where id = 5",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "EqualUnique",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "FieldQuery": "select exists (select 1 from dual where 1 != 1) from `user` where 1 != 1",
+    "Query": "select exists (select 1 from dual) from `user` where id = 5",
+    "Table": "`user`",
+    "Values": [
+      "INT64(5)"
+    ],
+    "Vindex": "user_index"
+  }
+}
diff --git a/go/vt/vtgate/planbuilder/testdata/select_cases_with_default.txt b/go/vt/vtgate/planbuilder/testdata/select_cases_with_default.txt
new file mode 100644
index 0000000000..057916f9a6
--- /dev/null
+++ b/go/vt/vtgate/planbuilder/testdata/select_cases_with_default.txt
@@ -0,0 +1,61 @@
+# EXISTS subquery when the default ks is different than the inner query
+"select exists(select * from user where id = 5)"
+{
+  "QueryType": "SELECT",
+  "Original": "select exists(select * from user where id = 5)",
+  "Instructions": {
+    "OperatorType": "Subquery",
+    "Variant": "PulloutExists",
+    "PulloutVars": [
+      "__sq_has_values1",
+      "__sq1"
+    ],
+    "Inputs": [
+      {
+        "OperatorType": "Route",
+        "Variant": "EqualUnique",
+        "Keyspace": {
+          "Name": "user",
+          "Sharded": true
+        },
+        "FieldQuery": "select * from `user` where 1 != 1",
+        "Query": "select * from `user` where id = 5",
+        "Table": "`user`",
+        "Values": [
+          "INT64(5)"
+        ],
+        "Vindex": "user_index"
+      },
+      {
+        "OperatorType": "Route",
+        "Variant": "Reference",
+        "Keyspace": {
+          "Name": "second_user",
+          "Sharded": true
+        },
+        "FieldQuery": "select :__sq_has_values1 from dual where 1 != 1",
+        "Query": "select :__sq_has_values1 from dual",
+        "Table": "dual"
+      }
+    ]
+  }
+}
+{
+  "QueryType": "SELECT",
+  "Original": "select exists(select * from user where id = 5)",
+  "Instructions": {
+    "OperatorType": "Route",
+    "Variant": "EqualUnique",
+    "Keyspace": {
+      "Name": "user",
+      "Sharded": true
+    },
+    "FieldQuery": "select exists (select * from `user` where 1 != 1) from dual where 1 != 1",
+    "Query": "select exists (select * from `user` where id = 5) from dual",
+    "Table": "dual",
+    "Values": [
+      "INT64(5)"
+    ],
+    "Vindex": "user_index"
+  }
+}
diff --git a/go/vt/vttablet/endtoend/connkilling/connkiller_test.go b/go/vt/vttablet/endtoend/connkilling/connkiller_test.go
index e9a0e51cea..80c913035b 100644
--- a/go/vt/vttablet/endtoend/connkilling/connkiller_test.go
+++ b/go/vt/vttablet/endtoend/connkilling/connkiller_test.go
@@ -20,6 +20,8 @@ All tests in this package come with a three second time out for OLTP session
 package connkilling
 
 import (
+	"context"
+	"strings"
 	"testing"
 	"time"
 
@@ -35,7 +37,10 @@ func TestTxKillerKillsTransactionsInReservedConnections(t *testing.T) {
 	_, err := client.ReserveBeginExecute("select 42", nil, nil, nil)
 	require.NoError(t, err)
 
-	assertIsKilledWithin6Seconds(t, client)
+	assertIsKilledWithin(t, client, queryPollOpts{
+		StopAfter: 6 * time.Second,
+		PollEvery: time.Second,
+	})
 }
 
 func TestTxKillerDoesNotKillReservedConnectionsInUse(t *testing.T) {
@@ -73,7 +78,10 @@ func TestTxKillerCountsTimeFromTxStartedNotStatefulConnCreated(t *testing.T) {
 	_, err = client.Execute("select 43", nil)
 	require.NoError(t, err)
 
-	assertIsKilledWithin6Seconds(t, client)
+	assertIsKilledWithin(t, client, queryPollOpts{
+		StopAfter: 6 * time.Second,
+		PollEvery: time.Second,
+	})
 }
 
 func TestTxKillerKillsTransactionThreeSecondsAfterCreation(t *testing.T) {
@@ -83,7 +91,10 @@ func TestTxKillerKillsTransactionThreeSecondsAfterCreation(t *testing.T) {
 	_, err := client.BeginExecute("select 42", nil, nil)
 	require.NoError(t, err)
 
-	assertIsKilledWithin6Seconds(t, client)
+	assertIsKilledWithin(t, client, queryPollOpts{
+		StopAfter: 6 * time.Second,
+		PollEvery: time.Second,
+	})
 }
 
 func assertIsNotKilledOver5Second(t *testing.T, client *framework.QueryClient) {
@@ -94,18 +105,59 @@ func assertIsNotKilledOver5Second(t *testing.T, client *framework.QueryClient) {
 	}
 }
 
-func assertIsKilledWithin6Seconds(t *testing.T, client *framework.QueryClient) {
+type queryPollOpts struct {
+	StopAfter time.Duration
+	PollEvery time.Duration
+}
+
+func assertIsKilledWithin(t testing.TB, client *framework.QueryClient, opts queryPollOpts) {
+	t.Helper()
+
 	var err error
-	// when it is used once per second
-	for i := 0; i < 6; i++ {
+
+	doQuery := func() {
 		_, err = client.Execute("select 43", nil)
 		if err != nil {
-			break
+			if strings.Contains(err.Error(), "in use: for tx killer rollback") {
+				t.Logf("tx is mid-rollback and not killed yet, ignoring this error: %s", err)
+				err = nil
+			}
 		}
-		time.Sleep(1 * time.Second)
 	}
 
-	// then it should still be killed. transactions are tracked per tx-creation time and not last-used time
-	require.Error(t, err)
-	require.Contains(t, err.Error(), "exceeded timeout: 3s")
+	ctx, cancel := context.WithTimeout(context.Background(), opts.StopAfter)
+	defer cancel()
+
+	defer func() {
+		// Check one last time after we've hit StopAfter
+		if err == nil {
+			doQuery()
+		}
+
+		// then it should still be killed. transactions are tracked per tx-creation time and not last-used time
+		require.Error(t, err)
+		require.Contains(t, err.Error(), "exceeded timeout: 3s")
+	}()
+
+	ticker := time.NewTicker(opts.PollEvery)
+	defer func() {
+		ticker.Stop()
+		select {
+		case <-ticker.C:
+		default:
+		}
+	}()
+
+poll:
+	for {
+		select {
+		case <-ctx.Done():
+			break poll
+		case <-ticker.C:
+			doQuery()
+			if err != nil {
+				break poll
+			}
+		}
+	}
 }