listener: Introduce batch abstraction

This cleans up the way the batch buffer is handled, normalising the
way reads for UDP and HTTP are done and separating batch buffer
concerns from the main listener code.

This paves the way for upcoming fixes.

listener/batch.go

@@ -0,0 +1,62 @@
+// Copyright 2018 Jump Trading
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package listener
+
+import "io"
+
+// newBatch returns a new batch buffer with the initial capacity
+// specified.
+func newBatch(capacity int) *batch {
+	return &batch{
+		buf: make([]byte, 0, capacity),
+	}
+}
+
+// batch implements a fixed buffer for storing a listener's current
+// batch. It is structured to minimise allocations and copies. Bytes
+// are read from an io.Reader (typically a network connection)
+// directly into an internal preallocated byte slice.
+//
+// Some ideas are borrowed from bytes.Buffer - the main difference is
+// the readOnceFrom method which reads just once from an
+// io.Reader. This is required to avoid grouping UDP reads together.
+type batch struct {
+	buf []byte
+}
+
+// readOnceFrom reads into the batch just once from an io.Reader.
+func (b *batch) readOnceFrom(r io.Reader) (int, error) {
+	n, err := r.Read(b.buf[len(b.buf):cap(b.buf)])
+	if n > 0 {
+		b.buf = b.buf[:len(b.buf)+n]
+	}
+	return n, err
+}
+
+// remaining returns the number of bytes still unused in the batch.
+func (b *batch) remaining() int {
+	return cap(b.buf) - len(b.buf)
+}
+
+// reset clears the batch so that it no longer holds data.
+func (b *batch) reset() {
+	b.buf = b.buf[:0]
+}
+
+// bytes returns the underlying batch byte slice. The returned slice
+// is valid only until the next modifying call to the batch.
+func (b *batch) bytes() []byte {
+	return b.buf
+}

listener/listener.go

@@ -41,7 +41,7 @@ const (
 	statFailedNATSPublish = "failed_nats_publish"
 
 	// The maximum possible UDP read size.
-	udpMaxDatagramSize = 65536
+	maxUDPDatagramSize = 65536
 )
 
 var statsInterval = 3 * time.Second
@@ -104,9 +104,7 @@ type Listener struct {
 	stats  *stats.Stats
 	probes probes.Probes
 
-	buf                []byte
-	batchSize          int
-	batchSizeThreshold int
+	batch *batch
 
 	wg   sync.WaitGroup
 	stop chan struct{}
@@ -135,13 +133,7 @@ func newListener(c *config.Config) (*Listener, error) {
 			statFailedNATSPublish,
 		),
 		probes: probes.Listen(c.ProbePort),
-		buf:    make([]byte, c.ListenerBatchBytes),
-
-		// If more than batchSizeThreshold bytes has been written to
-		// the current batch buffer, the batch will be sent. We allow
-		// for the maximum UDP datagram size to be read from the
-		// socket (unlikely but possible).
-		batchSizeThreshold: c.ListenerBatchBytes - udpMaxDatagramSize,
+		batch:  newBatch(c.ListenerBatchBytes),
 	}
 
 	nc, err := nats.Connect(l.c.NATSAddress, nats.MaxReconnects(-1))
@@ -192,14 +184,16 @@ func (l *Listener) listenUDP(sc *net.UDPConn) {
 	l.probes.SetReady(true)
 	for {
 		sc.SetReadDeadline(time.Now().Add(time.Second))
-		sz, _, err := sc.ReadFromUDP(l.buf[l.batchSize:])
+		bytesRead, err := l.batch.readOnceFrom(sc)
 		if err != nil && !isTimeout(err) {
 			l.stats.Inc(statReadErrors)
 		}
-
-		// Attempt to process the read even on error as Read may
-		// still have read some bytes successfully.
-		l.processRead(sz)
+		if bytesRead > 0 {
+			if l.c.Debug {
+				log.Printf("listener read %d bytes", bytesRead)
+			}
+			l.processRead()
+		}
 
 		select {
 		case <-l.stop:
@@ -213,17 +207,18 @@ func (l *Listener) setupHTTP() *http.Server {
 	mux := http.NewServeMux()
 	mux.HandleFunc("/write", func(w http.ResponseWriter, r *http.Request) {
 		for {
-			sz, err := r.Body.Read(l.buf[l.batchSize:])
-
-			// Attempt to process the read even on error has Read may
-			// still have read some bytes successfully.
-			l.processRead(sz)
-
+			bytesRead, err := l.batch.readOnceFrom(r.Body)
+			if bytesRead > 0 {
+				if l.c.Debug {
+					log.Printf("HTTP listener read %d bytes", bytesRead)
+				}
+				l.processRead()
+			}
 			if err != nil {
 				if err != io.EOF {
 					l.stats.Inc(statReadErrors)
 				}
-				break
+				return
 			}
 		}
 	})
@@ -250,27 +245,26 @@ func (l *Listener) listenHTTP(server *http.Server) {
 	server.Close()
 }
 
-func (l *Listener) processRead(sz int) {
-	if sz < 1 {
-		return // Empty read
-	}
-
+func (l *Listener) processRead() {
 	statReceived := l.stats.Inc(statReceived)
-	l.batchSize += sz
 
-	if l.c.Debug {
-		log.Printf("listener read %d bytes\n", sz)
-	}
+	// Send when the configured number of reads have been batched or
+	// the batch buffer is almost full.
+
+	// If the batch has less capacity left than the size of a maximum
+	// UDP datagram, then force a send to avoid growing the batch
+	// unnecessarily (allocations hurt performance). UDP datagrams of
+	// this size are practically unlikely but it's a nice number to
+	// use.
+	batchNearlyFull := l.batch.remaining() <= maxUDPDatagramSize
 
-	// Send when sufficient reads have been batched or the batch
-	// buffer is almost full.
-	if statReceived%l.c.BatchMessages == 0 || l.batchSize > l.batchSizeThreshold {
+	if statReceived%l.c.BatchMessages == 0 || batchNearlyFull {
 		l.stats.Inc(statSent)
-		if err := l.nc.Publish(l.c.NATSSubject[0], l.buf[:l.batchSize]); err != nil {
+		if err := l.nc.Publish(l.c.NATSSubject[0], l.batch.bytes()); err != nil {
 			l.stats.Inc(statFailedNATSPublish)
 			l.handleNatsError(err)
 		}
-		l.batchSize = 0
+		l.batch.reset()
 	}
 }