TT-17746: Improve resolver performance

kv/internal/resolve/memo.go

@@ -0,0 +1,50 @@
+package resolve
+
+import (
+	"context"
+	"sync"
+)
+
+// memoResult is a memoized resolution outcome (value or error).
+type memoResult struct {
+	val string
+	err error
+}
+
+// memo is the per-ResolveAll-call resolution cache.
+type memo struct {
+	m  map[refKey]memoResult
+	mu sync.Mutex
+}
+
+func (mm *memo) get(k refKey) (memoResult, bool) {
+	mm.mu.Lock()
+	defer mm.mu.Unlock()
+
+	mr, ok := mm.m[k]
+
+	return mr, ok
+}
+
+func (mm *memo) set(k refKey, v memoResult) {
+	mm.mu.Lock()
+	defer mm.mu.Unlock()
+
+	mm.m[k] = v
+}
+
+type memoCtxKey struct{}
+
+// withMemo attaches a fresh, per-call resolution memo to ctx.
+func withMemo(ctx context.Context) context.Context {
+	return context.WithValue(ctx, memoCtxKey{}, &memo{m: make(map[refKey]memoResult)})
+}
+
+// memoFrom returns the per-call memo, or nil when ctx carries none (e.g. a
+// direct Resolve call outside ResolveAll), in which case fetches are not
+// memoized.
+func memoFrom(ctx context.Context) *memo {
+	//nolint:errcheck
+	m, _ := ctx.Value(memoCtxKey{}).(*memo)
+	return m
+}

kv/internal/resolve/refs.go

@@ -0,0 +1,119 @@
+package resolve
+
+import (
+	"fmt"
+	"strings"
+)
+
+// refKey identifies a resolution target.
+type refKey struct {
+	store    string
+	path     string
+	fragment string
+}
+
+// parseWholeValue parses a whole-value reference of the form
+// "kv://store/path#frag".
+//
+// The bool reports whether input is a whole-value reference at all (i.e. has
+// the "kv://" prefix). When false, input is an inline/literal string and err is
+// nil. When true, err is non-nil if the reference is malformed.
+func parseWholeValue(input string) (refKey, bool, error) {
+	if !strings.HasPrefix(input, "kv://") {
+		return refKey{}, false, nil
+	}
+
+	trimmed := strings.TrimPrefix(input, "kv://")
+
+	slashIdx := strings.IndexByte(trimmed, '/')
+	if slashIdx < 0 {
+		return refKey{}, true, fmt.Errorf(
+			"%w: missing path separator in %q",
+			ErrMalformedReference,
+			input,
+		)
+	}
+
+	store := trimmed[:slashIdx]
+	path, fragment, _ := strings.Cut(trimmed[slashIdx+1:], "#")
+
+	if store == "" || path == "" {
+		return refKey{}, true, fmt.Errorf(
+			"%w: empty store name or path in %q",
+			ErrMalformedReference,
+			input,
+		)
+	}
+
+	return refKey{store: store, path: path, fragment: fragment}, true, nil
+}
+
+// parseInlineToken parses a single inline token. match is the full "$kv{...}"
+// text (used verbatim in error messages); its contents are "store:path#fragment".
+func parseInlineToken(match string) (refKey, error) {
+	// strip "$kv{" prefix and "}" suffix
+	inner := match[len("$kv{") : len(match)-1]
+
+	colonIdx := strings.IndexByte(inner, ':')
+	if colonIdx < 0 {
+		return refKey{}, fmt.Errorf(
+			"%w: missing store separator in %q",
+			ErrMalformedReference,
+			match,
+		)
+	}
+
+	store := inner[:colonIdx]
+	path, fragment, _ := strings.Cut(inner[colonIdx+1:], "#")
+
+	if store == "" || path == "" {
+		return refKey{}, fmt.Errorf(
+			"%w: empty store name or path in %q",
+			ErrMalformedReference,
+			match,
+		)
+	}
+
+	return refKey{store: store, path: path, fragment: fragment}, nil
+}
+
+// collectRefs walks a decoded JSON document and returns every distinct,
+// well-formed reference target. It is best-effort: malformed references are
+// skipped here.
+func collectRefs(node any) map[refKey]struct{} {
+	refs := make(map[refKey]struct{})
+	collectInto(node, refs)
+
+	return refs
+}
+
+func collectInto(node any, into map[refKey]struct{}) {
+	switch v := node.(type) {
+	case string:
+		collectRefsFromString(v, into)
+	case map[string]any:
+		for _, value := range v {
+			collectInto(value, into)
+		}
+	case []any:
+		for _, value := range v {
+			collectInto(value, into)
+		}
+	}
+}
+
+func collectRefsFromString(input string, into map[refKey]struct{}) {
+	if ref, ok, err := parseWholeValue(input); ok {
+		if err == nil {
+			into[ref] = struct{}{}
+		}
+
+		return
+	}
+
+	for _, match := range inlineRe.FindAllString(input, -1) {
+		if ref, err := parseInlineToken(match); err == nil {
+			into[ref] = struct{}{}
+		}
+	}
+}

