Merge pull request #19 from abligh/datagrams-without-scanner

Do not use scanner to process datagrams; implement datagram channel

Author: mcuadros

server.go

@@ -16,11 +16,17 @@ var (
 	RFC6587 = &format.RFC6587{} // RFC6587: http://www.ietf.org/rfc/rfc6587.txt
 )
 
+const (
+	datagramChannelBufferSize = 10
+	datagramReadBufferSize    = 64 * 1024
+)
+
 type Server struct {
 	listeners               []*net.TCPListener
 	connections             []net.Conn
 	wait                    sync.WaitGroup
 	doneTcp                 chan bool
+	datagramChannel         chan DatagramMessage
 	format                  format.Format
 	handler                 Handler
 	lastError               error
@@ -58,6 +64,7 @@ func (s *Server) ListenUDP(addr string) error {
 	if err != nil {
 		return err
 	}
+	connection.SetReadBuffer(datagramReadBufferSize)
 
 	s.connections = append(s.connections, connection)
 	return nil
@@ -74,6 +81,7 @@ func (s *Server) ListenUnixgram(addr string) error {
 	if err != nil {
 		return err
 	}
+	connection.SetReadBuffer(datagramReadBufferSize)
 
 	s.connections = append(s.connections, connection)
 	return nil
@@ -110,8 +118,12 @@ func (s *Server) Boot() error {
 		s.goAcceptConnection(listener)
 	}
 
+	if len(s.connections) > 0 {
+		s.goParseDatagrams()
+	}
+
 	for _, connection := range s.connections {
-		s.goScanConnection(connection, false)
+		s.goReceiveDatagrams(connection)
 	}
 
 	return nil
@@ -132,56 +144,44 @@ func (s *Server) goAcceptConnection(listener *net.TCPListener) {
 				continue
 			}
 
-			s.goScanConnection(connection, true)
+			s.goScanConnection(connection)
 		}
 
 		s.wait.Done()
 	}(listener)
 }
 
