Repository: jumptrading/influx-spout
Pull request #76 (2018-05-15 05:01:35)

listener: Prevent large HTTP POSTs from being broken up 
This fixes an issue where a large HTTP POST, or a POST which arrives when the batch buffer is nearly full could get split, corrupting one of the lines contained within it.

Part of the fix involves cleaning up the way the batch buffer is handled, normalising the way reads for UDP and HTTP are done and separating batch buffer handling concerns from the main listener code.

The batch buffer will now be grown if required to support any HTTP POST that arrives. Growing the batch buffer is expensive and the listener is arranged to minimise the chance of the buffer needing to grow at all. Growing the buffer is a fallback.


Patch:
diff --git a/listener/batch.go b/listener/batch.go
new file mode 100644
index 0000000..8c82876
--- /dev/null
+++ b/listener/batch.go
@@ -0,0 +1,102 @@
+// 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
+}
+
+// readFrom reads everything from an io.Reader, growing the batch if
+// required.
+func (b *batch) readFrom(r io.Reader) (int, error) {
+	total := 0
+	for {
+		// If there's not much capacity left, grow the buffer.
+		if b.remaining() <= 512 {
+			b.grow()
+		}
+		n, err := r.Read(b.buf[len(b.buf):cap(b.buf)])
+		if n > 0 {
+			b.buf = b.buf[:len(b.buf)+n]
+			total += n
+		}
+
+		if err != nil {
+			if err == io.EOF {
+				err = nil
+			}
+			return total, err
+		}
+	}
+}
+
+// size returns the number of bytes currently stored in the batch.
+func (b *batch) size() int {
+	return len(b.buf)
+}
+
+// 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
+}
+
+// grow doubles the size of the batch's internal buffer. This is
+// expensive and should be avoided where possible.
+func (b *batch) grow() {
+	newBuf := make([]byte, int(cap(b.buf)*2))
+	copy(newBuf, b.buf)
+	newBuf = newBuf[:len(b.buf)]
+	b.buf = newBuf
+}
diff --git a/listener/listener.go b/listener/listener.go
index b61bb6b..3e8efda 100644
--- a/listener/listener.go
+++ b/listener/listener.go
@@ -18,7 +18,6 @@ package listener
 
 import (
 	"fmt"
-	"io"
 	"log"
 	"net"
 	"net/http"
@@ -41,7 +40,7 @@ const (
 	statFailedNATSPublish = "failed_nats_publish"
 
 	// The maximum possible UDP read size.
-	udpMaxDatagramSize = 65536
+	maxUDPDatagramSize = 65536
 )
 
 var statsInterval = 3 * time.Second
@@ -104,9 +103,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 +132,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 +183,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:
@@ -212,19 +205,15 @@ func (l *Listener) listenUDP(sc *net.UDPConn) {
 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)
-
-			if err != nil {
-				if err != io.EOF {
-					l.stats.Inc(statReadErrors)
-				}
-				break
+		bytesRead, err := l.batch.readFrom(r.Body)
+		if err != nil {
+			l.stats.Inc(statReadErrors)
+		}
+		if bytesRead > 0 {
+			if l.c.Debug {
+				log.Printf("HTTP listener read %d bytes", bytesRead)
 			}
+			l.processRead()
 		}
 	})
 	return &http.Server{
@@ -250,27 +239,24 @@ 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 size is within a (maximum) UDP datagram of the
+	// configured target batch size, then force a send to avoid
+	// growing the batch unnecessarily (allocations hurt performance).
+	batchNearlyFull := l.c.ListenerBatchBytes-l.batch.size() <= 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()
 	}
 }
 


Test patch:
diff --git a/listener/listener_medium_test.go b/listener/listener_medium_test.go
index 2b8bdd7..d4f0626 100644
--- a/listener/listener_medium_test.go
+++ b/listener/listener_medium_test.go
@@ -210,6 +210,39 @@ func TestHTTPListener(t *testing.T) {
 	assertMonitor(t, monitorCh, numLines, numLines)
 }
 
+func TestHTTPListenerBigPOST(t *testing.T) {
+	conf := testConfig()
+	conf.ListenerBatchBytes = 1024
+	// Use a batch size > 1. Even though a single write will be made,
+	// the batch should still get sent because the buffer size limit
+	// is exceeded.
+	conf.BatchMessages = 10
+
+	listener, err := StartHTTPListener(conf)
+	require.NoError(t, err)
+	spouttest.AssertReadyProbe(t, conf.ProbePort)
+	defer listener.Stop()
+
+	listenerCh, unsubListener := subListener(t)
+	defer unsubListener()
+
+	monitorCh, unsubMonitor := subMonitor(t)
+	defer unsubMonitor()
+
+	// Send a post that's bigger than the configured batch size. This
+	// will force the batch buffer to grow.
+	buf := make([]byte, conf.ListenerBatchBytes+200)
+
+	url := fmt.Sprintf("http://localhost:%d/write", listenPort)
+	_, err = http.Post(url, "text/plain", bytes.NewBuffer(buf))
+	require.NoError(t, err)
+
+	assertBatch(t, listenerCh, string(buf))
+	assertNoMore(t, listenerCh)
+
+	assertMonitor(t, monitorCh, 1, 1)
+}
+
 func BenchmarkListenerLatency(b *testing.B) {
 	listener := startListener(b, testConfig())
 	defer listener.Stop()