From f8baf1a76d82dd5825669241958c82fa73c9b8cf Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@gmail.com>
Date: Wed, 24 Feb 2021 17:57:50 +0200
Subject: [PATCH] lib/promrelabel: optimize `labeldrop` and `labelkeep`
 relabeling for `prefix.*` and `prefix.+` regexps

---
 lib/promrelabel/relabel.go             | 16 +++++-
 lib/promrelabel/relabel_test.go        | 46 +++++++++++++++-
 lib/promrelabel/relabel_timing_test.go | 72 +++++++++++++++++++++++++-
 3 files changed, 129 insertions(+), 5 deletions(-)

diff --git a/lib/promrelabel/relabel.go b/lib/promrelabel/relabel.go
index a161e55cc..477cf2bea 100644
--- a/lib/promrelabel/relabel.go
+++ b/lib/promrelabel/relabel.go
@@ -334,7 +334,21 @@ func (prc *parsedRelabelConfig) matchString(s string) bool {
 	if complete {
 		return prefix == s
 	}
-	return strings.HasPrefix(s, prefix) && prc.Regex.MatchString(s)
+	if !strings.HasPrefix(s, prefix) {
+		return false
+	}
+	reStr := prc.regexOriginal.String()
+	if strings.HasPrefix(reStr, prefix) {
+		// Fast path for `foo.*` and `bar.+` regexps
+		reSuffix := reStr[len(prefix):]
+		switch reSuffix {
+		case ".+", "(.+)":
+			return len(s) > len(prefix)
+		case ".*", "(.*)":
+			return true
+		}
+	}
+	return prc.Regex.MatchString(s)
 }
 
 func (prc *parsedRelabelConfig) expandCaptureGroups(template, source string, match []int) string {
diff --git a/lib/promrelabel/relabel_test.go b/lib/promrelabel/relabel_test.go
index cc06887e7..c8033877a 100644
--- a/lib/promrelabel/relabel_test.go
+++ b/lib/promrelabel/relabel_test.go
@@ -713,7 +713,7 @@ func TestApplyRelabelConfigs(t *testing.T) {
 			},
 		})
 	})
-	t.Run("labeldrop-regexp", func(t *testing.T) {
+	t.Run("labeldrop-prefix", func(t *testing.T) {
 		f(`
 - action: labeldrop
   regex: "dropme.*"
@@ -730,7 +730,49 @@ func TestApplyRelabelConfigs(t *testing.T) {
 		})
 		f(`
 - action: labeldrop
-  regex: "dropme.*"
+  regex: "dropme(.+)"
+`, []prompbmarshal.Label{
+			{
+				Name:  "xxx",
+				Value: "yyy",
+			},
+			{
+				Name:  "dropme-please",
+				Value: "aaa",
+			},
+			{
+				Name:  "foo",
+				Value: "bar",
+			},
+		}, false, []prompbmarshal.Label{
+			{
+				Name:  "foo",
+				Value: "bar",
+			},
+			{
+				Name:  "xxx",
+				Value: "yyy",
+			},
+		})
+	})
+	t.Run("labeldrop-regexp", func(t *testing.T) {
+		f(`
+- action: labeldrop
+  regex: ".*dropme.*"
+`, []prompbmarshal.Label{
+			{
+				Name:  "aaa",
+				Value: "bbb",
+			},
+		}, true, []prompbmarshal.Label{
+			{
+				Name:  "aaa",
+				Value: "bbb",
+			},
+		})
+		f(`
+- action: labeldrop
+  regex: ".*dropme.*"
 `, []prompbmarshal.Label{
 			{
 				Name:  "xxx",
diff --git a/lib/promrelabel/relabel_timing_test.go b/lib/promrelabel/relabel_timing_test.go
index 4904bb0ef..bdaf05986 100644
--- a/lib/promrelabel/relabel_timing_test.go
+++ b/lib/promrelabel/relabel_timing_test.go
@@ -495,10 +495,44 @@ func BenchmarkApplyRelabelConfigs(b *testing.B) {
 			}
 		})
 	})
-	b.Run("labeldrop-match-regexp", func(b *testing.B) {
+	b.Run("labeldrop-match-prefix", func(b *testing.B) {
 		pcs := mustParseRelabelConfigs(`
 - action: labeldrop
   regex: "id.*"
+`)
+		labelsOrig := []prompbmarshal.Label{
+			{
+				Name:  "__name__",
+				Value: "metric",
+			},
+			{
+				Name:  "id",
+				Value: "foobar-random-string-here",
+			},
+		}
+		b.ReportAllocs()
+		b.SetBytes(1)
+		b.RunParallel(func(pb *testing.PB) {
+			var labels []prompbmarshal.Label
+			for pb.Next() {
+				labels = append(labels[:0], labelsOrig...)
+				labels = pcs.Apply(labels, 0, true)
+				if len(labels) != 1 {
+					panic(fmt.Errorf("unexpected number of labels; got %d; want %d; labels:\n%#v", len(labels), 1, labels))
+				}
+				if labels[0].Name != "__name__" {
+					panic(fmt.Errorf("unexpected label name; got %q; want %q", labels[0].Name, "__name__"))
+				}
+				if labels[0].Value != "metric" {
+					panic(fmt.Errorf("unexpected label value; got %q; want %q", labels[0].Value, "metric"))
+				}
+			}
+		})
+	})
+	b.Run("labeldrop-match-regexp", func(b *testing.B) {
+		pcs := mustParseRelabelConfigs(`
+- action: labeldrop
+  regex: ".*id.*"
 `)
 		labelsOrig := []prompbmarshal.Label{
 			{
@@ -591,10 +625,44 @@ func BenchmarkApplyRelabelConfigs(b *testing.B) {
 			}
 		})
 	})
-	b.Run("labelkeep-match-regexp", func(b *testing.B) {
+	b.Run("labelkeep-match-prefix", func(b *testing.B) {
 		pcs := mustParseRelabelConfigs(`
 - action: labelkeep
   regex: "id.*"
+`)
+		labelsOrig := []prompbmarshal.Label{
+			{
+				Name:  "__name__",
+				Value: "metric",
+			},
+			{
+				Name:  "id",
+				Value: "foobar-random-string-here",
+			},
+		}
+		b.ReportAllocs()
+		b.SetBytes(1)
+		b.RunParallel(func(pb *testing.PB) {
+			var labels []prompbmarshal.Label
+			for pb.Next() {
+				labels = append(labels[:0], labelsOrig...)
+				labels = pcs.Apply(labels, 0, true)
+				if len(labels) != 1 {
+					panic(fmt.Errorf("unexpected number of labels; got %d; want %d; labels:\n%#v", len(labels), 1, labels))
+				}
+				if labels[0].Name != "id" {
+					panic(fmt.Errorf("unexpected label name; got %q; want %q", labels[0].Name, "id"))
+				}
+				if labels[0].Value != "foobar-random-string-here" {
+					panic(fmt.Errorf("unexpected label value; got %q; want %q", labels[0].Value, "foobar-random-string-here"))
+				}
+			}
+		})
+	})
+	b.Run("labelkeep-match-regexp", func(b *testing.B) {
+		pcs := mustParseRelabelConfigs(`
+- action: labelkeep
+  regex: ".*id.*"
 `)
 		labelsOrig := []prompbmarshal.Label{
 			{