kv/internal/resolve/refs_test.go

@@ -0,0 +1,163 @@
+package resolve
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestParseWholeValue(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name    string
+		input   string
+		wantOK  bool
+		wantErr bool
+		want    refKey
+	}{
+		{
+			name:   "path only",
+			input:  "kv://vault/secret/data/app",
+			wantOK: true,
+			want:   refKey{store: "vault", path: "secret/data/app"},
+		},
+		{
+			name:   "with fragment",
+			input:  "kv://vault/secret/data/app#password",
+			wantOK: true,
+			want:   refKey{store: "vault", path: "secret/data/app", fragment: "password"},
+		},
+		{
+			name:   "fragment splits on first hash only",
+			input:  "kv://vault/a/b#x/y#z",
+			wantOK: true,
+			want:   refKey{store: "vault", path: "a/b", fragment: "x/y#z"},
+		},
+		{
+			name:   "not a whole-value reference",
+			input:  "prefix-$kv{vault:secret}-suffix",
+			wantOK: false,
+		},
+		{
+			name:   "plain literal",
+			input:  "just-a-string",
+			wantOK: false,
+		},
+		{
+			name:    "missing path separator",
+			input:   "kv://vault",
+			wantOK:  true,
+			wantErr: true,
+		},
+		{
+			name:    "empty store",
+			input:   "kv:///secret/data/app",
+			wantOK:  true,
+			wantErr: true,
+		},
+		{
+			name:    "empty path",
+			input:   "kv://vault/",
+			wantOK:  true,
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ref, ok, err := parseWholeValue(tc.input)
+
+			require.Equal(t, tc.wantOK, ok)
+
+			if tc.wantErr {
+				require.ErrorIs(t, err, ErrMalformedReference)
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, tc.want, ref)
+		})
+	}
+}
+
+func TestParseInlineToken(t *testing.T) {
+	t.Parallel()
+
+	tests := []struct {
+		name    string
+		match   string
+		wantErr bool
+		want    refKey
+	}{
+		{
+			name:  "path only",
+			match: "$kv{vault:secret/data/app}",
+			want:  refKey{store: "vault", path: "secret/data/app"},
+		},
+		{
+			name:  "with fragment",
+			match: "$kv{vault:secret/data/app#password}",
+			want:  refKey{store: "vault", path: "secret/data/app", fragment: "password"},
+		},
+		{
+			name:    "missing store separator",
+			match:   "$kv{no-colon-here}",
+			wantErr: true,
+		},
+		{
+			name:    "empty store",
+			match:   "$kv{:secret}",
+			wantErr: true,
+		},
+		{
+			name:    "empty path",
+			match:   "$kv{vault:}",
+			wantErr: true,
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			ref, err := parseInlineToken(tc.match)
+
+			if tc.wantErr {
+				require.ErrorIs(t, err, ErrMalformedReference)
+				return
+			}
+
+			require.NoError(t, err)
+			require.Equal(t, tc.want, ref)
+		})
+	}
+}
+
+func TestCollectRefsMatchesResolverParsing(t *testing.T) {
+	t.Parallel()
+
+	doc := map[string]any{
+		"whole":     "kv://vault/secret/data/app#password",
+		"inline":    "postgres://$kv{vault:db/creds#host}:$kv{env:DB_PORT}/prod",
+		"dupe":      "kv://vault/secret/data/app#password", // same target as "whole"
+		"malformed": "kv://vault",                          // no path — skipped
+		"literal":   "no references here",
+		"nested": map[string]any{
+			"list": []any{"kv://consul/services/web", "plain"},
+		},
+	}
+
+	got := collectRefs(doc)
+
+	want := map[refKey]struct{}{
+		{store: "vault", path: "secret/data/app", fragment: "password"}: {},
+		{store: "vault", path: "db/creds", fragment: "host"}:            {},
+		{store: "env", path: "DB_PORT"}:                                 {},
+		{store: "consul", path: "services/web"}:                         {},
+	}
+
+	require.Equal(t, want, got)
+}