From 9f297f8c8c8a21f68917edb2cd896d62f1166b09 Mon Sep 17 00:00:00 2001
From: Travis Patterson <travis.patterson@grafana.com>
Date: Mon, 6 May 2024 09:21:38 -0600
Subject: [PATCH] fix: Invalidate caches when pipeline wrappers are disabled

---
 pkg/storage/chunk/cache/resultscache/cache.go |  2 +-
 .../resultscache/pipelinewrapper_keygen.go    | 23 +++++++++++++
 .../pipelinewrapper_keygen_test.go            | 32 +++++++++++++++++++
 3 files changed, 56 insertions(+), 1 deletion(-)
 create mode 100644 pkg/storage/chunk/cache/resultscache/pipelinewrapper_keygen.go
 create mode 100644 pkg/storage/chunk/cache/resultscache/pipelinewrapper_keygen_test.go

diff --git a/pkg/storage/chunk/cache/resultscache/cache.go b/pkg/storage/chunk/cache/resultscache/cache.go
index 549e0b72983c..d6e153cf693b 100644
--- a/pkg/storage/chunk/cache/resultscache/cache.go
+++ b/pkg/storage/chunk/cache/resultscache/cache.go
@@ -87,7 +87,7 @@ func NewResultsCache(
 		next:                 next,
 		cache:                c,
 		limits:               limits,
-		keyGen:               keyGen,
+		keyGen:               NewPipelineWrapperKeygen(keyGen),
 		cacheGenNumberLoader: cacheGenNumberLoader,
 		retentionEnabled:     retentionEnabled,
 		extractor:            extractor,
diff --git a/pkg/storage/chunk/cache/resultscache/pipelinewrapper_keygen.go b/pkg/storage/chunk/cache/resultscache/pipelinewrapper_keygen.go
new file mode 100644
index 000000000000..e3681d961eff
--- /dev/null
+++ b/pkg/storage/chunk/cache/resultscache/pipelinewrapper_keygen.go
@@ -0,0 +1,23 @@
+package resultscache
+
+import (
+	"context"
+	"github.com/grafana/loki/v3/pkg/util/httpreq"
+)
+
+type PipelineWrapperKeyGenerator struct {
+	inner KeyGenerator
+}
+
+func NewPipelineWrapperKeygen(inner KeyGenerator) KeyGenerator {
+	return &PipelineWrapperKeyGenerator{inner: inner}
+}
+
+func (kg *PipelineWrapperKeyGenerator) GenerateCacheKey(ctx context.Context, userID string, r Request) string {
+	innerKey := kg.inner.GenerateCacheKey(ctx, userID, r)
+
+	if httpreq.ExtractHeader(ctx, httpreq.LokiDisablePipelineWrappersHeader) == "true" {
+		return "pipeline-disabled:" + innerKey
+	}
+	return innerKey
+}
diff --git a/pkg/storage/chunk/cache/resultscache/pipelinewrapper_keygen_test.go b/pkg/storage/chunk/cache/resultscache/pipelinewrapper_keygen_test.go
new file mode 100644
index 000000000000..621a77d85907
--- /dev/null
+++ b/pkg/storage/chunk/cache/resultscache/pipelinewrapper_keygen_test.go
@@ -0,0 +1,32 @@
+package resultscache
+
+import (
+	"context"
+	"github.com/grafana/loki/v3/pkg/util/httpreq"
+	"github.com/stretchr/testify/require"
+	"testing"
+)
+
+func TestPipelineWrapperKeygen(t *testing.T) {
+	kg := &stubKeygen{key: "cache-key"}
+	keygen := NewPipelineWrapperKeygen(kg)
+
+	t.Run("it does nothing if pipeline wrappers aren't disabled", func(t *testing.T) {
+		key := keygen.GenerateCacheKey(context.Background(), "", nil)
+		require.Equal(t, "cache-key", key)
+	})
+
+	t.Run("it changes the key when pipeline wrappers are disabled", func(t *testing.T) {
+		ctx := httpreq.InjectHeader(context.Background(), httpreq.LokiDisablePipelineWrappersHeader, "true")
+		key := keygen.GenerateCacheKey(ctx, "", nil)
+		require.Equal(t, "pipeline-disabled:cache-key", key)
+	})
+}
+
+type stubKeygen struct {
+	key string
+}
+
+func (k *stubKeygen) GenerateCacheKey(_ context.Context, _ string, _ Request) string {
+	return k.key
+}