fix scanner pool consuming large memory (#22679)

fix scanner pool consuming large memory

Approved by: @XuPeng-SH

Author: iamlinjunhong

pkg/sql/parsers/dialect/mysql/scanner.go

@@ -27,6 +27,10 @@ import (
 
 const eofChar = 0x100
 
+// maxPoolSQLSize is the size threshold beyond which a scanner will not be kept
+// in the pool and large fields will be cleared to release memory.
+const maxPoolSQLSize = 1 << 20 // 1 MiB
+
 var scannerPool = sync.Pool{
 	New: func() any {
 		return &Scanner{}
@@ -50,18 +54,35 @@ type Scanner struct {
 	strBuilder bytes.Buffer
 }
 
-func (s *Scanner) setSql(sql string) {
-	// This is a mysql scanner, so we set the dialect type to mysql
-	s.dialectType = dialect.MYSQL
+func (s *Scanner) reset(clearLargeOnly bool, oversized bool) {
+	// Reset light-weight state shared by both setSql and PutScanner
 	s.LastToken = ""
 	s.LastError = nil
 	s.posVarIndex = 0
 	s.MysqlSpecialComment = nil
+	s.CommentFlag = false
 	s.Pos = 0
 	s.Line = 0
 	s.Col = 0
 	s.PrePos = 0
+
+	if clearLargeOnly {
+		if oversized {
+			// Oversized by SQL size: drop both to avoid retaining huge memory.
+			s.buf = ""
+			s.strBuilder = bytes.Buffer{}
+		}
+	}
+}
+
+func (s *Scanner) setSql(sql string) {
+	// This is a mysql scanner, so we set the dialect type to mysql
+	s.dialectType = dialect.MYSQL
+	// Reset transient fields but do not aggressively clear buffers here so that
+	// small capacities can be reused.
+	s.reset(false, false)
 	s.buf = sql
+	// Reset length to 0; this keeps capacity for small cases.
 	s.strBuilder.Reset()
 }
 
@@ -72,6 +93,12 @@ func NewScanner(dialectType dialect.DialectType, sql string) *Scanner {
 }
 
 func PutScanner(scanner *Scanner) {
+	oversized := len(scanner.buf) > maxPoolSQLSize
+	// Reset shared state. Only clear buffers/strings when oversized.
+	scanner.reset(true, oversized)
+	if oversized {
+		return
+	}
 	scannerPool.Put(scanner)
 }
 

pkg/sql/parsers/dialect/mysql/scanner_test.go

@@ -391,3 +391,79 @@ func TestHexadecimalLiteral(t *testing.T) {
 		}
 	}
 }
+
+func TestScannerPoolCleanupAndThreshold(t *testing.T) {
+	// Case 1: normal small SQL should be pooled and fields cleared
+	s := NewScanner(dialect.MYSQL, "select 1")
+	// grow strBuilder a little to ensure it is cleared on Put
+	s.strBuilder.WriteString("abc")
+	PutScanner(s)
+
+	// Fetch again to see if we receive a cleared scanner from pool
+	s2 := NewScanner(dialect.MYSQL, "select 2")
+	if s2.LastToken != "" || s2.LastError != nil || s2.MysqlSpecialComment != nil || s2.Pos != 0 || s2.Line != 0 || s2.Col != 0 || s2.PrePos != 0 {
+		t.Fatalf("pooled scanner should be reset: %+v", s2)
+	}
+	if s2.strBuilder.Len() != 0 {
+		t.Fatalf("strBuilder should be cleared")
+	}
+	PutScanner(s2)
+
+	// Case 2: big SQL (>1MiB) should NOT be pooled
+	big := make([]byte, (1<<20)+10)
+	for i := range big {
+		big[i] = 'a'
+	}
+	sbig := NewScanner(dialect.MYSQL, string(big))
+	// also grow internal builder to simulate expansion
+	sbig.strBuilder.Grow(1 << 20)
+	PutScanner(sbig)
+
+	// Next Get should not necessarily return the same oversized instance; at least, it must be a clean one
+	s3 := NewScanner(dialect.MYSQL, "select 3")
+	if s3.buf != "select 3" {
+		t.Fatalf("unexpected scanner buf after Get")
+	}
+	PutScanner(s3)
+}
+
+func TestPutScannerSmallKeepsBuffers(t *testing.T) {
+	// Small SQL should keep buf and builder content when returned to pool
+	sql := "select 1"
+	s := NewScanner(dialect.MYSQL, sql)
+	s.strBuilder.WriteString("xyz")
+	PutScanner(s)
+
+	if s.buf == "" {
+		t.Fatalf("small scanner buf should not be cleared on PutScanner")
+	}
+	if s.strBuilder.Len() == 0 {
+		t.Fatalf("small scanner strBuilder should retain content on PutScanner")
+	}
+
+	// When taking from pool next time, setSql will Reset the builder length
+	s2 := NewScanner(dialect.MYSQL, "select 2")
+	if s2.strBuilder.Len() != 0 {
+		t.Fatalf("builder length must be reset on setSql")
+	}
+	PutScanner(s2)
+}
+
+func TestPutScannerOversizedClearsBuffers(t *testing.T) {
+	// Big SQL should be cleared and dropped
+	big := make([]byte, (1<<20)+123)
+	for i := range big {
+		big[i] = 'b'
+	}
+	s := NewScanner(dialect.MYSQL, string(big))
+	s.strBuilder.Grow(1 << 20)
+	s.strBuilder.WriteString("payload")
+	PutScanner(s)
+
+	if s.buf != "" {
+		t.Fatalf("oversized scanner buf should be cleared on PutScanner")
+	}
+	if s.strBuilder.Len() != 0 {
+		t.Fatalf("oversized scanner strBuilder should be zeroed on PutScanner")
+	}
+}