From fa46c28c5f42170d890bafefef312f5ea484025d Mon Sep 17 00:00:00 2001
From: Aliaksandr Valialkin <valyala@victoriametrics.com>
Date: Fri, 30 Sep 2022 10:43:31 +0300
Subject: [PATCH] lib/promrelabel: optimize `action: labelmap` for non-trivial
 regexs

---
 lib/promrelabel/config.go      |  7 +++--
 lib/promrelabel/config_test.go |  5 ++++
 lib/promrelabel/relabel.go     | 47 ++++++++++++++++++++++------------
 3 files changed, 40 insertions(+), 19 deletions(-)

diff --git a/lib/promrelabel/config.go b/lib/promrelabel/config.go
index 346e0320a1..65c383dd3b 100644
--- a/lib/promrelabel/config.go
+++ b/lib/promrelabel/config.go
@@ -6,6 +6,7 @@ import (
 	"strconv"
 	"strings"
 
+	"github.com/VictoriaMetrics/VictoriaMetrics/lib/bytesutil"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/envtemplate"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/fs"
 	"github.com/VictoriaMetrics/VictoriaMetrics/lib/logger"
@@ -346,7 +347,7 @@ func parseRelabelConfig(rc *RelabelConfig) (*parsedRelabelConfig, error) {
 			return nil, fmt.Errorf("`labels` config cannot be applied to `action=%s`; it is applied only to `action=graphite`", action)
 		}
 	}
-	return &parsedRelabelConfig{
+	prc := &parsedRelabelConfig{
 		SourceLabels:  sourceLabels,
 		Separator:     separator,
 		TargetLabel:   targetLabel,
@@ -365,7 +366,9 @@ func parseRelabelConfig(rc *RelabelConfig) (*parsedRelabelConfig, error) {
 		hasCaptureGroupInTargetLabel:   strings.Contains(targetLabel, "$"),
 		hasCaptureGroupInReplacement:   strings.Contains(replacement, "$"),
 		hasLabelReferenceInReplacement: strings.Contains(replacement, "{{"),
-	}, nil
+	}
+	prc.stringReplacer = bytesutil.NewFastStringTransformer(prc.replaceFullStringSlow)
+	return prc, nil
 }
 
 func isDefaultRegex(expr string) bool {
diff --git a/lib/promrelabel/config_test.go b/lib/promrelabel/config_test.go
index 18066c1e26..db0919a05d 100644
--- a/lib/promrelabel/config_test.go
+++ b/lib/promrelabel/config_test.go
@@ -161,6 +161,11 @@ func TestParseRelabelConfigsSuccess(t *testing.T) {
 		if err != nil {
 			t.Fatalf("unexpected error: %s", err)
 		}
+		if pcs != nil {
+			for _, prc := range pcs.prcs {
+				prc.stringReplacer = nil
+			}
+		}
 		if !reflect.DeepEqual(pcs, pcsExpected) {
 			t.Fatalf("unexpected pcs; got\n%#v\nwant\n%#v", pcs, pcsExpected)
 		}
diff --git a/lib/promrelabel/relabel.go b/lib/promrelabel/relabel.go
index b8d293aa8b..f0a9b760d7 100644
--- a/lib/promrelabel/relabel.go
+++ b/lib/promrelabel/relabel.go
@@ -35,6 +35,8 @@ type parsedRelabelConfig struct {
 	hasCaptureGroupInTargetLabel   bool
 	hasCaptureGroupInReplacement   bool
 	hasLabelReferenceInReplacement bool
+
+	stringReplacer *bytesutil.FastStringTransformer
 }
 
 // String returns human-readable representation for prc.
@@ -305,8 +307,8 @@ func (prc *parsedRelabelConfig) apply(labels []prompbmarshal.Label, labelsOffset
 	case "labelmap":
 		// Replace label names with the `replacement` if they match `regex`
 		for _, label := range src {
-			labelName, ok := prc.replaceFullString(label.Name, prc.Replacement, prc.hasCaptureGroupInReplacement)
-			if ok {
+			labelName := prc.replaceFullStringFast(label.Name)
+			if labelName != label.Name {
 				labels = setLabelValue(labels, labelsOffset, labelName, label.Value)
 			}
 		}
@@ -360,16 +362,23 @@ func (prc *parsedRelabelConfig) apply(labels []prompbmarshal.Label, labelsOffset
 	}
 }
 
-func (prc *parsedRelabelConfig) replaceFullString(s, replacement string, hasCaptureGroupInReplacement bool) (string, bool) {
+// replaceFullStringFast replaces s with the replacement if s matches '^regex$'.
+//
+// s is returned as is if it doesn't match '^regex$'.
+func (prc *parsedRelabelConfig) replaceFullStringFast(s string) string {
 	prefix, complete := prc.regexOriginal.LiteralPrefix()
-	if complete && !hasCaptureGroupInReplacement {
+	replacement := prc.Replacement
+	if complete && !prc.hasCaptureGroupInReplacement {
 		if s == prefix {
-			return replacement, true
+			// Fast path - s matches literal regex
+			return replacement
 		}
-		return s, false
+		// Fast path - s doesn't match literal regex
+		return s
 	}
 	if !strings.HasPrefix(s, prefix) {
-		return s, false
+		// Fast path - s doesn't match literl prefix from regex
+		return s
 	}
 	if replacement == "$1" {
 		// Fast path for commonly used rule for deleting label prefixes such as:
@@ -383,29 +392,33 @@ func (prc *parsedRelabelConfig) replaceFullString(s, replacement string, hasCapt
 			reSuffix := reStr[len(prefix):]
 			switch reSuffix {
 			case "(.*)":
-				return suffix, true
+				return suffix
 			case "(.+)":
 				if len(suffix) > 0 {
-					return suffix, true
+					return suffix
 				}
-				return s, false
+				return s
 			}
 		}
 	}
+	// Slow path - handle the rest of cases.
+	return prc.stringReplacer.Transform(s)
+}
+
+// replaceFullStringSlow replaces s with the replacement if s matches '^regex$'.
+//
+// s is returned as is if it doesn't match '^regex$'.
+func (prc *parsedRelabelConfig) replaceFullStringSlow(s string) string {
 	if re := prc.regex; re.HasPrefix() && !re.MatchString(s) {
 		// Fast path - regex mismatch
-		return s, false
+		return s
 	}
 	// Slow path - regexp processing
 	match := prc.RegexAnchored.FindStringSubmatchIndex(s)
 	if match == nil {
-		return s, false
+		return s
 	}
-	bb := relabelBufPool.Get()
-	bb.B = prc.RegexAnchored.ExpandString(bb.B[:0], replacement, s, match)
-	result := string(bb.B)
-	relabelBufPool.Put(bb)
-	return result, true
+	return prc.expandCaptureGroups(prc.Replacement, s, match)
 }
 
 func (prc *parsedRelabelConfig) replaceStringSubmatches(s, replacement string, hasCaptureGroupInReplacement bool) (string, bool) {