Repository: atlassian/gostatsd
Pull request #353 (2020-10-28 04:18:20)

Refactor lexer


Patch:
diff --git a/pkg/statsd/lexer.go b/internal/lexer/lexer.go
similarity index 80%
rename from pkg/statsd/lexer.go
rename to internal/lexer/lexer.go
index f6b4ca6..be575cc 100644
--- a/pkg/statsd/lexer.go
+++ b/internal/lexer/lexer.go
@@ -1,4 +1,4 @@
-package statsd
+package lexer
 
 import (
 	"bytes"
@@ -10,7 +10,8 @@ import (
 	"github.com/atlassian/gostatsd/internal/pool"
 )
 
-type lexer struct {
+type Lexer struct {
+	// any field added must be considered in Lexer.reset
 	input         []byte
 	len           uint32
 	start         uint32
@@ -24,7 +25,7 @@ type lexer struct {
 	err           error
 	sampling      float64
 
-	metricPool *pool.MetricPool
+	MetricPool *pool.MetricPool
 }
 
 // assumes we don't have \x00 bytes in input.
@@ -54,7 +55,7 @@ var alertError = []byte("error")
 var alertWarning = []byte("warning")
 var alertSuccess = []byte("success")
 
-func (l *lexer) next() byte {
+func (l *Lexer) next() byte {
 	if l.pos >= l.len {
 		return eof
 	}
@@ -63,7 +64,24 @@ func (l *lexer) next() byte {
 	return b
 }
 
-func (l *lexer) run(input []byte, namespace string) (*gostatsd.Metric, *gostatsd.Event, error) {
+func (l *Lexer) reset() {
+	// l.input = nil       // re-initialized by Run
+	// l.len = 0           // re-initialized by Run
+	// l.eventTitleLen = 0 // re-initialized by lexDatadogSpecial before lexEventBody
+	// l.eventTextLen = 0  // re-initialized by lexDatadogSpecial before lexEventBody
+	// l.namespace = ""    // re-initialized by Run
+	// l.sampling = 1      // re-initialized by Run
+
+	l.start = 0
+	l.pos = 0
+	l.m = nil
+	l.e = nil
+	l.tags = nil
+	l.err = nil
+}
+
+func (l *Lexer) Run(input []byte, namespace string) (*gostatsd.Metric, *gostatsd.Event, error) {
+	l.reset()
 	l.input = input
 	l.namespace = namespace
 	l.len = uint32(len(l.input))
@@ -95,10 +113,10 @@ func (l *lexer) run(input []byte, namespace string) (*gostatsd.Metric, *gostatsd
 	return l.m, l.e, nil
 }
 
-type stateFn func(*lexer) stateFn
+type stateFn func(*Lexer) stateFn
 
 // check the first byte for special Datadog type.
-func lexSpecial(l *lexer) stateFn {
+func lexSpecial(l *Lexer) stateFn {
 	switch b := l.next(); b {
 	case '_':
 		return lexDatadogSpecial
@@ -107,7 +125,7 @@ func lexSpecial(l *lexer) stateFn {
 		return nil
 	default:
 		l.pos--
-		l.m = l.metricPool.Get()
+		l.m = l.MetricPool.Get()
 		// Pull the tags from the metric, because it may have an empty buffer we can reuse.
 		l.tags = l.m.Tags
 		return lexKeySep
@@ -115,7 +133,7 @@ func lexSpecial(l *lexer) stateFn {
 }
 
 // lex until we find the colon separator between key and value.
-func lexKeySep(l *lexer) stateFn {
+func lexKeySep(l *Lexer) stateFn {
 	for {
 		switch b := l.next(); b {
 		case '/':
@@ -142,7 +160,7 @@ func lexKeySep(l *lexer) stateFn {
 }
 
 // lex Datadog special type.
-func lexDatadogSpecial(l *lexer) stateFn {
+func lexDatadogSpecial(l *Lexer) stateFn {
 	switch b := l.next(); b {
 	// _e{title.length,text.length}:title|text|d:date_happened|h:hostname|p:priority|t:alert_type|#tag1,tag2
 	case 'e':
@@ -158,7 +176,7 @@ func lexDatadogSpecial(l *lexer) stateFn {
 	}
 }
 
-func lexEventBody(l *lexer) stateFn {
+func lexEventBody(l *Lexer) stateFn {
 	if l.len-l.pos < l.eventTitleLen+1+l.eventTextLen {
 		l.err = errNotEnoughData
 		return nil
@@ -174,7 +192,7 @@ func lexEventBody(l *lexer) stateFn {
 	return lexEventAttributes
 }
 
-func lexEventAttributes(l *lexer) stateFn {
+func lexEventAttributes(l *Lexer) stateFn {
 	switch b := l.next(); b {
 	case '|':
 		return lexEventAttribute
@@ -185,11 +203,11 @@ func lexEventAttributes(l *lexer) stateFn {
 	return nil
 }
 
-func lexEventAttribute(l *lexer) stateFn {
+func lexEventAttribute(l *Lexer) stateFn {
 	// d:date_happened|h:hostname|p:priority|t:alert_type|#tag1,tag2
 	switch b := l.next(); b {
 	case 'd':
-		return lexAssert(':', lexUint(func(l *lexer, value uint64) stateFn {
+		return lexAssert(':', lexUint(func(l *Lexer, value uint64) stateFn {
 			if value > math.MaxInt64 {
 				l.err = errOverflow
 				return nil
@@ -198,17 +216,17 @@ func lexEventAttribute(l *lexer) stateFn {
 			return lexEventAttributes
 		}))
 	case 'h':
-		return lexAssert(':', lexUntil('|', func(l *lexer, data []byte) stateFn {
+		return lexAssert(':', lexUntil('|', func(l *Lexer, data []byte) stateFn {
 			l.e.Source = gostatsd.Source(data)
 			return lexEventAttributes
 		}))
 	case 'k':
-		return lexAssert(':', lexUntil('|', func(l *lexer, data []byte) stateFn {
+		return lexAssert(':', lexUntil('|', func(l *Lexer, data []byte) stateFn {
 			l.e.AggregationKey = string(data)
 			return lexEventAttributes
 		}))
 	case 'p':
-		return lexAssert(':', lexUntil('|', func(l *lexer, data []byte) stateFn {
+		return lexAssert(':', lexUntil('|', func(l *Lexer, data []byte) stateFn {
 			if bytes.Equal(data, priorityLow) {
 				l.e.Priority = gostatsd.PriLow
 			} else if bytes.Equal(data, priorityNormal) {
@@ -220,12 +238,12 @@ func lexEventAttribute(l *lexer) stateFn {
 			return lexEventAttributes
 		}))
 	case 's':
-		return lexAssert(':', lexUntil('|', func(l *lexer, data []byte) stateFn {
+		return lexAssert(':', lexUntil('|', func(l *Lexer, data []byte) stateFn {
 			l.e.SourceTypeName = string(data)
 			return lexEventAttributes
 		}))
 	case 't':
-		return lexAssert(':', lexUntil('|', func(l *lexer, data []byte) stateFn {
+		return lexAssert(':', lexUntil('|', func(l *Lexer, data []byte) stateFn {
 			if bytes.Equal(data, alertError) {
 				l.e.AlertType = gostatsd.AlertError
 			} else if bytes.Equal(data, alertWarning) {
@@ -250,7 +268,7 @@ func lexEventAttribute(l *lexer) stateFn {
 }
 
 func lexUint32(target *uint32, next stateFn) stateFn {
-	return lexUint(func(l *lexer, value uint64) stateFn {
+	return lexUint(func(l *Lexer, value uint64) stateFn {
 		if value > math.MaxUint32 {
 			l.err = errOverflow
 			return nil
@@ -260,8 +278,8 @@ func lexUint32(target *uint32, next stateFn) stateFn {
 	})
 }
 
-func lexUint(handler func(*lexer, uint64) stateFn) stateFn {
-	return func(l *lexer) stateFn {
+func lexUint(handler func(*Lexer, uint64) stateFn) stateFn {
+	return func(l *Lexer) stateFn {
 		var value uint64
 		start := l.pos
 	loop:
@@ -291,7 +309,7 @@ func lexUint(handler func(*lexer, uint64) stateFn) stateFn {
 
 // lexAssert returns a function that checks if the next byte matches the provided byte and returns next in that case.
 func lexAssert(nextByte byte, next stateFn) stateFn {
-	return func(l *lexer) stateFn {
+	return func(l *Lexer) stateFn {
 		switch b := l.next(); b {
 		case nextByte:
 			return next
@@ -304,8 +322,8 @@ func lexAssert(nextByte byte, next stateFn) stateFn {
 
 // lexUntil invokes handler with all bytes up to the stop byte or an eof.
 // The stop byte is not consumed.
-func lexUntil(stop byte, handler func(*lexer, []byte) stateFn) stateFn {
-	return func(l *lexer) stateFn {
+func lexUntil(stop byte, handler func(*Lexer, []byte) stateFn) stateFn {
+	return func(l *Lexer) stateFn {
 		start := l.pos
 		p := bytes.IndexByte(l.input[l.pos:], stop)
 		switch p {
@@ -319,7 +337,7 @@ func lexUntil(stop byte, handler func(*lexer, []byte) stateFn) stateFn {
 }
 
 // lex the key.
-func lexKey(l *lexer) stateFn {
+func lexKey(l *Lexer) stateFn {
 	if l.start == l.pos-1 {
 		l.err = errEmptyKey
 		return nil
@@ -333,7 +351,7 @@ func lexKey(l *lexer) stateFn {
 }
 
 // lex until we find the pipe separator between value and modifier.
-func lexValueSep(l *lexer) stateFn {
+func lexValueSep(l *Lexer) stateFn {
 	for {
 		// cheap check here. ParseFloat will do it.
 		switch b := l.next(); b {
@@ -347,14 +365,14 @@ func lexValueSep(l *lexer) stateFn {
 }
 
 // lex the value.
-func lexValue(l *lexer) stateFn {
+func lexValue(l *Lexer) stateFn {
 	l.m.StringValue = string(l.input[l.start : l.pos-1])
 	l.start = l.pos
 	return lexType
 }
 
 // lex the type.
-func lexType(l *lexer) stateFn {
+func lexType(l *Lexer) stateFn {
 	b := l.next()
 	switch b {
 	case 'c':
@@ -389,7 +407,7 @@ func lexType(l *lexer) stateFn {
 }
 
 // lex the possible separator between type and sampling rate.
-func lexTypeSep(l *lexer) stateFn {
+func lexTypeSep(l *Lexer) stateFn {
 	b := l.next()
 	switch b {
 	case eof:
@@ -403,7 +421,7 @@ func lexTypeSep(l *lexer) stateFn {
 }
 
 // lex the sample rate or the tags.
-func lexSampleRateOrTags(l *lexer) stateFn {
+func lexSampleRateOrTags(l *Lexer) stateFn {
 	b := l.next()
 	switch b {
 	case '@':
@@ -426,7 +444,7 @@ func lexSampleRateOrTags(l *lexer) stateFn {
 }
 
 // lex the sample rate.
-func lexSampleRate(l *lexer) stateFn {
+func lexSampleRate(l *Lexer) stateFn {
 	v, err := strconv.ParseFloat(string(l.input[l.start:l.pos-1]), 64)
 	if err != nil {
 		l.err = err
@@ -440,8 +458,8 @@ func lexSampleRate(l *lexer) stateFn {
 }
 
 // lex the tags.
-func lexTags(l *lexer) stateFn {
-	return lexUntil(',', func(l *lexer, data []byte) stateFn {
+func lexTags(l *Lexer) stateFn {
+	return lexUntil(',', func(l *Lexer, data []byte) stateFn {
 		if len(data) > 0 {
 			l.tags = append(l.tags, string(data))
 		}
diff --git a/pkg/statsd/lexer_fuzz.go b/internal/lexer/lexer_fuzz.go
similarity index 74%
rename from pkg/statsd/lexer_fuzz.go
rename to internal/lexer/lexer_fuzz.go
index de1c4c9..202eaa3 100644
--- a/pkg/statsd/lexer_fuzz.go
+++ b/internal/lexer/lexer_fuzz.go
@@ -1,13 +1,17 @@
 // +build gofuzz
 
-package statsd
+package lexer
 
 import (
 	"fmt"
+
+	"github.com/atlassian/gostatsd/internal/pool"
 )
 
 func Fuzz(data []byte) int {
-	l := lexer{}
+	l := Lexer{
+		MetricPool: pool.NewMetricPool(0),
+	}
 	metric, event, err := l.run(data, "")
 	if err != nil {
 		return 0
diff --git a/pkg/statsd/parser.go b/pkg/statsd/parser.go
index 206eeed..4eccd18 100644
--- a/pkg/statsd/parser.go
+++ b/pkg/statsd/parser.go
@@ -12,6 +12,7 @@ import (
 	"golang.org/x/time/rate"
 
 	"github.com/atlassian/gostatsd"
+	"github.com/atlassian/gostatsd/internal/lexer"
 	"github.com/atlassian/gostatsd/internal/pool"
 	"github.com/atlassian/gostatsd/pkg/stats"
 )
@@ -94,6 +95,10 @@ func (dp *DatagramParser) RunMetricsContext(ctx context.Context) {
 func (dp *DatagramParser) Run(ctx context.Context) {
 	dp.initLogRawMetric(ctx)
 
+	l := &lexer.Lexer{
+		MetricPool: dp.metricPool,
+	}
+
 	for {
 		select {
 		case <-ctx.Done():
@@ -104,7 +109,7 @@ func (dp *DatagramParser) Run(ctx context.Context) {
 			accumB, accumE := uint64(0), uint64(0)
 			for _, dg := range dgs {
 				// TODO: Dispatch Events in Run, not handleDatagram, so it's consistent with Metrics
-				parsedMetrics, eventCount, badLineCount := dp.handleDatagram(ctx, dg.Timestamp, dg.IP, dg.Msg)
+				parsedMetrics, eventCount, badLineCount := dp.handleDatagram(ctx, l, dg.Timestamp, dg.IP, dg.Msg)
 				dg.DoneFunc()
 				metrics = append(metrics, parsedMetrics...)
 				accumE += eventCount
@@ -139,7 +144,7 @@ func (dp *DatagramParser) logBadLineRateLimited(line []byte, ip gostatsd.Source,
 
 // handleDatagram handles the contents of a datagram and parsers it in to Metrics (which are returned), or
 // Events (which are sent to the pipeline via DispatchEvent).
-func (dp *DatagramParser) handleDatagram(ctx context.Context, now gostatsd.Nanotime, ip gostatsd.Source, msg []byte) (metrics []*gostatsd.Metric, eventCount uint64, badLineCount uint64) {
+func (dp *DatagramParser) handleDatagram(ctx context.Context, l *lexer.Lexer, now gostatsd.Nanotime, ip gostatsd.Source, msg []byte) (metrics []*gostatsd.Metric, eventCount uint64, badLineCount uint64) {
 	var numEvents, numBad uint64
 	for {
 		idx := bytes.IndexByte(msg, '\n')
@@ -155,7 +160,7 @@ func (dp *DatagramParser) handleDatagram(ctx context.Context, now gostatsd.Nanot
 			line = msg[:idx]
 			msg = msg[idx+1:]
 		}
-		metric, event, err := dp.parseLine(line)
+		metric, event, err := dp.parseLine(l, line)
 		if err != nil {
 			// logging as debug to avoid spamming logs when a bad actor sends
 			// badly formatted messages
@@ -197,11 +202,8 @@ func (dp *DatagramParser) handleDatagram(ctx context.Context, now gostatsd.Nanot
 }
 
 // parseLine with lexer.
-func (dp *DatagramParser) parseLine(line []byte) (*gostatsd.Metric, *gostatsd.Event, error) {
-	l := lexer{
-		metricPool: dp.metricPool,
-	}
-	return l.run(line, dp.namespace)
+func (dp *DatagramParser) parseLine(l *lexer.Lexer, line []byte) (*gostatsd.Metric, *gostatsd.Event, error) {
+	return l.Run(line, dp.namespace)
 }
 
 func (dp *DatagramParser) initLogRawMetric(ctx context.Context) {


Test patch:
diff --git a/pkg/statsd/lexer_test.go b/internal/lexer/lexer_test.go
similarity index 91%
rename from pkg/statsd/lexer_test.go
rename to internal/lexer/lexer_test.go
index 7fde201..7c7e31a 100644
--- a/pkg/statsd/lexer_test.go
+++ b/internal/lexer/lexer_test.go
@@ -1,4 +1,4 @@
-package statsd
+package lexer
 
 import (
 	"testing"
@@ -142,10 +142,10 @@ func TestInvalidEventsLexer(t *testing.T) {
 }
 
 func parseLine(input []byte, namespace string) (*gostatsd.Metric, *gostatsd.Event, error) {
-	l := lexer{
-		metricPool: pool.NewMetricPool(0),
+	l := Lexer{
+		MetricPool: pool.NewMetricPool(0),
 	}
-	return l.run(input, namespace)
+	return l.Run(input, namespace)
 }
 
 func compareMetric(t *testing.T, tests map[string]gostatsd.Metric, namespace string) {
@@ -164,46 +164,50 @@ func compareMetric(t *testing.T, tests map[string]gostatsd.Metric, namespace str
 
 var parselineBlackhole *gostatsd.Metric
 
-func benchmarkLexer(dp *DatagramParser, input string, b *testing.B) {
+func benchmarkLexer(ns string, input string, b *testing.B) {
 	slice := []byte(input)
 	var r *gostatsd.Metric
-	dp.metricPool = pool.NewMetricPool(0)
+
+	l := &Lexer{
+		MetricPool: pool.NewMetricPool(0),
+	}
+
 	b.ReportAllocs()
 	b.ResetTimer()
 	for n := 0; n < b.N; n++ {
-		r, _, _ = dp.parseLine(slice)
+		r, _, _ = l.Run(slice, ns)
 		r.Done()
 	}
 	parselineBlackhole = r
 }
 
 func BenchmarkParseCounter(b *testing.B) {
-	benchmarkLexer(&DatagramParser{}, "foo.bar.baz:2|c", b)
+	benchmarkLexer("", "foo.bar.baz:2|c", b)
 }
 func BenchmarkParseCounterWithSampleRate(b *testing.B) {
-	benchmarkLexer(&DatagramParser{}, "smp.rte:5|c|@0.1", b)
+	benchmarkLexer("", "smp.rte:5|c|@0.1", b)
 }
 func BenchmarkParseCounterWithTags(b *testing.B) {
-	benchmarkLexer(&DatagramParser{}, "smp.rte:5|c|#foo:bar,baz", b)
+	benchmarkLexer("", "smp.rte:5|c|#foo:bar,baz", b)
 }
 func BenchmarkParseCounterWithTagsAndSampleRate(b *testing.B) {
-	benchmarkLexer(&DatagramParser{}, "smp.rte:5|c|@0.1|#foo:bar,baz", b)
+	benchmarkLexer("", "smp.rte:5|c|@0.1|#foo:bar,baz", b)
 }
 func BenchmarkParseGauge(b *testing.B) {
-	benchmarkLexer(&DatagramParser{}, "abc.def.g:3|g", b)
+	benchmarkLexer("", "abc.def.g:3|g", b)
 }
 func BenchmarkParseTimer(b *testing.B) {
-	benchmarkLexer(&DatagramParser{}, "def.g:10|ms", b)
+	benchmarkLexer("", "def.g:10|ms", b)
 }
 func BenchmarkParseSet(b *testing.B) {
-	benchmarkLexer(&DatagramParser{}, "uniq.usr:joe|s", b)
+	benchmarkLexer("", "uniq.usr:joe|s", b)
 }
 func BenchmarkParseCounterWithDefaultTags(b *testing.B) {
-	benchmarkLexer(&DatagramParser{}, "foo.bar.baz:2|c", b)
+	benchmarkLexer("", "foo.bar.baz:2|c", b)
 }
 func BenchmarkParseCounterWithDefaultTagsAndTags(b *testing.B) {
-	benchmarkLexer(&DatagramParser{}, "foo.bar.baz:2|c|#foo:bar,baz", b)
+	benchmarkLexer("", "foo.bar.baz:2|c|#foo:bar,baz", b)
 }
 func BenchmarkParseCounterWithDefaultTagsAndTagsAndNameSpace(b *testing.B) {
-	benchmarkLexer(&DatagramParser{namespace: "stats"}, "foo.bar.baz:2|c|#foo:bar,baz", b)
+	benchmarkLexer("stats", "foo.bar.baz:2|c|#foo:bar,baz", b)
 }
diff --git a/pkg/statsd/parser_test.go b/pkg/statsd/parser_test.go
index 67f0353..54896f6 100644
--- a/pkg/statsd/parser_test.go
+++ b/pkg/statsd/parser_test.go
@@ -12,6 +12,8 @@ import (
 
 	"github.com/atlassian/gostatsd"
 	"github.com/atlassian/gostatsd/internal/fixtures"
+	"github.com/atlassian/gostatsd/internal/lexer"
+	"github.com/atlassian/gostatsd/internal/pool"
 )
 
 type metricAndEvent struct {
@@ -21,6 +23,12 @@ type metricAndEvent struct {
 
 const fakeIP = gostatsd.Source("127.0.0.1")
 
+func lex() *lexer.Lexer {
+	return &lexer.Lexer{
+		MetricPool: pool.NewMetricPool(0),
+	}
+}
+
 func newTestParser(ignoreHost bool) (*DatagramParser, *countingHandler) {
 	ch := &countingHandler{}
 	return NewDatagramParser(nil, "", ignoreHost, 0, ch, rate.Limit(0), false, logrus.New()), ch
@@ -38,7 +46,7 @@ func TestParseEmptyDatagram(t *testing.T) {
 		t.Run(strconv.Itoa(pos), func(t *testing.T) {
 			t.Parallel()
 			mr, ch := newTestParser(false)
-			_, _, _ = mr.handleDatagram(context.Background(), 0, gostatsd.UnknownSource, inp)
+			_, _, _ = mr.handleDatagram(context.Background(), lex(), 0, gostatsd.UnknownSource, inp)
 			assert.Zero(t, len(ch.events), ch.events)
 			assert.Zero(t, len(ch.metrics), ch.metrics)
 		})
@@ -95,7 +103,7 @@ func TestParseDatagram(t *testing.T) {
 		t.Run(datagram, func(t *testing.T) {
 			t.Parallel()
 			mr, ch := newTestParser(false)
-			metrics, _, _ := mr.handleDatagram(context.Background(), 0, fakeIP, []byte(datagram))
+			metrics, _, _ := mr.handleDatagram(context.Background(), lex(), 0, fakeIP, []byte(datagram))
 			mm := gostatsd.NewMetricMap()
 			for _, m := range metrics {
 				mm.Receive(m)
@@ -176,7 +184,7 @@ func TestParseDatagramIgnoreHost(t *testing.T) {
 		t.Run(datagram, func(t *testing.T) {
 			t.Parallel()
 			mr, ch := newTestParser(true)
-			metrics, _, _ := mr.handleDatagram(context.Background(), 0, fakeIP, []byte(datagram))
+			metrics, _, _ := mr.handleDatagram(context.Background(), lex(), 0, fakeIP, []byte(datagram))
 			for i, e := range ch.events {
 				if e.DateHappened <= 0 {
 					t.Errorf("%q: DateHappened should be positive", e)