more memory-efficient implementation of line reading

irc/ircconn.go

@@ -1,9 +1,9 @@
 package irc
 
 import (
-	"bufio"
 	"bytes"
 	"errors"
+	"io"
 	"net"
 	"unicode/utf8"
 
@@ -14,7 +14,8 @@ import (
 )
 
 const (
-	maxReadQBytes = ircmsg.MaxlenTagsFromClient + MaxLineLen + 1024
+	maxReadQBytes     = ircmsg.MaxlenTagsFromClient + MaxLineLen + 1024
+	initialBufferSize = 1024
 )
 
 var (
@@ -41,8 +42,13 @@ type IRCConn interface {
 
 // IRCStreamConn is an IRCConn over a regular stream connection.
 type IRCStreamConn struct {
-	conn   *utils.WrappedConn
-	reader *bufio.Reader
+	conn *utils.WrappedConn
+
+	buf        []byte
+	start      int // start of valid (i.e., read but not yet consumed) data in the buffer
+	end        int // end of valid data in the buffer
+	searchFrom int // start of valid data in the buffer not yet searched for \n
+	eof        bool
 }
 
 func NewIRCStreamConn(conn *utils.WrappedConn) *IRCStreamConn {
@@ -67,21 +73,58 @@ func (cc *IRCStreamConn) WriteLines(buffers [][]byte) (err error) {
 	return
 }
 
-func (cc *IRCStreamConn) ReadLine() (line []byte, err error) {
-	// lazy initialize the reader in case the IP is banned
-	if cc.reader == nil {
-		cc.reader = bufio.NewReaderSize(cc.conn, maxReadQBytes)
-	}
+func (cc *IRCStreamConn) ReadLine() ([]byte, error) {
+	for {
+		// try to find a terminated line in the buffered data already read
+		nlidx := bytes.IndexByte(cc.buf[cc.searchFrom:cc.end], '\n')
+		if nlidx != -1 {
+			// got a complete line
+			line := cc.buf[cc.start : cc.searchFrom+nlidx]
+			cc.start = cc.searchFrom + nlidx + 1
+			cc.searchFrom = cc.start
+			if globalUtf8EnforcementSetting && !utf8.Valid(line) {
+				return line, errInvalidUtf8
+			} else {
+				return line, nil
+			}
+		}
 
-	var isPrefix bool
-	line, isPrefix, err = cc.reader.ReadLine()
-	if isPrefix {
-		return nil, errReadQ
+		if cc.start == 0 && len(cc.buf) == maxReadQBytes {
+			return nil, errReadQ // out of space, can't expand or slide
+		}
+
+		if cc.eof {
+			return nil, io.EOF
+		}
+
+		if len(cc.buf) < maxReadQBytes && (len(cc.buf)-(cc.end-cc.start) < initialBufferSize/2) {
+			// allocate a new buffer, copy any remaining data
+			newLen := utils.RoundUpToPowerOfTwo(len(cc.buf) + 1)
+			if newLen > maxReadQBytes {
+				newLen = maxReadQBytes
+			} else if newLen < initialBufferSize {
+				newLen = initialBufferSize
+			}
+			newBuf := make([]byte, newLen)
+			copy(newBuf, cc.buf[cc.start:cc.end])
+			cc.buf = newBuf
+		} else if cc.start != 0 {
+			// slide remaining data back to the front of the buffer
+			copy(cc.buf, cc.buf[cc.start:cc.end])
+		}
+		cc.end = cc.end - cc.start
+		cc.start = 0
+
+		cc.searchFrom = cc.end
+		n, err := cc.conn.Read(cc.buf[cc.end:])
+		cc.end += n
+		if n != 0 && err == io.EOF {
+			// we may have received new \n-terminated lines, try to parse them
+			cc.eof = true
+		} else if err != nil {
+			return nil, err
+		}
 	}
-	if globalUtf8EnforcementSetting && !utf8.Valid(line) {
-		err = errInvalidUtf8
-	}
-	return
 }
 
 func (cc *IRCStreamConn) Close() (err error) {