Merge pull request #218 from slingamn/socketwriter.1

refactor irc.Socket

irc/client.go

@@ -87,7 +87,9 @@ type Client struct {
 // NewClient returns a client with all the appropriate info setup.
 func NewClient(server *Server, conn net.Conn, isTLS bool) *Client {
 	now := time.Now()
-	socket := NewSocket(conn, server.MaxSendQBytes)
+	limits := server.Limits()
+	fullLineLenLimit := limits.LineLen.Tags + limits.LineLen.Rest
+	socket := NewSocket(conn, fullLineLenLimit*2, server.MaxSendQBytes())
 	go socket.RunSocketWriter()
 	client := &Client{
 		atime:          now,
@@ -253,7 +255,11 @@ func (client *Client) run() {
 
 		line, err = client.socket.Read()
 		if err != nil {
-			client.Quit("connection closed")
+			quitMessage := "connection closed"
+			if err == errReadQ {
+				quitMessage = "readQ exceeded"
+			}
+			client.Quit(quitMessage)
 			break
 		}
 

irc/config.go

@@ -216,7 +216,7 @@ type Config struct {
 		ProxyAllowedFrom    []string       `yaml:"proxy-allowed-from"`
 		WebIRC              []webircConfig `yaml:"webirc"`
 		MaxSendQString      string         `yaml:"max-sendq"`
-		MaxSendQBytes       uint64
+		MaxSendQBytes       int
 		ConnectionLimiter   connection_limits.LimiterConfig   `yaml:"connection-limits"`
 		ConnectionThrottler connection_limits.ThrottlerConfig `yaml:"connection-throttling"`
 	}
@@ -530,10 +530,11 @@ func LoadConfig(filename string) (config *Config, err error) {
 		}
 	}
 
-	config.Server.MaxSendQBytes, err = bytefmt.ToBytes(config.Server.MaxSendQString)
+	maxSendQBytes, err := bytefmt.ToBytes(config.Server.MaxSendQString)
 	if err != nil {
 		return nil, fmt.Errorf("Could not parse maximum SendQ size (make sure it only contains whole numbers): %s", err.Error())
 	}
+	config.Server.MaxSendQBytes = int(maxSendQBytes)
 
 	// get language files
 	config.Languages.Data = make(map[string]languages.LangData)

irc/errors.go

@@ -40,6 +40,7 @@ var (
 var (
 	errNoPeerCerts = errors.New("Client did not provide a certificate")
 	errNotTLS      = errors.New("Not a TLS connection")
+	errReadQ       = errors.New("ReadQ Exceeded")
 )
 
 // String Errors

irc/getters.go

@@ -6,8 +6,17 @@ package irc
 import (
 	"github.com/oragono/oragono/irc/isupport"
 	"github.com/oragono/oragono/irc/modes"
+	"sync/atomic"
 )
 
+func (server *Server) MaxSendQBytes() int {
+	return int(atomic.LoadUint32(&server.maxSendQBytes))
+}
+
+func (server *Server) SetMaxSendQBytes(m int) {
+	atomic.StoreUint32(&server.maxSendQBytes, uint32(m))
+}
+
 func (server *Server) ISupport() *isupport.List {
 	server.configurableStateMutex.RLock()
 	defer server.configurableStateMutex.RUnlock()

irc/server.go

@@ -109,7 +109,7 @@ type Server struct {
 	limits                     Limits
 	listeners                  map[string]*ListenerWrapper
 	logger                     *logger.Manager
-	MaxSendQBytes              uint64
+	maxSendQBytes              uint32
 	monitorManager             *MonitorManager
 	motdLines                  []string
 	name                       string
@@ -928,16 +928,7 @@ func (server *Server) applyConfig(config *Config, initial bool) error {
 	server.configurableStateMutex.Unlock()
 
 	// set new sendqueue size
-	if config.Server.MaxSendQBytes != server.MaxSendQBytes {
-		server.configurableStateMutex.Lock()
-		server.MaxSendQBytes = config.Server.MaxSendQBytes
-		server.configurableStateMutex.Unlock()
-
-		// update on all clients
-		for _, sClient := range server.clients.AllClients() {
-			sClient.socket.MaxSendQBytes = config.Server.MaxSendQBytes
-		}
-	}
+	server.SetMaxSendQBytes(config.Server.MaxSendQBytes)
 
 	server.loadMOTD(config.Server.MOTD, config.Server.MOTDFormatting)
 

irc/socket.go

@@ -9,6 +9,7 @@ import (
 	"crypto/sha256"
 	"crypto/tls"
 	"encoding/hex"
+	"errors"
 	"io"
 	"net"
 	"strings"
@@ -18,47 +19,44 @@ import (
 
 var (
 	handshakeTimeout, _ = time.ParseDuration("5s")
+	errSendQExceeded    = errors.New("SendQ exceeded")
 )
 
 // Socket represents an IRC socket.
 type Socket struct {
+	sync.Mutex
+
 	conn   net.Conn
 	reader *bufio.Reader
 
-	MaxSendQBytes uint64
-
-	closed      bool
-	closedMutex sync.Mutex
-
-	finalData      string // what to send when we die
-	finalDataMutex sync.Mutex
+	maxSendQBytes int
 
+	// coordination system for asynchronous writes
+	buffer           []byte
 	lineToSendExists chan bool
-	linesToSend      []string
-	linesToSendMutex sync.Mutex
+
+	closed        bool
+	sendQExceeded bool
+	finalData     string // what to send when we die
 }
 
 // NewSocket returns a new Socket.
-func NewSocket(conn net.Conn, maxSendQBytes uint64) Socket {
+func NewSocket(conn net.Conn, maxReadQBytes int, maxSendQBytes int) Socket {
 	return Socket{
 		conn:             conn,
-		reader:           bufio.NewReader(conn),
-		MaxSendQBytes:    maxSendQBytes,
-		lineToSendExists: make(chan bool),
+		reader:           bufio.NewReaderSize(conn, maxReadQBytes),
+		maxSendQBytes:    maxSendQBytes,
+		lineToSendExists: make(chan bool, 1),
 	}
 }
 
 // Close stops a Socket from being able to send/receive any more data.
 func (socket *Socket) Close() {
-	socket.closedMutex.Lock()
-	defer socket.closedMutex.Unlock()
-	if socket.closed {
-		return
-	}
+	socket.Lock()
 	socket.closed = true
+	socket.Unlock()
 
-	// force close loop to happen if it hasn't already
-	go socket.timedFillLineToSendExists(200 * time.Millisecond)
+	socket.wakeWriter()
 }
 
 // CertFP returns the fingerprint of the certificate provided by the client.
@@ -94,10 +92,13 @@ func (socket *Socket) Read() (string, error) {
 		return "", io.EOF
 	}
 
-	lineBytes, err := socket.reader.ReadBytes('\n')
+	lineBytes, isPrefix, err := socket.reader.ReadLine()
+	if isPrefix {
+		return "", errReadQ
+	}
 
 	// convert bytes to string
-	line := string(lineBytes[:])
+	line := string(lineBytes)
 
 	// read last message properly (such as ERROR/QUIT/etc), just fail next reads/writes
 	if err == io.EOF {
@@ -110,128 +111,84 @@ func (socket *Socket) Read() (string, error) {
 		return "", err
 	}
 
-	return strings.TrimRight(line, "\r\n"), nil
+	return line, nil
 }
 
 // Write sends the given string out of Socket.
-func (socket *Socket) Write(data string) error {
-	if socket.IsClosed() {
-		return io.EOF
+func (socket *Socket) Write(data string) (err error) {
+	socket.Lock()
+	if socket.closed {
+		err = io.EOF
+	} else if len(data)+len(socket.buffer) > socket.maxSendQBytes {
+		socket.sendQExceeded = true
+		err = errSendQExceeded
+	} else {
+		socket.buffer = append(socket.buffer, data...)
 	}
+	socket.Unlock()
 
-	socket.linesToSendMutex.Lock()
-	socket.linesToSend = append(socket.linesToSend, data)
-	socket.linesToSendMutex.Unlock()
-
-	go socket.timedFillLineToSendExists(15 * time.Second)
-
-	return nil
+	socket.wakeWriter()
+	return
 }
 
-// timedFillLineToSendExists either sends the note or times out.
-func (socket *Socket) timedFillLineToSendExists(duration time.Duration) {
-	lineToSendTimeout := time.NewTimer(duration)
-	defer lineToSendTimeout.Stop()
+// wakeWriter wakes up the goroutine that actually performs the write, without blocking
+func (socket *Socket) wakeWriter() {
+	// nonblocking send to the channel, no-op if it's full
 	select {
 	case socket.lineToSendExists <- true:
-		// passed data successfully
-	case <-lineToSendTimeout.C:
-		// timed out send
+	default:
 	}
 }
 
 // SetFinalData sets the final data to send when the SocketWriter closes.
 func (socket *Socket) SetFinalData(data string) {
-	socket.finalDataMutex.Lock()
+	socket.Lock()
+	defer socket.Unlock()
 	socket.finalData = data
-	socket.finalDataMutex.Unlock()
 }
 
 // IsClosed returns whether the socket is closed.
 func (socket *Socket) IsClosed() bool {
-	socket.closedMutex.Lock()
-	defer socket.closedMutex.Unlock()
+	socket.Lock()
+	defer socket.Unlock()
 	return socket.closed
 }
 
 // RunSocketWriter starts writing messages to the outgoing socket.
 func (socket *Socket) RunSocketWriter() {
-	for {
+	localBuffer := make([]byte, 0)
+	shouldStop := false
+	for !shouldStop {
 		// wait for new lines
 		select {
 		case <-socket.lineToSendExists:
-			socket.linesToSendMutex.Lock()
-
-			// check if we're closed
-			if socket.IsClosed() {
-				socket.linesToSendMutex.Unlock()
-				break
-			}
-
-			// check whether new lines actually exist or not
-			if len(socket.linesToSend) < 1 {
-				socket.linesToSendMutex.Unlock()
-				continue
-			}
-
-			// check sendq
-			var sendQBytes uint64
-			for _, line := range socket.linesToSend {
-				sendQBytes += uint64(len(line))
-				if socket.MaxSendQBytes < sendQBytes {
-					// don't unlock mutex because this break is just to escape this for loop
-					break
-				}
-			}
-			if socket.MaxSendQBytes < sendQBytes {
-				socket.SetFinalData("\r\nERROR :SendQ Exceeded\r\n")
-				socket.linesToSendMutex.Unlock()
-				break
-			}
-
-			// get all existing data
-			data := strings.Join(socket.linesToSend, "")
-			socket.linesToSend = []string{}
-
-			socket.linesToSendMutex.Unlock()
-
-			// write data
-			if 0 < len(data) {
-				_, err := socket.conn.Write([]byte(data))
-				if err != nil {
-					break
-				}
-			}
-		}
-		if socket.IsClosed() {
-			// error out or we've been closed
-			break
+			// retrieve the buffered data, clear the buffer
+			socket.Lock()
+			localBuffer = append(localBuffer, socket.buffer...)
+			socket.buffer = socket.buffer[:0]
+			socket.Unlock()
+
+			_, err := socket.conn.Write(localBuffer)
+			localBuffer = localBuffer[:0]
+
+			socket.Lock()
+			shouldStop = (err != nil) || socket.closed || socket.sendQExceeded
+			socket.Unlock()
 		}
 	}
-	// force closure of socket
-	socket.closedMutex.Lock()
-	if !socket.closed {
-		socket.closed = true
-	}
-	socket.closedMutex.Unlock()
 
-	// write error lines
-	socket.finalDataMutex.Lock()
-	if 0 < len(socket.finalData) {
-		socket.conn.Write([]byte(socket.finalData))
+	// mark the socket closed (if someone hasn't already), then write error lines
+	socket.Lock()
+	socket.closed = true
+	finalData := socket.finalData
+	if socket.sendQExceeded {
+		finalData = "\r\nERROR :SendQ Exceeded\r\n"
+	}
+	socket.Unlock()
+	if finalData != "" {
+		socket.conn.Write([]byte(finalData))
 	}
-	socket.finalDataMutex.Unlock()
 
 	// close the connection
 	socket.conn.Close()
-
-	// empty the lineToSendExists channel
-	for 0 < len(socket.lineToSendExists) {
-		<-socket.lineToSendExists
-	}
-}
-
-// WriteLine writes the given line out of Socket.
-func (socket *Socket) WriteLine(line string) error {
-	return socket.Write(line + "\r\n")
 }