Reduce network pressure on multi row insert (#14064)

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

go/test/endtoend/cluster/cluster_util.go

@@ -127,7 +127,7 @@ func VerifyRowsInTablet(t *testing.T, vttablet *Vttablet, ksName string, expecte
 }
 
 // PanicHandler handles the panic in the testcase.
-func PanicHandler(t *testing.T) {
+func PanicHandler(t testing.TB) {
 	err := recover()
 	if t == nil {
 		return

go/test/endtoend/utils/utils.go

@@ -169,7 +169,7 @@ func ExecCompareMySQL(t *testing.T, vtConn, mysqlConn *mysql.Conn, query string)
 
 // ExecAllowError executes the given query without failing the test if it produces
 // an error. The error is returned to the client, along with the result set.
-func ExecAllowError(t *testing.T, conn *mysql.Conn, query string) (*sqltypes.Result, error) {
+func ExecAllowError(t testing.TB, conn *mysql.Conn, query string) (*sqltypes.Result, error) {
 	t.Helper()
 	return conn.ExecuteFetch(query, 1000, true)
 }

go/test/endtoend/vtgate/queries/benchmark/benchmark_test.go

@@ -0,0 +1,80 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package dml
+
+import (
+	"fmt"
+	"math/rand"
+	"strconv"
+	"strings"
+	"testing"
+
+	"vitess.io/vitess/go/test/endtoend/utils"
+)
+
+type testQuery struct {
+	tableName string
+	cols      []string
+	intTyp    []bool
+}
+
+func (tq *testQuery) getInsertQuery(rows int) string {
+	var allRows []string
+	for i := 0; i < rows; i++ {
+		var row []string
+		for _, isInt := range tq.intTyp {
+			if isInt {
+				row = append(row, strconv.Itoa(i))
+				continue
+			}
+			row = append(row, "'"+getRandomString(50)+"'")
+		}
+		allRows = append(allRows, "("+strings.Join(row, ",")+")")
+	}
+	return fmt.Sprintf("insert into %s(%s) values %s", tq.tableName, strings.Join(tq.cols, ","), strings.Join(allRows, ","))
+}
+
+func getRandomString(size int) string {
+	var str strings.Builder
+
+	for i := 0; i < size; i++ {
+		str.WriteByte(byte(rand.Intn(27) + 97))
+	}
+	return str.String()
+}
+
+func BenchmarkShardedTblNoLookup(b *testing.B) {
+	conn, closer := start(b)
+	defer closer()
+
+	cols := []string{"id", "c1", "c2", "c3", "c4", "c5", "c6", "c7", "c8", "c9", "c10", "c11", "c12"}
+	intType := make([]bool, len(cols))
+	intType[0] = true
+	tq := &testQuery{
+		tableName: "tbl_no_lkp_vdx",
+		cols:      cols,
+		intTyp:    intType,
+	}
+	for _, rows := range []int{1, 10, 100, 500, 1000, 5000, 10000} {
+		insStmt := tq.getInsertQuery(rows)
+		b.Run(fmt.Sprintf("16-shards-%d-rows", rows), func(b *testing.B) {
+			for i := 0; i < b.N; i++ {
+				_ = utils.Exec(b, conn, insStmt)
+			}
+		})
+	}
+}

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

@@ -0,0 +1,124 @@
+/*
+Copyright 2023 The Vitess Authors.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package dml
+
+import (
+	"context"
+	_ "embed"
+	"flag"
+	"os"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+
+	"vitess.io/vitess/go/mysql"
+	"vitess.io/vitess/go/test/endtoend/cluster"
+	"vitess.io/vitess/go/test/endtoend/utils"
+)
+
+var (
+	clusterInstance *cluster.LocalProcessCluster
+	vtParams        mysql.ConnParams
+	mysqlParams     mysql.ConnParams
+	sKs             = "sks"
+	uKs             = "uks"
+	cell            = "test"
+
+	//go:embed sharded_schema.sql
+	sSchemaSQL string
+
+	//go:embed vschema.json
+	sVSchema string
+)
+
+var (
+	shards4 = []string{
+		"-40", "40-80", "80-c0", "c0-",
+	}
+
+	shards8 = []string{
+		"-20", "20-40", "40-60", "60-80", "80-a0", "a0-c0", "c0-e0", "e0-",
+	}
+
+	shards16 = []string{
+		"-10", "10-20", "20-30", "30-40", "40-50", "50-60", "60-70", "70-80", "80-90", "90-a0", "a0-b0", "b0-c0", "c0-d0", "d0-e0", "e0-f0", "f0-",
+	}
+
+	shards32 = []string{
+		"-05", "05-10", "10-15", "15-20", "20-25", "25-30", "30-35", "35-40", "40-45", "45-50", "50-55", "55-60", "60-65", "65-70", "70-75", "75-80",
+		"80-85", "85-90", "90-95", "95-a0", "a0-a5", "a5-b0", "b0-b5", "b5-c0", "c0-c5", "c5-d0", "d0-d5", "d5-e0", "e0-e5", "e5-f0", "f0-f5", "f5-",
+	}
+)
+
+func TestMain(m *testing.M) {
+	defer cluster.PanicHandler(nil)
+	flag.Parse()
+
+	exitCode := func() int {
+		clusterInstance = cluster.NewCluster(cell, "localhost")
+		defer clusterInstance.Teardown()
+
+		// Start topo server
+		err := clusterInstance.StartTopo()
+		if err != nil {
+			return 1
+		}
+
+		// Start sharded keyspace
+		sKeyspace := &cluster.Keyspace{
+			Name:      sKs,
+			SchemaSQL: sSchemaSQL,
+			VSchema:   sVSchema,
+		}
+
+		err = clusterInstance.StartKeyspace(*sKeyspace, shards4, 0, false)
+		if err != nil {
+			return 1
+		}
+
+		// Start vtgate
+		err = clusterInstance.StartVtgate()
+		if err != nil {
+			return 1
+		}
+
+		vtParams = clusterInstance.GetVTParams(sKs)
+
+		return m.Run()
+	}()
+	os.Exit(exitCode)
+}
+
+func start(b *testing.B) (*mysql.Conn, func()) {
+	conn, err := mysql.Connect(context.Background(), &vtParams)
+	require.NoError(b, err)
+
+	deleteAll := func() {
+		tables := []string{"tbl_no_lkp_vdx"}
+		for _, table := range tables {
+			_, _ = utils.ExecAllowError(b, conn, "delete from "+table)
+		}
+	}
+
+	deleteAll()
+
+	return conn, func() {
+		deleteAll()
+		conn.Close()
+		cluster.PanicHandler(b)
+	}
+}

go/test/endtoend/vtgate/queries/benchmark/sharded_schema.sql

@@ -0,0 +1,16 @@
+create table tbl_no_lkp_vdx
+(
+    id  bigint,
+    c1 varchar(50),
+    c2 varchar(50),
+    c3 varchar(50),
+    c4 varchar(50),
+    c5 varchar(50),
+    c6 varchar(50),
+    c7 varchar(50),
+    c8 varchar(50),
+    c9 varchar(50),
+    c10 varchar(50),
+    c11 varchar(50),
+    c12 varchar(50)
+) Engine = InnoDB;

go/test/endtoend/vtgate/queries/benchmark/vschema.json

@@ -0,0 +1,18 @@
+{
+  "sharded": true,
+  "vindexes": {
+    "xxhash": {
+      "type": "xxhash"
+    }
+  },
+  "tables": {
+    "tbl_no_lkp_vdx": {
+      "column_vindexes": [
+        {
+          "column": "id",
+          "name": "xxhash"
+        }
+      ]
+    }
+  }
+}

go/vt/vtgate/autocommit_test.go

@@ -274,7 +274,6 @@ func TestAutocommitInsertLookup(t *testing.T) {
 		BindVariables: map[string]*querypb.BindVariable{
 			"_Id_0":   sqltypes.Int64BindVariable(1),
 			"_name_0": sqltypes.StringBindVariable("myname"),
-			"__seq0":  sqltypes.Int64BindVariable(1),
 		},
 	}})
 	testCommitCount(t, "sbc1", sbc1, 1)
@@ -292,15 +291,13 @@ func TestAutocommitInsertMultishardAutoCommit(t *testing.T) {
 			Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_0, 2)",
 			BindVariables: map[string]*querypb.BindVariable{
 				"_user_id_0": sqltypes.Int64BindVariable(1),
-				"_user_id_1": sqltypes.Int64BindVariable(3),
 			},
 		}})
 		testCommitCount(t, "sbc1", sbc1, 0)
 
 		assertQueries(t, sbc2, []*querypb.BoundQuery{{
 			Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_1, 4)",
 			BindVariables: map[string]*querypb.BindVariable{
-				"_user_id_0": sqltypes.Int64BindVariable(1),
 				"_user_id_1": sqltypes.Int64BindVariable(3),
 			},
 		}})
@@ -321,7 +318,6 @@ func TestAutocommitInsertMultishardAutoCommit(t *testing.T) {
 		assertQueries(t, sbc2, []*querypb.BoundQuery{{
 			Sql: "insert /*vt+ MULTI_SHARD_AUTOCOMMIT=1 */ into user_extra(user_id, v) values (:_user_id_1, 4)",
 			BindVariables: map[string]*querypb.BindVariable{
-				"_user_id_0": sqltypes.Int64BindVariable(1),
 				"_user_id_1": sqltypes.Int64BindVariable(3),
 			},
 		}})
@@ -339,15 +335,13 @@ func TestAutocommitInsertMultishard(t *testing.T) {
 		Sql: "insert into user_extra(user_id, v) values (:_user_id_0, 2)",
 		BindVariables: map[string]*querypb.BindVariable{
 			"_user_id_0": sqltypes.Int64BindVariable(1),
-			"_user_id_1": sqltypes.Int64BindVariable(3),
 		},
 	}})
 	testCommitCount(t, "sbc1", sbc1, 1)
 
 	assertQueries(t, sbc2, []*querypb.BoundQuery{{
 		Sql: "insert into user_extra(user_id, v) values (:_user_id_1, 4)",
 		BindVariables: map[string]*querypb.BindVariable{
-			"_user_id_0": sqltypes.Int64BindVariable(1),
 			"_user_id_1": sqltypes.Int64BindVariable(3),
 		},
 	}})

