feat: authoritative query timeout for vttablet from vtgate (#16735)

Signed-off-by: Harshit Gangal <[email protected]>

changelog/21.0/21.0.0/summary.md

@@ -14,6 +14,7 @@
   - **[Allow Cross Cell Promotion in PRS](#allow-cross-cell)**
   - **[Support for recursive CTEs](#recursive-cte)**
   - **[VTGate Tablet Balancer](#tablet-balancer)**
+  - **[Query Timeout Override](#query-timeout)**
 
 ## <a id="major-changes"/>Major Changes
 
@@ -113,4 +114,19 @@ When a VTGate routes a query and has multiple available tablets for a given shar
 
 The tablet balancer is enabled by a new flag `--enable-balancer` and configured by `--balancer-vtgate-cells` and `--balancer-keyspaces`.
 
-See [RFC for details](https://github.com/vitessio/vitess/issues/12241).
+See [RFC for details](https://github.com/vitessio/vitess/issues/12241).
+
+### <a id="query-timeout"/>Query Timeout Override
+VTGate sends an authoritative query timeout to VTTablet when the `QUERY_TIMEOUT_MS` comment directive, 
+`query_timeout` session system variable, or `query-timeout` flag is set. 
+The order of precedence is: `QUERY_TIMEOUT_MS` > `query_timeout` > `query-timeout`. 
+VTTablet overrides its default query timeout with the value received from VTGate. 
+All timeouts are specified in milliseconds.
+
+When a query is executed inside a transaction, this behavior does not apply; instead, 
+the smaller of the transaction timeout or the query timeout from VTGate is used.
+
+A query can also be set to have no timeout by using the `QUERY_TIMEOUT_MS` comment directive with a value of `0`.
+
+Example usage:
+`select /*vt+ QUERY_TIMEOUT_MS=30 */ col from tbl`

go/test/endtoend/vtgate/queries/timeout/main_test.go

@@ -63,8 +63,9 @@ func TestMain(m *testing.M) {
 
 		clusterInstance.VtTabletExtraArgs = append(clusterInstance.VtTabletExtraArgs,
 			"--queryserver-config-max-result-size", "1000000",
-			"--queryserver-config-query-timeout", "200s",
-			"--queryserver-config-query-pool-timeout", "200s")
+			"--queryserver-config-query-timeout", "2s",
+			"--queryserver-config-transaction-timeout", "3s",
+			"--queryserver-config-query-pool-timeout", "2s")
 		// Start Unsharded keyspace
 		ukeyspace := &cluster.Keyspace{
 			Name:      uks,

go/test/endtoend/vtgate/queries/timeout/timeout_test.go

@@ -17,11 +17,13 @@ limitations under the License.
 package misc
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 
+	"vitess.io/vitess/go/mysql"
 	"vitess.io/vitess/go/test/endtoend/cluster"
 	"vitess.io/vitess/go/test/endtoend/utils"
 )
@@ -66,6 +68,9 @@ func TestQueryTimeoutWithDual(t *testing.T) {
 	assert.Error(t, err)
 	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=15 */ sleep(0.001) from dual")
 	assert.NoError(t, err)
+	// infinite query timeout overriding all defaults
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=0 */ sleep(5) from dual")
+	assert.NoError(t, err)
 }
 
 func TestQueryTimeoutWithTables(t *testing.T) {
@@ -124,3 +129,53 @@ func TestQueryTimeoutWithShardTargeting(t *testing.T) {
 		})
 	}
 }
+
+func TestQueryTimeoutWithoutVTGateDefault(t *testing.T) {
+	// disable query timeout
+	clusterInstance.VtGateExtraArgs = append(clusterInstance.VtGateExtraArgs,
+		"--query-timeout", "0")
+	require.NoError(t,
+		clusterInstance.RestartVtgate())
+
+	// update vtgate params
+	vtParams = clusterInstance.GetVTParams(keyspaceName)
+
+	mcmp, closer := start(t)
+	defer closer()
+
+	// tablet query timeout of 2s
+	_, err := utils.ExecAllowError(t, mcmp.VtConn, "select sleep(5) from dual")
+	assert.Error(t, err)
+
+	// infinite timeout using query hint
+	utils.Exec(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=0 */ sleep(5) from dual")
+
+	// checking again without query hint, tablet query timeout of 2s should be applied
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select sleep(5) from dual")
+	assert.Error(t, err)
+
+	// set timeout of 20ms
+	utils.Exec(t, mcmp.VtConn, "set query_timeout=20")
+
+	// query timeout of 20ms should be applied
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select sleep(1) from dual")
+	assert.Error(t, err)
+
+	// infinite timeout using query hint will override session timeout.
+	utils.Exec(t, mcmp.VtConn, "select /*vt+ QUERY_TIMEOUT_MS=0 */ sleep(5) from dual")
+
+	// open second session
+	conn2, err := mysql.Connect(context.Background(), &vtParams)
+	require.NoError(t, err)
+	defer conn2.Close()
+
+	// tablet query timeout of 2s should be applied, as session timeout is not set on this connection.
+	utils.Exec(t, conn2, "select sleep(1) from dual")
+	_, err = utils.ExecAllowError(t, conn2, "select sleep(5) from dual")
+	assert.Error(t, err)
+
+	// reset session on first connection, tablet query timeout of 2s should be applied.
+	utils.Exec(t, mcmp.VtConn, "set query_timeout=0")
+	_, err = utils.ExecAllowError(t, mcmp.VtConn, "select sleep(5) from dual")
+	assert.Error(t, err)
+}