Merge pull request pkg#343 from drakkan/allocations

eikenb · web-flow · commit 18dc4db7a456 · 2020-03-10T15:48:33.000-07:00
fileget: allocate a slice with enough capacity
diff --git a/packet.go b/packet.go
@@ -584,6 +584,13 @@ func (p *sshFxpReadPacket) UnmarshalBinary(b []byte) error {
 	return nil
 }
 
+func (p *sshFxpReadPacket) getDataSlice() []byte {
+	dataLen := clamp(p.Len, maxTxPacket)
+	// we allocate a slice with a bigger capacity so we avoid a new allocation in sshFxpDataPacket.MarshalBinary
+	// and in sendPacket, we need 9 bytes in MarshalBinary and 4 bytes in sendPacket
+	return make([]byte, dataLen, dataLen+9+4)
+}
+
 type sshFxpRenamePacket struct {
 	ID      uint32
 	Oldpath string
@@ -819,11 +826,14 @@ type sshFxpDataPacket struct {
 	Data   []byte
 }
 
+// MarshalBinary encodes the receiver into a binary form and returns the result.
+// To avoid a new allocation the Data slice must have a capacity >= Length + 9
 func (p sshFxpDataPacket) MarshalBinary() ([]byte, error) {
-	b := []byte{sshFxpData}
-	b = marshalUint32(b, p.ID)
-	b = marshalUint32(b, p.Length)
-	b = append(b, p.Data[:p.Length]...)
+	b := append(p.Data, make([]byte, 9)...)
+	copy(b[9:], p.Data[:p.Length])
+	b[0] = sshFxpData
+	binary.BigEndian.PutUint32(b[1:5], p.ID)
+	binary.BigEndian.PutUint32(b[5:9], p.Length)
 	return b, nil
 }
 
diff --git a/request.go b/request.go
@@ -215,8 +215,7 @@ func fileget(h FileReader, r *Request, pkt requestPacket) responsePacket {
 		return statusFromError(pkt, errors.New("unexpected read packet"))
 	}
 
-	_, offset, length := packetData(pkt)
-	data := make([]byte, clamp(length, maxTxPacket))
+	data, offset, _ := packetData(pkt)
 	n, err := reader.ReadAt(data, offset)
 	// only return EOF erro if no data left to read
 	if err != nil && (err != io.EOF || n == 0) {
@@ -250,6 +249,7 @@ func packetData(p requestPacket) (data []byte, offset int64, length uint32) {
 	case *sshFxpReadPacket:
 		length = p.Len
 		offset = int64(p.Offset)
+		data = p.getDataSlice()
 	case *sshFxpWritePacket:
 		data = p.Data
 		length = p.Length
diff --git a/server.go b/server.go
@@ -34,7 +34,6 @@ type Server struct {
 	openFiles     map[string]*os.File
 	openFilesLock sync.RWMutex
 	handleCount   int
-	maxTxPacket   uint32
 }
 
 func (svr *Server) nextHandle(f *os.File) string {
@@ -87,7 +86,6 @@ func NewServer(rwc io.ReadWriteCloser, options ...ServerOption) (*Server, error)
 		debugStream: ioutil.Discard,
 		pktMgr:      newPktMgr(svrConn),
 		openFiles:   make(map[string]*os.File),
-		maxTxPacket: 1 << 15,
 	}
 
 	for _, o := range options {
@@ -258,7 +256,7 @@ func handlePacket(s *Server, p orderedRequest) error {
 		f, ok := s.getHandle(p.Handle)
 		if ok {
 			err = nil
-			data := make([]byte, clamp(p.Len, s.maxTxPacket))
+			data := p.getDataSlice()
 			n, _err := f.ReadAt(data, int64(p.Offset))
 			if _err != nil && (_err != io.EOF || n == 0) {
 				err = _err
@@ -267,6 +265,7 @@ func handlePacket(s *Server, p orderedRequest) error {
 				ID:     p.ID,
 				Length: uint32(n),
 				Data:   data[:n],
+				// do not use data[:n:n] here to clamp the capacity, we allocated extra capacity above to avoid reallocations
 			}
 		}
 		if err != nil {
