more memory-efficient implementation of line reading

irc/history/history.go

@@ -321,18 +321,6 @@ func (list *Buffer) next(index int) int {
 	}
 }
 
-// return n such that v <= n and n == 2**i for some i
-func roundUpToPowerOfTwo(v int) int {
-	// http://graphics.stanford.edu/~seander/bithacks.html
-	v -= 1
-	v |= v >> 1
-	v |= v >> 2
-	v |= v >> 4
-	v |= v >> 8
-	v |= v >> 16
-	return v + 1
-}
-
 func (list *Buffer) maybeExpand() {
 	if list.window == 0 {
 		return // autoresize is disabled
@@ -352,7 +340,7 @@ func (list *Buffer) maybeExpand() {
 		return // oldest element is old enough to overwrite
 	}
 
-	newSize := roundUpToPowerOfTwo(length + 1)
+	newSize := utils.RoundUpToPowerOfTwo(length + 1)
 	if list.maximumSize < newSize {
 		newSize = list.maximumSize
 	}

irc/history/history_test.go

@@ -241,17 +241,6 @@ func TestDisabledByResize(t *testing.T) {
 	assertEqual(len(items), 0, t)
 }
 
-func TestRoundUp(t *testing.T) {
-	assertEqual(roundUpToPowerOfTwo(2), 2, t)
-	assertEqual(roundUpToPowerOfTwo(3), 4, t)
-	assertEqual(roundUpToPowerOfTwo(64), 64, t)
-	assertEqual(roundUpToPowerOfTwo(65), 128, t)
-	assertEqual(roundUpToPowerOfTwo(100), 128, t)
-	assertEqual(roundUpToPowerOfTwo(1000), 1024, t)
-	assertEqual(roundUpToPowerOfTwo(1025), 2048, t)
-	assertEqual(roundUpToPowerOfTwo(269435457), 536870912, t)
-}
-
 func BenchmarkInsert(b *testing.B) {
 	buf := NewHistoryBuffer(1024, 0)
 	b.ResetTimer()

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 globalUtf8EnforcementSetting && !utf8.Valid(line) {
-		err = errInvalidUtf8
+		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
+		}
 	}
-	return
 }
 
 func (cc *IRCStreamConn) Close() (err error) {

irc/socket.go

@@ -69,10 +69,6 @@ func (socket *Socket) Read() (string, error) {
 
 	if err == io.EOF {
 		socket.Close()
-		// process last message properly (such as ERROR/QUIT/etc), just fail next reads/writes
-		if line != "" {
-			err = nil
-		}
 	}
 
 	return line, err

irc/utils/math.go

@@ -0,0 +1,16 @@
+// Copyright (c) 2020 Shivaram Lingamneni <slingamn@cs.stanford.edu>
+// released under the MIT license
+
+package utils
+
+// return n such that v <= n and n == 2**i for some i
+func RoundUpToPowerOfTwo(v int) int {
+	// http://graphics.stanford.edu/~seander/bithacks.html
+	v -= 1
+	v |= v >> 1
+	v |= v >> 2
+	v |= v >> 4
+	v |= v >> 8
+	v |= v >> 16
+	return v + 1
+}

irc/utils/math_test.go

@@ -0,0 +1,19 @@
+// Copyright (c) 2020 Shivaram Lingamneni <slingamn@cs.stanford.edu>
+// released under the MIT license
+
+package utils
+
+import (
+	"testing"
+)
+
+func TestRoundUp(t *testing.T) {
+	assertEqual(RoundUpToPowerOfTwo(2), 2, t)
+	assertEqual(RoundUpToPowerOfTwo(3), 4, t)
+	assertEqual(RoundUpToPowerOfTwo(64), 64, t)
+	assertEqual(RoundUpToPowerOfTwo(65), 128, t)
+	assertEqual(RoundUpToPowerOfTwo(100), 128, t)
+	assertEqual(RoundUpToPowerOfTwo(1000), 1024, t)
+	assertEqual(RoundUpToPowerOfTwo(1025), 2048, t)
+	assertEqual(RoundUpToPowerOfTwo(269435457), 536870912, t)
+}