From e452308e3768825e2f22dfae1e9b9bb80b8c48f2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Giedrius=20Statkevi=C4=8Dius?=
 <giedrius.statkevicius@vinted.com>
Date: Mon, 28 Oct 2024 12:09:23 +0200
Subject: [PATCH] discovery/kubernetes: optimize resolvePodRef
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

resolvePodRef is in a hot path:

```
ROUTINE ======================== github.com/prometheus/prometheus/discovery/kubernetes.(*Endpoints).resolvePodRef in discovery/kubernetes/endpoints.go
    2.50TB     2.66TB (flat, cum) 22.28% of Total
         .          .    447:func (e *Endpoints) resolvePodRef(ref *apiv1.ObjectReference) *apiv1.Pod {
         .          .    448:   if ref == nil || ref.Kind != "Pod" {
         .          .    449:           return nil
         .          .    450:   }
    2.50TB     2.50TB    451:   p := &apiv1.Pod{}
         .          .    452:   p.Namespace = ref.Namespace
         .          .    453:   p.Name = ref.Name
         .          .    454:
         .   156.31GB    455:   obj, exists, err := e.podStore.Get(p)
         .          .    456:   if err != nil {
         .          .    457:           level.Error(e.logger).Log("msg", "resolving pod ref failed", "err", err)
         .          .    458:           return nil
         .          .    459:   }
         .          .    460:   if !exists {
```

This is some low hanging fruit that we can easily optimize. The key of
an object has format "namespace/name" so generate that inside of
Prometheus itself and use pooling.

```
goos: linux
goarch: amd64
pkg: github.com/prometheus/prometheus/discovery/kubernetes
cpu: Intel(R) Core(TM) i9-10885H CPU @ 2.40GHz
                 │   olddisc    │               newdisc               │
                 │    sec/op    │   sec/op     vs base                │
ResolvePodRef-16   516.3n ± 17%   289.5n ± 7%  -43.92% (p=0.000 n=10)

                 │   olddisc    │              newdisc               │
                 │     B/op     │    B/op     vs base                │
ResolvePodRef-16   1168.00 ± 0%   24.00 ± 0%  -97.95% (p=0.000 n=10)

                 │  olddisc   │            newdisc             │
                 │ allocs/op  │ allocs/op   vs base            │
ResolvePodRef-16   2.000 ± 0%   2.000 ± 0%  ~ (p=1.000 n=10) ¹
¹ all samples are equal
```

Signed-off-by: Giedrius Statkevičius <giedrius.statkevicius@vinted.com>
---
 discovery/kubernetes/endpoints.go      | 32 +++++++++++++++++++++++++-
 discovery/kubernetes/endpoints_test.go | 21 +++++++++++++++++
 2 files changed, 52 insertions(+), 1 deletion(-)

diff --git a/discovery/kubernetes/endpoints.go b/discovery/kubernetes/endpoints.go
index 5ba9df6276..934f37ee45 100644
--- a/discovery/kubernetes/endpoints.go
+++ b/discovery/kubernetes/endpoints.go
@@ -20,6 +20,8 @@ import (
 	"log/slog"
 	"net"
 	"strconv"
+	"strings"
+	"sync"
 
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/common/model"
@@ -453,15 +455,43 @@ func (e *Endpoints) buildEndpoints(eps *apiv1.Endpoints) *targetgroup.Group {
 	return tg
 }
 
+var objKeyPool = sync.Pool{}
+
+func generateObjKey(namespace, name string) (string, *strings.Builder) {
+	var sb *strings.Builder
+
+	b := objKeyPool.Get()
+	if b == nil {
+		sb = &strings.Builder{}
+	} else {
+		sb = b.(*strings.Builder)
+	}
+
+	sb.Reset()
+	if namespace == "" {
+		_, _ = sb.WriteString(name)
+		return sb.String(), sb
+	}
+
+	_, _ = sb.WriteString(namespace)
+	_, _ = sb.WriteRune('/')
+	_, _ = sb.WriteString(name)
+	return sb.String(), sb
+}
+
 func (e *Endpoints) resolvePodRef(ref *apiv1.ObjectReference) *apiv1.Pod {
 	if ref == nil || ref.Kind != "Pod" {
 		return nil
 	}
+
 	p := &apiv1.Pod{}
 	p.Namespace = ref.Namespace
 	p.Name = ref.Name
 
-	obj, exists, err := e.podStore.Get(p)
+	key, sb := generateObjKey(p.Namespace, p.Name)
+	defer objKeyPool.Put(sb)
+
+	obj, exists, err := e.podStore.GetByKey(key)
 	if err != nil {
 		e.logger.Error("resolving pod ref failed", "err", err)
 		return nil
diff --git a/discovery/kubernetes/endpoints_test.go b/discovery/kubernetes/endpoints_test.go
index 4af6889602..a1ac6e5d48 100644
--- a/discovery/kubernetes/endpoints_test.go
+++ b/discovery/kubernetes/endpoints_test.go
@@ -18,10 +18,12 @@ import (
 	"testing"
 
 	"github.com/prometheus/common/model"
+	"github.com/stretchr/testify/require"
 	v1 "k8s.io/api/core/v1"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
 	"k8s.io/apimachinery/pkg/types"
+	"k8s.io/client-go/tools/cache"
 
 	"github.com/prometheus/prometheus/discovery/targetgroup"
 )
@@ -1257,3 +1259,22 @@ func TestEndpointsDiscoverySidecarContainer(t *testing.T) {
 		},
 	}.Run(t)
 }
+
+func BenchmarkResolvePodRef(b *testing.B) {
+	indexer := cache.NewIndexer(cache.DeletionHandlingMetaNamespaceKeyFunc, nil)
+	e := &Endpoints{
+		podStore: indexer,
+	}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		p := e.resolvePodRef(&v1.ObjectReference{
+			Kind:      "Pod",
+			Name:      "testpod",
+			Namespace: "foo",
+		})
+		require.Nil(b, p)
+	}
+}