backend, net: reduce memory allocation in forwarding packets (#394)

Co-authored-by: xhe <[email protected]>

Author: djshow832

pkg/proxy/backend/backend_conn_mgr.go

@@ -353,7 +353,7 @@ func (mgr *BackendConnManager) initSessionStates(backendIO *pnet.PacketIO, sessi
 
 func (mgr *BackendConnManager) querySessionStates(backendIO *pnet.PacketIO) (sessionStates, sessionToken string, err error) {
 	// Do not lock here because the caller already locks.
-	var result *gomysql.Result
+	var result *gomysql.Resultset
 	if result, _, err = mgr.cmdProcessor.query(backendIO, sqlQueryState); err != nil {
 		return
 	}

pkg/proxy/backend/cmd_processor.go

@@ -37,10 +37,10 @@ func NewCmdProcessor(logger *zap.Logger) *CmdProcessor {
 	}
 }
 
-func (cp *CmdProcessor) handleOKPacket(request, response []byte) *gomysql.Result {
-	r := pnet.ParseOKPacket(response)
-	cp.updateServerStatus(request, r.Status)
-	return r
+func (cp *CmdProcessor) handleOKPacket(request, response []byte) uint16 {
+	status := pnet.ParseOKPacket(response)
+	cp.updateServerStatus(request, status)
+	return status
 }
 
 func (cp *CmdProcessor) handleErrorPacket(data []byte) error {

pkg/proxy/backend/cmd_processor_exec.go

@@ -99,14 +99,14 @@ func forwardOnePacket(destIO, srcIO *pnet.PacketIO, flush bool) (data []byte, er
 
 func (cp *CmdProcessor) forwardUntilResultEnd(clientIO, backendIO *pnet.PacketIO, request []byte) (uint16, error) {
 	var serverStatus uint16
-	err := backendIO.ForwardUntil(clientIO, func(firstByte byte, length int) bool {
+	err := backendIO.ForwardUntil(clientIO, func(firstByte byte, length int) (end, needData bool) {
 		switch {
 		case pnet.IsErrorPacket(firstByte):
-			return true
+			return true, true
 		case cp.capability&pnet.ClientDeprecateEOF == 0:
-			return pnet.IsEOFPacket(firstByte, length)
+			return pnet.IsEOFPacket(firstByte, length), true
 		default:
-			return pnet.IsResultSetOKPacket(firstByte, length)
+			return pnet.IsResultSetOKPacket(firstByte, length), true
 		}
 	}, func(response []byte) error {
 		switch {
@@ -119,7 +119,7 @@ func (cp *CmdProcessor) forwardUntilResultEnd(clientIO, backendIO *pnet.PacketIO
 			serverStatus = cp.handleEOFPacket(request, response)
 			return clientIO.Flush()
 		default:
-			serverStatus = cp.handleOKPacket(request, response).Status
+			serverStatus = cp.handleOKPacket(request, response)
 			return clientIO.Flush()
 		}
 	})
@@ -146,9 +146,14 @@ func (cp *CmdProcessor) forwardPrepareCmd(clientIO, backendIO *pnet.PacketIO) er
 				expectedPackets++
 			}
 		}
-		for i := 0; i < expectedPackets; i++ {
-			// Ignore this status because PREPARE doesn't affect status.
-			if _, err = forwardOnePacket(clientIO, backendIO, false); err != nil {
+		// Ignore this status because PREPARE doesn't affect status.
+		if expectedPackets > 0 {
+			i := 0
+			err = backendIO.ForwardUntil(clientIO, func(firstByte byte, firstPktLen int) (end, needData bool) {
+				i++
+				return i >= expectedPackets, false
+			}, nil)
+			if err != nil {
 				return err
 			}
 		}
@@ -175,26 +180,35 @@ func (cp *CmdProcessor) forwardFieldListCmd(clientIO, backendIO *pnet.PacketIO,
 
 func (cp *CmdProcessor) forwardQueryCmd(clientIO, backendIO *pnet.PacketIO, request []byte) error {
 	for {
-		response, err := forwardOnePacket(clientIO, backendIO, false)
-		if err != nil {
-			return err
-		}
 		var serverStatus uint16
-		switch response[0] {
-		case mysql.OKHeader:
-			rs := cp.handleOKPacket(request, response)
-			serverStatus, err = rs.Status, clientIO.Flush()
-		case mysql.ErrHeader:
-			if err := clientIO.Flush(); err != nil {
-				return err
+		var first byte
+		err := backendIO.ForwardUntil(clientIO, func(firstByte byte, _ int) (end, needData bool) {
+			first = firstByte
+			switch firstByte {
+			case mysql.OKHeader, mysql.ErrHeader:
+				return true, true
+			default:
+				return true, false
 			}
-			// Subsequent statements won't be executed even if it's a multi-statement.
-			return cp.handleErrorPacket(response)
-		case mysql.LocalInFileHeader:
-			serverStatus, err = cp.forwardLoadInFile(clientIO, backendIO, request)
-		default:
-			serverStatus, err = cp.forwardResultSet(clientIO, backendIO, request)
-		}
+		}, func(response []byte) error {
+			var err error
+			switch first {
+			case mysql.OKHeader:
+				status := cp.handleOKPacket(request, response)
+				serverStatus, err = status, clientIO.Flush()
+			case mysql.ErrHeader:
+				if err = clientIO.Flush(); err != nil {
+					return err
+				}
+				// Subsequent statements won't be executed even if it's a multi-statement.
+				return cp.handleErrorPacket(response)
+			case mysql.LocalInFileHeader:
+				serverStatus, err = cp.forwardLoadInFile(clientIO, backendIO, request)
+			default:
+				serverStatus, err = cp.forwardResultSet(clientIO, backendIO, request)
+			}
+			return err
+		})
 		if err != nil {
 			return err
 		}
@@ -213,11 +227,14 @@ func (cp *CmdProcessor) forwardLoadInFile(clientIO, backendIO *pnet.PacketIO, re
 	// The client sends file data until an empty packet.
 	for {
 		var data []byte
-		// The file may be large, so always flush it.
-		if data, err = forwardOnePacket(backendIO, clientIO, true); err != nil {
+		// Do not call PacketIO.ForwardUntil. It peeks 5 bytes but there may be only 4 bytes here.
+		if data, err = forwardOnePacket(backendIO, clientIO, false); err != nil {
 			return
 		}
 		if len(data) == 0 {
+			if err := backendIO.Flush(); err != nil {
+				return 0, err
+			}
 			break
 		}
 	}
@@ -227,8 +244,7 @@ func (cp *CmdProcessor) forwardLoadInFile(clientIO, backendIO *pnet.PacketIO, re
 	}
 	switch response[0] {
 	case mysql.OKHeader:
-		rs := cp.handleOKPacket(request, response)
-		return rs.Status, nil
+		return cp.handleOKPacket(request, response), nil
 	case mysql.ErrHeader:
 		return serverStatus, cp.handleErrorPacket(response)
 	}
@@ -238,22 +254,22 @@ func (cp *CmdProcessor) forwardLoadInFile(clientIO, backendIO *pnet.PacketIO, re
 
 func (cp *CmdProcessor) forwardResultSet(clientIO, backendIO *pnet.PacketIO, request []byte) (uint16, error) {
 	if cp.capability&pnet.ClientDeprecateEOF == 0 {
-		var response []byte
+		var serverStatus uint16
 		// read columns
-		for {
-			var err error
-			if response, err = forwardOnePacket(clientIO, backendIO, false); err != nil {
-				return 0, err
-			}
-			if pnet.IsEOFPacket(response[0], len(response)) {
-				break
+		err := backendIO.ForwardUntil(clientIO, func(firstByte byte, firstPktLen int) (end, needData bool) {
+			return pnet.IsEOFPacket(firstByte, firstPktLen), true
+		}, func(response []byte) error {
+			serverStatus = binary.LittleEndian.Uint16(response[3:])
+			// If a cursor exists, only columns are sent this time. The client will then send COM_STMT_FETCH to fetch rows.
+			// Otherwise, columns and rows are both sent once.
+			if serverStatus&mysql.ServerStatusCursorExists > 0 {
+				serverStatus = cp.handleEOFPacket(request, response)
+				return clientIO.Flush()
 			}
-		}
-		serverStatus := binary.LittleEndian.Uint16(response[3:])
-		// If a cursor exists, only columns are sent this time. The client will then send COM_STMT_FETCH to fetch rows.
-		// Otherwise, columns and rows are both sent once.
-		if serverStatus&mysql.ServerStatusCursorExists > 0 {
-			return cp.handleEOFPacket(request, response), clientIO.Flush()
+			return nil
+		})
+		if err != nil {
+			return serverStatus, err
 		}
 	}
 	// Deprecate EOF or no cursor.

pkg/proxy/backend/cmd_processor_query.go

@@ -16,7 +16,7 @@ import (
 // query is called when the proxy sends requests to the backend by itself,
 // such as querying session states, committing the current transaction.
 // It only supports limited cases, excluding loading file, cursor fetch, multi-statements, etc.
-func (cp *CmdProcessor) query(packetIO *pnet.PacketIO, sql string) (result *gomysql.Result, response []byte, err error) {
+func (cp *CmdProcessor) query(packetIO *pnet.PacketIO, sql string) (result *gomysql.Resultset, response []byte, err error) {
 	// send request
 	packetIO.ResetSequence()
 	data := hack.Slice(sql)
@@ -33,13 +33,15 @@ func (cp *CmdProcessor) query(packetIO *pnet.PacketIO, sql string) (result *gomy
 	}
 	switch response[0] {
 	case mysql.OKHeader:
-		result = cp.handleOKPacket(request, response)
+		cp.handleOKPacket(request, response)
 	case mysql.ErrHeader:
 		err = cp.handleErrorPacket(response)
 	case mysql.LocalInFileHeader:
 		err = errors.WithStack(mysql.ErrMalformPacket)
 	default:
-		result, err = cp.readResultSet(packetIO, response)
+		var rs *gomysql.Result
+		rs, err = cp.readResultSet(packetIO, response)
+		result = rs.Resultset
 	}
 	return
 }
@@ -109,8 +111,7 @@ func (cp *CmdProcessor) readResultRows(packetIO *pnet.PacketIO, result *gomysql.
 			}
 		} else {
 			if pnet.IsResultSetOKPacket(data[0], len(data)) {
-				rs := pnet.ParseOKPacket(data)
-				result.Status = rs.Status
+				result.Status = pnet.ParseOKPacket(data)
 				break
 			}
 		}

pkg/proxy/backend/mock_client_test.go

@@ -361,8 +361,7 @@ func (mc *mockClient) readResultSet(packetIO *pnet.PacketIO) error {
 			if mc.capability&pnet.ClientDeprecateEOF == 0 {
 				serverStatus = binary.LittleEndian.Uint16(pkt[3:])
 			} else {
-				rs := pnet.ParseOKPacket(pkt)
-				serverStatus = rs.Status
+				serverStatus = pnet.ParseOKPacket(pkt)
 			}
 		}
 		if serverStatus&mysql.ServerMoreResultsExists == 0 {

pkg/proxy/backend/mock_proxy_test.go

@@ -39,7 +39,7 @@ type mockProxy struct {
 
 	*proxyConfig
 	// outputs that received from the server.
-	rs *gomysql.Result
+	rs *gomysql.Resultset
 	// execution results
 	err         error
 	logger      *zap.Logger

pkg/proxy/net/mysql.go

@@ -385,17 +385,15 @@ func WriteServerVersion(conn net.Conn, serverVersion string) error {
 	return c.WritePacket(data)
 }
 
-// ParseOKPacket transforms an OK packet into a Result object.
-func ParseOKPacket(data []byte) *gomysql.Result {
-	var n int
+// ParseOKPacket parses an OK packet and only returns server status.
+func ParseOKPacket(data []byte) uint16 {
 	var pos = 1
-	r := new(gomysql.Result)
-	r.AffectedRows, _, n = ParseLengthEncodedInt(data[pos:])
-	pos += n
-	r.InsertId, _, n = ParseLengthEncodedInt(data[pos:])
-	pos += n
-	r.Status = binary.LittleEndian.Uint16(data[pos:])
-	return r
+	// skip affected rows
+	pos += SkipLengthEncodedInt(data[pos:])
+	// skip insert id
+	pos += SkipLengthEncodedInt(data[pos:])
+	// return status
+	return binary.LittleEndian.Uint16(data[pos:])
 }
 
 // ParseErrorPacket transforms an error packet into a MyError object.

pkg/proxy/net/packetio.go

@@ -325,7 +325,8 @@ func (p *PacketIO) WritePacket(data []byte, flush bool) (err error) {
 	return nil
 }
 
-func (p *PacketIO) ForwardUntil(dest *PacketIO, isEnd func(firstByte byte, firstPktLen int) bool, process func(response []byte) error) error {
+func (p *PacketIO) ForwardUntil(dest *PacketIO, isEnd func(firstByte byte, firstPktLen int) (end, needData bool),
+	process func(response []byte) error) error {
 	p.readWriter.BeginRW(rwRead)
 	dest.readWriter.BeginRW(rwWrite)
 	p.limitReader.R = p.readWriter
@@ -335,28 +336,48 @@ func (p *PacketIO) ForwardUntil(dest *PacketIO, isEnd func(firstByte byte, first
 			return errors.Wrap(ErrReadConn, err)
 		}
 		length := int(header[0]) | int(header[1])<<8 | int(header[2])<<16
-		if isEnd(header[4], length) {
+		end, needData := isEnd(header[4], length)
+		var data []byte
+		// Just call ReadFrom if the caller doesn't need the data, even if it's the last packet.
+		if end && needData {
 			// TODO: allocate a buffer from pool and return the buffer after `process`.
-			data, err := p.ReadPacket()
+			data, err = p.ReadPacket()
 			if err != nil {
 				return errors.Wrap(ErrReadConn, err)
 			}
 			if err := dest.WritePacket(data, false); err != nil {
 				return errors.Wrap(ErrWriteConn, err)
 			}
-			return process(data)
 		} else {
-			sequence, pktSequence := header[3], p.readWriter.Sequence()
-			if sequence != pktSequence {
-				return ErrInvalidSequence.GenWithStack("invalid sequence, expected %d, actual %d", pktSequence, sequence)
+			for {
+				sequence, pktSequence := header[3], p.readWriter.Sequence()
+				if sequence != pktSequence {
+					return ErrInvalidSequence.GenWithStack("invalid sequence, expected %d, actual %d", pktSequence, sequence)
+				}
+				p.readWriter.SetSequence(sequence + 1)
+				// Sequence may be different (e.g. with compression) so we can't just copy the data to the destination.
+				dest.readWriter.SetSequence(dest.readWriter.Sequence() + 1)
+				p.limitReader.N = int64(length + 4)
+				if _, err := dest.readWriter.ReadFrom(&p.limitReader); err != nil {
+					return errors.Wrap(ErrRelayConn, err)
+				}
+				// For large packets, continue.
+				if length < MaxPayloadLen {
+					break
+				}
+				if header, err = p.readWriter.Peek(4); err != nil {
+					return errors.Wrap(ErrReadConn, err)
+				}
+				length = int(header[0]) | int(header[1])<<8 | int(header[2])<<16
 			}
-			p.readWriter.SetSequence(sequence + 1)
-			// Sequence may be different (e.g. with compression) so we can't just copy the data to the destination.
-			dest.readWriter.SetSequence(dest.readWriter.Sequence() + 1)
-			p.limitReader.N = int64(length + 4)
-			if _, err := dest.readWriter.ReadFrom(&p.limitReader); err != nil {
-				return errors.Wrap(ErrRelayConn, err)
+		}
+
+		if end {
+			if process != nil {
+				// data == nil iff needData == false
+				return process(data)
 			}
+			return nil
 		}
 	}
 }