go/vt/vtgate/engine/insert.go

@@ -25,6 +25,7 @@ import (
 	"sync"
 	"time"
 
+	"vitess.io/vitess/go/slice"
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/key"
 	querypb "vitess.io/vitess/go/vt/proto/query"
@@ -75,7 +76,7 @@ type (
 
 		// Prefix, Mid and Suffix are for sharded insert plans.
 		Prefix string
-		Mid    []string
+		Mid    sqlparser.Values
 		Suffix string
 
 		// Option to override the standard behavior and allow a multi-shard insert
@@ -132,7 +133,7 @@ func NewInsert(
 	vindexValues [][][]evalengine.Expr,
 	table *vindexes.Table,
 	prefix string,
-	mid []string,
+	mid sqlparser.Values,
 	suffix string,
 ) *Insert {
 	ins := &Insert{
@@ -758,17 +759,23 @@ func (ins *Insert) getInsertShardedRoute(
 
 	queries := make([]*querypb.BoundQuery, len(rss))
 	for i := range rss {
+		shardBindVars := map[string]*querypb.BindVariable{}
 		var mids []string
 		for _, indexValue := range indexesPerRss[i] {
 			index, _ := strconv.ParseInt(string(indexValue.Value), 0, 64)
 			if keyspaceIDs[index] != nil {
-				mids = append(mids, ins.Mid[index])
+				mids = append(mids, sqlparser.String(ins.Mid[index]))
+				for _, expr := range ins.Mid[index] {
+					if arg, ok := expr.(*sqlparser.Argument); ok {
+						shardBindVars[arg.Name] = bindVars[arg.Name]
+					}
+				}
 			}
 		}
 		rewritten := ins.Prefix + strings.Join(mids, ",") + ins.Suffix
 		queries[i] = &querypb.BoundQuery{
 			Sql:           rewritten,
-			BindVariables: bindVars,
+			BindVariables: shardBindVars,
 		}
 	}
 
@@ -985,7 +992,10 @@ func (ins *Insert) description() PrimitiveDescription {
 		other["VindexOffsetFromSelect"] = valuesOffsets
 	}
 	if len(ins.Mid) > 0 {
-		shardQuery := fmt.Sprintf("%s%s%s", ins.Prefix, strings.Join(ins.Mid, ", "), ins.Suffix)
+		mids := slice.Map(ins.Mid, func(from sqlparser.ValTuple) string {
+			return sqlparser.String(from)
+		})
+		shardQuery := fmt.Sprintf("%s%s%s", ins.Prefix, strings.Join(mids, ", "), ins.Suffix)
 		if shardQuery != ins.Query {
 			other["ShardedQuery"] = shardQuery
 		}