Introduce timed flush

pkg/extractor/batchers/batcher.go

@@ -16,6 +16,9 @@ import (
 // ReadAheadBufferSize is the default size of the read-ahead buffer
 const ReadAheadBufferSize = 128 * 1024
 
+// AutoFlushTimeout sets time before an auto-flushing reader will write a batch
+const AutoFlushTimeout = 250 * time.Millisecond
+
 type Batcher struct {
 	c chan extractor.InputBatch
 
@@ -166,3 +169,43 @@ func (s *Batcher) syncReaderToBatcher(sourceName string, reader io.Reader, batch
 		s.incReadBytes(readerMetrics.CountReset())
 	}
 }
+
+// syncReaderToBatcherWithTimeFlush is similar to `syncReaderToBatcher`, except if it gets a new line
+// it will flush the batch if n time has elapsed since the last flush, irregardless of how many items are in the current batch
+// Good for potentially slow or more interactive workloads (tail, stdin, etc)
+func (s *Batcher) syncReaderToBatcherWithTimeFlush(sourceName string, reader io.Reader, batchSize int, autoFlush time.Duration) {
+	readerMetrics := newReaderMetrics(reader)
+	readahead := readahead.NewImmediate(readerMetrics, ReadAheadBufferSize)
+	readahead.OnError(func(e error) {
+		s.incErrors()
+		logger.Printf("Error reading %s: %v", sourceName, e)
+	})
+
+	batch := make([]extractor.BString, 0, batchSize)
+	var batchStart uint64 = 1
+	lastBatchFlush := time.Now()
+
+	for readahead.Scan() {
+		batch = append(batch, readahead.Bytes())
+		if len(batch) >= batchSize || time.Since(lastBatchFlush) >= autoFlush {
+			s.c <- extractor.InputBatch{
+				Batch:      batch,
+				Source:     sourceName,
+				BatchStart: batchStart,
+			}
+			batchStart += uint64(len(batch))
+			batch = make([]extractor.BString, 0, batchSize)
+
+			s.incReadBytes(readerMetrics.CountReset())
+			lastBatchFlush = time.Now()
+		}
+	}
+	if len(batch) > 0 {
+		s.c <- extractor.InputBatch{
+			Batch:      batch,
+			Source:     sourceName,
+			BatchStart: batchStart,
+		}
+		s.incReadBytes(readerMetrics.CountReset())
+	}
+}

pkg/extractor/batchers/batcher_test.go

@@ -3,6 +3,7 @@ package batchers
 import (
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/stretchr/testify/assert"
 )
@@ -41,3 +42,21 @@ line3`
 	assert.Equal(t, s.errorCount, 0)
 	assert.Equal(t, s.ReadBytes(), uint64(17))
 }
+
+func TestBatcherWithAutoFlush(t *testing.T) {
+	s := newBatcher(10)
+
+	testData := `line1
+line2
+line3`
+
+	s.syncReaderToBatcherWithTimeFlush("string", strings.NewReader(testData), 2, 1*time.Second)
+
+	b1 := <-s.BatchChan()
+	b2 := <-s.BatchChan()
+
+	assert.Len(t, b1.Batch, 2)
+	assert.Len(t, b2.Batch, 1)
+	assert.Equal(t, s.errorCount, 0)
+	assert.Equal(t, s.ReadBytes(), uint64(17))
+}

pkg/extractor/batchers/readerBatcher.go

@@ -11,7 +11,7 @@ func OpenReaderToChan(sourceName string, reader io.ReadCloser, batchSize int) *B
 		defer reader.Close()
 		defer out.close()
 		out.startFileReading(sourceName)
-		out.syncReaderToBatcher(sourceName, reader, batchSize)
+		out.syncReaderToBatcherWithTimeFlush(sourceName, reader, batchSize, AutoFlushTimeout)
 	}()
 
 	return out

pkg/extractor/batchers/tailBatcher.go

@@ -29,7 +29,7 @@ func TailFilesToChan(filenames <-chan string, batchSize int, reopen, poll bool)
 					return
 				}
 
-				out.syncReaderToBatcher(filename, r, batchSize)
+				out.syncReaderToBatcherWithTimeFlush(filename, r, batchSize, AutoFlushTimeout)
 			}(filename)
 		}