fix: flush data when exceeds buffer

Author: ndakkoune

aws/logs_monitoring_go/internal/parsing/scanner.go

@@ -12,6 +12,8 @@ import (
 	"regexp"
 )
 
+const maxTokenSize = 512 * 1024
+
 type Scanner struct {
 	*bufio.Scanner
 	rg *regexp.Regexp
@@ -23,6 +25,8 @@ func NewScanner(r io.Reader, rg *regexp.Regexp) *Scanner {
 		rg:      rg,
 	}
 
+	s.Buffer(make([]byte, 0, bufio.MaxScanTokenSize), maxTokenSize)
+
 	if rg != nil {
 		s.Split(s.splitOnRegex)
 	} else {
@@ -33,7 +37,7 @@ func NewScanner(r io.Reader, rg *regexp.Regexp) *Scanner {
 }
 
 func (s *Scanner) splitOnRegex(data []byte, atEOF bool) (advance int, token []byte, err error) {
-	if atEOF && len(data) == 0 {
+	if len(data) == 0 {
 		return 0, nil, nil
 	}
 
@@ -43,15 +47,15 @@ func (s *Scanner) splitOnRegex(data []byte, atEOF bool) (advance int, token []by
 		return splitAt, data[:splitAt], nil
 	}
 
-	if atEOF {
+	if atEOF || len(data) >= maxTokenSize {
 		return len(data), data, nil
 	}
 
 	return 0, nil, nil
 }
 
 func (s *Scanner) splitOnLines(data []byte, atEOF bool) (advance int, token []byte, err error) {
-	if atEOF && len(data) == 0 {
+	if len(data) == 0 {
 		return 0, nil, nil
 	}
 
@@ -63,7 +67,7 @@ func (s *Scanner) splitOnLines(data []byte, atEOF bool) (advance int, token []by
 		return j, data[:i], nil
 	}
 
-	if atEOF {
+	if atEOF || len(data) >= maxTokenSize {
 		return len(data), data, nil
 	}
 

aws/logs_monitoring_go/internal/parsing/scanner_test.go

@@ -33,6 +33,8 @@ func TestScanner(t *testing.T) {
 		"lines_mixed_delimiters":    {input: "a\r\n\fb", rg: nil, want: []string{"a", "b"}},
 		"lines_only_delimiters":     {input: "\n\r\f", rg: nil, want: nil},
 
+		"regex_empty":                     {input: "", rg: dateRegex, want: nil},
+		"regex_not_matching_at_start":     {input: "ERROR something2024-01-15 ERROR something", rg: dateRegex, want: []string{"ERROR something", "2024-01-15 ERROR something"}},
 		"regex_single_entry":              {input: "2024-01-15 ERROR something", rg: dateRegex, want: []string{"2024-01-15 ERROR something"}},
 		"regex_two_entries_with_newline":  {input: "2024-01-15 ERROR\n2024-01-16 INFO", rg: dateRegex, want: []string{"2024-01-15 ERROR\n", "2024-01-16 INFO"}},
 		"regex_continuation_lines":        {input: "2024-01-15 ERROR\n    at com.foo\n2024-01-16 INFO", rg: dateRegex, want: []string{"2024-01-15 ERROR\n    at com.foo\n", "2024-01-16 INFO"}},
@@ -65,3 +67,41 @@ func TestScanner(t *testing.T) {
 		})
 	}
 }
+
+func TestScannerWithOversizedToken(t *testing.T) {
+	t.Parallel()
+
+	overflow := 100
+	oversized := strings.Repeat("a", maxTokenSize+overflow)
+
+	tests := map[string]struct {
+		rg *regexp.Regexp
+	}{
+		"lines": {rg: nil},
+		"regex": {rg: regexp.MustCompile(`\d{4}-\d{2}-\d{2}`)},
+	}
+
+	for name, tc := range tests {
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+			scanner := NewScanner(strings.NewReader(oversized), tc.rg)
+
+			var got []string
+			for scanner.Scan() {
+				got = append(got, scanner.Text())
+			}
+			if err := scanner.Err(); err != nil {
+				t.Fatalf("unexpected error: %v", err)
+			}
+			if len(got) != 2 {
+				t.Fatalf("want 2 tokens, got %d", len(got))
+			}
+			if len(got[0]) != maxTokenSize {
+				t.Errorf("first token: want %d bytes, got %d", maxTokenSize, len(got[0]))
+			}
+			if len(got[1]) != overflow {
+				t.Errorf("second token: want %d bytes, got %d", overflow, len(got[1]))
+			}
+		})
+	}
+}