proto: enable pooling for vreplication

Signed-off-by: Vicent Marti <[email protected]>
dkua · Jun 7, 2021 · 5b5ccd8 · 5b5ccd8
1 parent f051584
commit 5b5ccd8
Show file tree

Hide file tree

Showing 6 changed files with 91 additions and 14 deletions.
diff --git a/Makefile b/Makefile
@@ -232,7 +232,9 @@ $(PROTO_GO_OUTS): minimaltools install_protoc-gen-go proto/*.proto
 		--go_out=. --plugin protoc-gen-go="${GOBIN}/protoc-gen-go" \
 		--go-grpc_out=. --plugin protoc-gen-go-grpc="${GOBIN}/protoc-gen-go-grpc" \
 		--go-vtproto_out=. --plugin protoc-gen-go-vtproto="${GOBIN}/protoc-gen-go-vtproto" \
-		--go-vtproto_opt=features=marshal+unmarshal+size \
+		--go-vtproto_opt=features=marshal+unmarshal+size+pool \
+		--go-vtproto_opt=pool=vitess.io/vitess/go/vt/proto/query.Row \
+		--go-vtproto_opt=pool=vitess.io/vitess/go/vt/proto/binlogdata.VStreamRowsResponse \
 		-I${PWD}/dist/vt-protoc-3.6.1/include:proto proto/$${name}.proto; \
 	done
 	cp -Rf vitess.io/vitess/go/vt/proto/* go/vt/proto

diff --git a/go/vt/proto/binlogdata/binlogdata_vtproto.pb.go b/go/vt/proto/binlogdata/binlogdata_vtproto.pb.go
diff --git a/go/vt/proto/query/query_vtproto.pb.go b/go/vt/proto/query/query_vtproto.pb.go
diff --git a/go/vt/vttablet/grpcqueryservice/server.go b/go/vt/vttablet/grpcqueryservice/server.go
@@ -17,10 +17,10 @@ limitations under the License.
 package grpcqueryservice
 
 import (
-	"google.golang.org/grpc"
-
 	"context"
 
+	"google.golang.org/grpc"
+
 	"vitess.io/vitess/go/sqltypes"
 	"vitess.io/vitess/go/vt/callerid"
 	"vitess.io/vitess/go/vt/callinfo"

diff --git a/go/vt/vttablet/grpctabletconn/conn.go b/go/vt/vttablet/grpctabletconn/conn.go
@@ -675,18 +675,18 @@ func (conn *gRPCQueryClient) VStreamRows(ctx context.Context, target *querypb.Ta
 		return err
 	}
 	for {
-		r, err := stream.Recv()
+		r := binlogdatapb.VStreamRowsResponseFromVTPool()
+		err := stream.RecvMsg(r)
 		if err != nil {
 			return tabletconn.ErrorFromGRPC(err)
 		}
-		select {
-		case <-ctx.Done():
+		if ctx.Err() != nil {
 			return ctx.Err()
-		default:
 		}
 		if err := send(r); err != nil {
 			return err
 		}
+		r.ReturnToVTPool()
 	}
 }
 

diff --git a/go/vt/vttablet/tabletmanager/vreplication/vcopier.go b/go/vt/vttablet/tabletmanager/vreplication/vcopier.go
@@ -244,13 +244,13 @@ func (vc *vcopier) copyTable(ctx context.Context, tableName string, copyState ma
 			}
 			fieldEvent := &binlogdatapb.FieldEvent{
 				TableName: initialPlan.SendRule.Match,
-				Fields:    rows.Fields,
 			}
+			fieldEvent.Fields = append(fieldEvent.Fields, rows.Fields...)
 			vc.tablePlan, err = plan.buildExecutionPlan(fieldEvent)
 			if err != nil {
 				return err
 			}
-			pkfields = rows.Pkfields
+			pkfields = append(pkfields, rows.Pkfields...)
 			buf := sqlparser.NewTrackedBuffer(nil)
 			buf.Myprintf("update _vt.copy_state set lastpk=%a where vrepl_id=%s and table_name=%s", ":lastpk", strconv.Itoa(int(vc.vr.id)), encodeString(tableName))
 			updateCopyState = buf.ParsedQuery()