From 9d238291b1f664244cb9a8b68d29f3f5b7ec43c0 Mon Sep 17 00:00:00 2001
From: Callum Styan <callumstyan@gmail.com>
Date: Tue, 9 Jan 2024 15:07:38 -0800
Subject: [PATCH] labels/line formatter pooling panic fix (#11607)

Big thanks to @slim-bean, even after ~1 day of trying to find a cause
here I didn't track down this issue, that I was returning something to
the pool that hadn't actually been retrieved from the pool.

The test here in the 2nd commit will fail if you pass the `-race` flag
without the fix from the 3rd commit.

---------

Signed-off-by: Callum Styan <callumstyan@gmail.com>
(cherry picked from commit 3048e4af136feac123b48cd5634d3c89d822f907)
---
 pkg/logql/log/fmt.go      |  8 +++++--
 pkg/logql/log/fmt_test.go | 46 +++++++++++++++++++++++++++++++++++++++
 pkg/logql/log/labels.go   |  7 +++---
 3 files changed, 55 insertions(+), 6 deletions(-)

diff --git a/pkg/logql/log/fmt.go b/pkg/logql/log/fmt.go
index 9257834eee34..34a1bb32c5e4 100644
--- a/pkg/logql/log/fmt.go
+++ b/pkg/logql/log/fmt.go
@@ -222,8 +222,12 @@ func (lf *LineFormatter) Process(ts int64, line []byte, lbs *LabelsBuilder) ([]b
 	lf.currentTs = ts
 
 	// map now is taking from a pool
-	m := lbs.Map()
-	defer smp.Put(m)
+	m, ret := lbs.Map()
+	defer func() {
+		if ret { // if we return the base map from the labels builder we should not put it back in the pool
+			smp.Put(m)
+		}
+	}()
 	if err := lf.Template.Execute(lf.buf, m); err != nil {
 		lbs.SetErr(errTemplateFormat)
 		lbs.SetErrorDetails(err.Error())
diff --git a/pkg/logql/log/fmt_test.go b/pkg/logql/log/fmt_test.go
index 06cf20c37606..637caec29a46 100644
--- a/pkg/logql/log/fmt_test.go
+++ b/pkg/logql/log/fmt_test.go
@@ -2,6 +2,7 @@ package log
 
 import (
 	"fmt"
+	"sync"
 	"testing"
 	"time"
 
@@ -923,3 +924,48 @@ func TestInvalidUnixTimes(t *testing.T) {
 	_, err = unixToTime("464")
 	require.Error(t, err)
 }
+
+func TestMapPoolPanic(_ *testing.T) {
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	wgFinished := sync.WaitGroup{}
+
+	ls := labels.FromStrings("cluster", "us-central-0")
+	builder := NewBaseLabelsBuilder().ForLabels(ls, ls.Hash())
+	// this specific line format was part of the query that first alerted us to the panic caused by map pooling in the label/line formatter Process functions
+	tmpl := `[1m{{if .level }}{{alignRight 5 .level}}{{else if .severity}}{{alignRight 5 .severity}}{{end}}[0m [90m[{{alignRight 10 .resources_service_instance_id}}{{if .attributes_thread_name}}/{{alignRight 20 .attributes_thread_name}}{{else if eq "java" .resources_telemetry_sdk_language }}                    {{end}}][0m [36m{{if .instrumentation_scope_name }}{{alignRight 40 .instrumentation_scope_name}}{{end}}[0m {{.body}} {{if .traceid}} [37m[3m[traceid={{.traceid}}]{{end}}`
+	a := newMustLineFormatter(tmpl)
+	a.Process(0,
+		[]byte("logger=sqlstore.metrics traceID=XXXXXXXXXXXXXXXXXXXXXXXXXXXX t=2024-01-04T23:58:47.696779826Z level=debug msg=\"query finished\" status=success elapsedtime=1.523571ms sql=\"some SQL query\" error=null"),
+		builder,
+	)
+
+	for i := 0; i < 100; i++ {
+		wgFinished.Add(1)
+		go func() {
+			wg.Wait()
+			a := newMustLineFormatter(tmpl)
+			a.Process(0,
+				[]byte("logger=sqlstore.metrics traceID=XXXXXXXXXXXXXXXXXXXXXXXXXXXX t=2024-01-04T23:58:47.696779826Z level=debug msg=\"query finished\" status=success elapsedtime=1.523571ms sql=\"some SQL query\" error=null"),
+				builder,
+			)
+			wgFinished.Done()
+		}()
+	}
+	for i := 0; i < 100; i++ {
+		wgFinished.Add(1)
+		j := i
+		go func() {
+			wg.Wait()
+			m := smp.Get()
+			for k, v := range m {
+				m[k] = fmt.Sprintf("%s%d", v, j)
+			}
+			smp.Put(m)
+			wgFinished.Done()
+
+		}()
+	}
+	wg.Done()
+	wgFinished.Wait()
+}
diff --git a/pkg/logql/log/labels.go b/pkg/logql/log/labels.go
index 76a1ae0d7d5e..2144abdf5d7e 100644
--- a/pkg/logql/log/labels.go
+++ b/pkg/logql/log/labels.go
@@ -495,22 +495,21 @@ func (b *LabelsBuilder) IntoMap(m map[string]string) {
 	}
 }
 
-func (b *LabelsBuilder) Map() map[string]string {
+func (b *LabelsBuilder) Map() (map[string]string, bool) {
 	if !b.hasDel() && !b.hasAdd() && !b.HasErr() {
 		if b.baseMap == nil {
 			b.baseMap = b.base.Map()
 		}
-		return b.baseMap
+		return b.baseMap, false
 	}
 	b.buf = b.UnsortedLabels(b.buf)
 	// todo should we also cache maps since limited by the result ?
 	// Maps also don't create a copy of the labels.
 	res := smp.Get()
-	clear(res)
 	for _, l := range b.buf {
 		res[l.Name] = l.Value
 	}
-	return res
+	return res, true
 }
 
 // LabelsResult returns the LabelsResult from the builder.