[v17] Improve performance of utils.ReplaceRegexp

lib/services/trustedcluster.go

@@ -19,6 +19,7 @@
 package services
 
 import (
+	"errors"
 	"fmt"
 
 	"github.com/gravitational/trace"
@@ -133,7 +134,7 @@ func MapRoles(r types.RoleMap, remoteRoles []string) ([]string, error) {
 					}
 					seen[replacement] = struct{}{}
 					outRoles = append(outRoles, replacement)
-				case trace.IsNotFound(err):
+				case errors.Is(err, utils.ErrReplaceRegexNotFound):
 					continue
 				default:
 					return nil, trace.Wrap(err)

lib/utils/replace.go

@@ -43,13 +43,18 @@ func GlobToRegexp(in string) string {
 	return replaceWildcard.ReplaceAllString(regexp.QuoteMeta(in), "(.*)")
 }
 
+// ErrReplaceRegexNotFound is a marker error returned by
+// [ReplaceRegexp], [RegexpWithConfig], and [ReplaceRegexpWith] to
+// indicate no matches were found.
+var ErrReplaceRegexNotFound = &trace.NotFoundError{Message: "no match found"}
+
 // ReplaceRegexp replaces value in string, accepts regular expression and simplified
 // wildcard syntax, it has several important differences with standard lib
 // regexp replacer:
 // * Wildcard globs '*' are treated as regular expression .* expression
 // * Expression is treated as regular expression if it starts with ^ and ends with $
 // * Full match is expected, partial replacements ignored
-// * If there is no match, returns a NotFound error
+// * If there is no match, returns [ErrReplaceRegexNotFound]
 func ReplaceRegexp(expression string, replaceWith string, input string) (string, error) {
 	expr, err := RegexpWithConfig(expression, RegexpConfig{})
 	if err != nil {
@@ -58,9 +63,20 @@ func ReplaceRegexp(expression string, replaceWith string, input string) (string,
 	return ReplaceRegexpWith(expr, replaceWith, input)
 }
 
-// RegexpWithConfig compiles a regular expression given some configuration.
-// There are several important differences with standard lib (see ReplaceRegexp).
-func RegexpWithConfig(expression string, config RegexpConfig) (*regexp.Regexp, error) {
+type regexKey struct {
+	expression string
+	ignoreCase bool
+}
+
+// regexpCache interns compiled regular expressions to improve performance.
+var regexpCache = mustCache[regexKey, *regexp.Regexp](2000)
+
+func replaceRegexCached(expression string, config RegexpConfig) (*regexp.Regexp, error) {
+	key := regexKey{expression: expression, ignoreCase: config.IgnoreCase}
+	if expr, ok := regexpCache.Get(key); ok {
+		return expr, nil
+	}
+
 	if !strings.HasPrefix(expression, "^") || !strings.HasSuffix(expression, "$") {
 		// replace glob-style wildcards with regexp wildcards
 		// for plain strings, and quote all characters that could
@@ -74,15 +90,27 @@ func RegexpWithConfig(expression string, config RegexpConfig) (*regexp.Regexp, e
 	if err != nil {
 		return nil, trace.BadParameter(err.Error())
 	}
+
+	regexpCache.Add(key, expr)
 	return expr, nil
 }
 
-// ReplaceRegexp replaces string in a given regexp.
+// RegexpWithConfig compiles a regular expression given some configuration.
+// There are several important differences with standard lib (see ReplaceRegexp).
+func RegexpWithConfig(expression string, config RegexpConfig) (*regexp.Regexp, error) {
+	expr, err := replaceRegexCached(expression, config)
+	return expr, trace.Wrap(err)
+}
+
+// ReplaceRegexpWith replaces string in a given regexp.
 func ReplaceRegexpWith(expr *regexp.Regexp, replaceWith string, input string) (string, error) {
-	// if there is no match, return NotFound error
 	index := expr.FindStringIndex(input)
 	if index == nil {
-		return "", trace.NotFound("no match found")
+		// The returned error is intentionally not wrapped to avoid
+		// capturing stack traces. This method is used by authorization
+		// logic and the additional overhead of strack trace capturing
+		// is a performance bottleneck.
+		return "", ErrReplaceRegexNotFound
 	}
 	return expr.ReplaceAllString(input, replaceWith), nil
 }
@@ -367,10 +395,6 @@ func mustCache[K comparable, V any](size int) *lru.Cache[K, V] {
 	return cache
 }
 
-// exprCache interns compiled regular expressions created in MatchString
-// to improve performance.
-var exprCache = mustCache[string, *regexp.Regexp](1000)
-
 // MatchString will match an input against the given expression. The expression is cached for later use.
 func MatchString(input, expression string) (bool, error) {
 	expr, err := compileRegexCached(expression)
@@ -403,7 +427,8 @@ func CompileExpression(expression string) (*regexp.Regexp, error) {
 }
 
 func compileRegexCached(expression string) (*regexp.Regexp, error) {
-	if expr, ok := exprCache.Get(expression); ok {
+	key := regexKey{expression: expression}
+	if expr, ok := regexpCache.Get(key); ok {
 		return expr, nil
 	}
 
@@ -412,7 +437,7 @@ func compileRegexCached(expression string) (*regexp.Regexp, error) {
 		return nil, trace.Wrap(err)
 	}
 
-	exprCache.Add(expression, expr)
+	regexpCache.Add(key, expr)
 	return expr, nil
 }
 

lib/utils/replace_test.go

@@ -19,6 +19,7 @@
 package utils
 
 import (
+	"strconv"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -1319,3 +1320,31 @@ func TestKubeResourceCouldMatchRules(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkReplaceRegexp(b *testing.B) {
+	b.Run("same expression", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			replaced, err := ReplaceRegexp("*", "foo", "test")
+			require.NoError(b, err)
+			require.NotEmpty(b, replaced)
+		}
+	})
+
+	b.Run("unique expressions", func(b *testing.B) {
+		for i := 0; i < b.N; i++ {
+			r := strconv.Itoa(i)
+			replaced, err := ReplaceRegexp(r, r, r)
+			require.NoError(b, err)
+			require.NotEmpty(b, replaced)
+		}
+	})
+
+	b.Run("no matches", func(b *testing.B) {
+		expression := "$abc^"
+		for i := 0; i < b.N; i++ {
+			replaced, err := ReplaceRegexp(expression, strconv.Itoa(i), "test")
+			require.ErrorIs(b, err, ErrReplaceRegexNotFound)
+			require.Empty(b, replaced)
+		}
+	})
+}

lib/utils/utils_test.go

@@ -342,21 +342,21 @@ func TestReplaceRegexp(t *testing.T) {
 			expr:    "value",
 			replace: "value",
 			in:      "val",
-			err:     trace.NotFound(""),
+			err:     ErrReplaceRegexNotFound,
 		},
 		{
 			comment: "empty value is no match",
 			expr:    "",
 			replace: "value",
 			in:      "value",
-			err:     trace.NotFound(""),
+			err:     ErrReplaceRegexNotFound,
 		},
 		{
 			comment: "bad regexp results in bad parameter error",
 			expr:    "^(($",
 			replace: "value",
 			in:      "val",
-			err:     trace.BadParameter(""),
+			err:     &trace.BadParameterError{Message: "error parsing regexp: missing closing ): `^(($`"},
 		},
 		{
 			comment: "full match is supported",
@@ -415,7 +415,7 @@ func TestReplaceRegexp(t *testing.T) {
 				require.NoError(t, err)
 				require.Equal(t, testCase.out, out)
 			} else {
-				require.IsType(t, testCase.err, err)
+				require.ErrorIs(t, err, testCase.err)
 			}
 		})
 	}