-func (s *Server) goScanConnection(connection net.Conn, needClose bool) {
+func (s *Server) goScanConnection(connection net.Conn) {
 	scanner := bufio.NewScanner(connection)
 	if sf := s.format.GetSplitFunc(); sf != nil {
 		scanner.Split(sf)
 	}
 
 	var scanCloser *ScanCloser
-	if needClose {
-		scanCloser = &ScanCloser{scanner, connection}
-	} else {
-		scanCloser = &ScanCloser{scanner, nil}
-	}
+	scanCloser = &ScanCloser{scanner, connection}
 
 	s.wait.Add(1)
 	go s.scan(scanCloser)
 }
 
 func (s *Server) scan(scanCloser *ScanCloser) {
-	if scanCloser.closer == nil {
-		// UDP
-		for scanCloser.Scan() {
-			s.parser([]byte(scanCloser.Text()))
+loop:
+	for {
+		select {
+		case <-s.doneTcp:
+			break loop
+		default:
 		}
-	} else {
-		// TCP
-	loop:
-		for {
-			select {
-			case <-s.doneTcp:
-				break loop
-			default:
-			}
-			if s.readTimeoutMilliseconds > 0 {
-				scanCloser.closer.SetReadDeadline(time.Now().Add(time.Duration(s.readTimeoutMilliseconds) * time.Millisecond))
-			}
-			if scanCloser.Scan() {
-				s.parser([]byte(scanCloser.Text()))
-			} else {
-				break loop
-			}
+		if s.readTimeoutMilliseconds > 0 {
+			scanCloser.closer.SetReadDeadline(time.Now().Add(time.Duration(s.readTimeoutMilliseconds) * time.Millisecond))
+		}
+		if scanCloser.Scan() {
+			s.parser([]byte(scanCloser.Text()))
+		} else {
+			break loop
 		}
-		scanCloser.closer.Close()
 	}
+	scanCloser.closer.Close()
 
 	s.wait.Done()
 }
@@ -220,6 +220,9 @@ func (s *Server) Kill() error {
 	if s.doneTcp != nil {
 		close(s.doneTcp)
 	}
+	if s.datagramChannel != nil {
+		close(s.datagramChannel)
+	}
 	return nil
 }
 
@@ -237,3 +240,58 @@ type ScanCloser struct {
 	*bufio.Scanner
 	closer TimeoutCloser
 }
+
+type DatagramMessage struct {
+	message []byte
+	client  string
+}
+
+func (s *Server) goReceiveDatagrams(connection net.Conn) {
+	packetconn, ok := connection.(net.PacketConn)
+	if !ok {
+		panic("Connection is not a packet connection")
+	}
+	s.wait.Add(1)
+	go func() {
+		defer s.wait.Done()
+		for {
+			buf := make([]byte, 65536)
+			n, addr, err := packetconn.ReadFrom(buf)
+			if err == nil {
+				// Ignore trailing control characters and NULs
+				for ; (n > 0) && (buf[n-1] < 32); n-- {
+				}
+				if n > 0 {
+					s.datagramChannel <- DatagramMessage{buf[:n], addr.String()}
+				}
+			} else {
+				// there has been an error. Either the server has been killed
+				// or may be getting a transitory error due to (e.g.) the
+				// interface being shutdown in which case sleep() to avoid busy wait.
+				opError, ok := err.(*net.OpError)
+				if (ok) && !opError.Temporary() && !opError.Timeout() {
+					return
+				}
+				time.Sleep(10 * time.Millisecond)
+			}
+		}
+	}()
+}
+
+func (s *Server) goParseDatagrams() {
+	s.datagramChannel = make(chan DatagramMessage, datagramChannelBufferSize)
+
+	s.wait.Add(1)
+	go func() {
+		defer s.wait.Done()
+		for {
+			select {
+			case msg, ok := (<-s.datagramChannel):
+				if !ok {
+					return
+				}
+				s.parser(msg.message)
+			}
+		}
+	}()
+}

server_test.go

@@ -109,30 +109,26 @@ func (c *ConnMock) SetWriteDeadline(t time.Time) error {
 }
 
 func (s *ServerSuite) TestConnectionClose(c *C) {
-	for _, closeConnection := range []bool{true, false} {
-		handler := new(HandlerMock)
-		server := NewServer()
-		server.SetFormat(RFC3164)
-		server.SetHandler(handler)
-		con := ConnMock{ReadData: []byte(exampleSyslog)}
-		server.goScanConnection(&con, closeConnection)
-		server.Wait()
-		c.Check(con.isClosed, Equals, closeConnection)
-	}
+	handler := new(HandlerMock)
+	server := NewServer()
+	server.SetFormat(RFC3164)
+	server.SetHandler(handler)
+	con := ConnMock{ReadData: []byte(exampleSyslog)}
+	server.goScanConnection(&con)
+	server.Wait()
+	c.Check(con.isClosed, Equals, true)
 }
 
 func (s *ServerSuite) TestConnectionUDPKill(c *C) {
-	for _, closeConnection := range []bool{true, false} {
-		handler := new(HandlerMock)
-		server := NewServer()
-		server.SetFormat(RFC5424)
-		server.SetHandler(handler)
-		con := ConnMock{ReadData: []byte(exampleSyslog)}
-		server.goScanConnection(&con, closeConnection)
-		server.Kill()
-		server.Wait()
-		c.Check(con.isClosed, Equals, closeConnection)
-	}
+	handler := new(HandlerMock)
+	server := NewServer()
+	server.SetFormat(RFC5424)
+	server.SetHandler(handler)
+	con := ConnMock{ReadData: []byte(exampleSyslog)}
+	server.goScanConnection(&con)
+	server.Kill()
+	server.Wait()
+	c.Check(con.isClosed, Equals, true)
 }
 
 func (s *ServerSuite) TestTcpTimeout(c *C) {
@@ -143,7 +139,7 @@ func (s *ServerSuite) TestTcpTimeout(c *C) {
 	server.SetTimeout(10)
 	con := ConnMock{ReadData: []byte(exampleSyslog), ReturnTimeout: true}
 	c.Check(con.isReadDeadline, Equals, false)
-	server.goScanConnection(&con, true)
+	server.goScanConnection(&con)
 	server.Wait()
 	c.Check(con.isReadDeadline, Equals, true)
 	c.Check(handler.LastLogParts, IsNil)