[metrics]: Allow request metadata as Prometheus labels

README.md

@@ -75,7 +75,9 @@ reaches a different upstream with no change to the Worker.
   (the gRPC method, and the Namespace the proxy resolved from the request rather than from anything the caller claims),
   so it can decide per Namespace and per method rather than only whether the caller is who it says it is.
 - **Prometheus metrics.** Expose request latency and counts, routing decisions, and encryption activity on `/metrics`.
-  The listen address and the namespace prefixed onto every metric are set under `metrics:` in the config.
+  The listen address and the namespace prefixed onto every metric are set under `metrics:` in the config, which can
+  also name request metadata to carry onto the request-scoped metrics as extra labels, so they can be sliced by a
+  dimension only your callers know.
 - **Codec-transparent.** The gateway never parses payloads. It peeks the Namespace, picks an upstream, and relays raw
   frames in both directions.
 - **Multiple deployment options.** Ship as a Go binary, a container image, or a Helm chart.

e2e/metric_tags_test.go

@@ -0,0 +1,129 @@
+package e2e
+
+import (
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/stretchr/testify/require"
+	"go.temporal.io/api/common/v1"
+	"go.temporal.io/api/query/v1"
+	"go.temporal.io/api/workflowservice/v1"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/metadata"
+
+	"github.com/temporalio/temporal-proxy/internal/config"
+	"github.com/temporalio/temporal-proxy/internal/dataplane/dataplanetest"
+)
+
+// TestEndToEndMetricTags drives a payload-carrying call through the full stack
+// with a metrics tag configured, and proves the caller's metadata reaches the
+// labels of every request-scoped collector.
+//
+// The vault_ops assertion is the one that matters: that collector lives on the
+// per-upstream hop, on the far side of a unix socket that context values do not
+// cross. It passes only because the value is read from the metadata the gateway
+// forwards, and it is what fails if this is ever reworked into a single
+// interceptor stashing values in a context.
+func TestEndToEndMetricTags(t *testing.T) {
+	t.Parallel()
+
+	up := dataplanetest.NewUpstream(t)
+
+	cfg := dataplanetest.Config(up)
+	// Mixed case on purpose: gRPC canonicalizes metadata keys, so a config
+	// written this way still has to match what arrives on the wire.
+	cfg.Metrics.Tags = []config.MetricTag{{Header: "X-Tenant", Label: "tenant"}}
+	// Encryption is what makes the per-upstream hop emit vault_ops at all.
+	cfg.Encryption = config.Encryption{
+		Enabled: true,
+		Default: &config.KeyPolicy{URI: testingKeyURI(t), Duration: time.Hour},
+	}
+
+	f := dataplanetest.StartApp(t, cfg)
+
+	ctx := metadata.AppendToOutgoingContext(f.Context(), "x-tenant", "acme")
+	_, err := f.Client().QueryWorkflow(ctx, queryWithPayload(), grpc.WaitForReady(true))
+	require.NoError(t, err)
+
+	// Hop 1: the gateway's own interceptor and the routing decision.
+	requireLabel(t, f, "test_server_requests_total", "tenant", "acme")
+	requireLabel(t, f, "test_router_decisions_total", "tenant", "acme")
+
+	// Hop 2: across the socket.
+	requireLabel(t, f, "test_encryption_vault_ops_total", "tenant", "acme")
+}
+
+// TestEndToEndMetricTagsWithoutTheHeader proves a configured label is always
+// present, empty when the caller sent nothing, rather than the series being
+// reported without it.
+func TestEndToEndMetricTagsWithoutTheHeader(t *testing.T) {
+	t.Parallel()
+
+	up := dataplanetest.NewUpstream(t)
+
+	cfg := dataplanetest.Config(up)
+	cfg.Metrics.Tags = []config.MetricTag{{Header: "x-tenant", Label: "tenant"}}
+	cfg.Encryption = config.Encryption{
+		Enabled: true,
+		Default: &config.KeyPolicy{URI: testingKeyURI(t), Duration: time.Hour},
+	}
+
+	f := dataplanetest.StartApp(t, cfg)
+
+	_, err := f.Client().QueryWorkflow(f.Context(), queryWithPayload(), grpc.WaitForReady(true))
+	require.NoError(t, err)
+
+	requireLabel(t, f, "test_server_requests_total", "tenant", "")
+	requireLabel(t, f, "test_router_decisions_total", "tenant", "")
+	requireLabel(t, f, "test_encryption_vault_ops_total", "tenant", "")
+}
+
+// queryWithPayload is a QueryWorkflow request carrying one payload, so the
+// encryption interceptor on the per-upstream hop has something to seal.
+func queryWithPayload() *workflowservice.QueryWorkflowRequest {
+	return &workflowservice.QueryWorkflowRequest{
+		Namespace: "ns1",
+		Execution: &common.WorkflowExecution{WorkflowId: "wf-1"},
+		Query: &query.WorkflowQuery{
+			QueryType: "state",
+			QueryArgs: &common.Payloads{Payloads: []*common.Payload{{
+				Metadata: map[string][]byte{wireEncoding: []byte("json/plain")},
+				Data:     []byte(`"hi"`),
+			}}},
+		},
+	}
+}
+
+// requireLabel asserts that some series in the named metric family carries
+// label=want. It compares against the gathered registry rather than a scrape,
+// which is equivalent: the metrics module hands that same Gatherer to its
+// /metrics handler.
+func requireLabel(t *testing.T, f *dataplanetest.Fixture, family, label, want string) {
+	t.Helper()
+
+	found, err := testutil.GatherAndLint(f.Gatherer(), family)
+	require.NoError(t, err)
+	require.Empty(t, found, "collector %s failed prometheus linting", family)
+
+	mfs, err := f.Gatherer().Gather()
+	require.NoError(t, err)
+
+	for _, mf := range mfs {
+		if mf.GetName() != family {
+			continue
+		}
+
+		for _, m := range mf.GetMetric() {
+			for _, lp := range m.GetLabel() {
+				if lp.GetName() == label && lp.GetValue() == want {
+					return
+				}
+			}
+		}
+
+		t.Fatalf("metric family %s has no series with %s=%q, got %v", family, label, want, mf.GetMetric())
+	}
+
+	t.Fatalf("metric family %s was never registered", family)
+}

internal/config/metrics.go

@@ -1,23 +1,189 @@
 package config
 
-import "github.com/temporalio/temporal-proxy/pkg/validation"
+import (
+	"errors"
+	"fmt"
+	"regexp"
+	"slices"
+	"strings"
 
-// Metrics configures the Prometheus endpoint. HostPort is the address the
-// /metrics handler listens on, and Namespace is the prefix stamped onto every
-// collector: a Prometheus namespace, unrelated to a Temporal namespace. Load
-// defaults both, so neither is empty in a loaded config.
-type Metrics struct {
-	HostPort  string `yaml:"hostPort"`
-	Namespace string `yaml:"namespace"`
+	"github.com/temporalio/temporal-proxy/pkg/validation"
+)
+
+const (
+	binaryHeaderSuffix   = "-bin"  // gRPC binary metadata marker
+	reservedHeaderPrefix = "grpc-" // gRPC reserved metadata prefix
+	reservedLabelPrefix  = "__"    // Prometheus reserved label prefix
+)
+
+var (
+	// promReservedLabels are the label names Prometheus keeps for a histogram's
+	// bucket bound and a summary's quantile.
+	promReservedLabels = []string{"le", "quantile"}
+
+	// mdKeyRegex matches the characters gRPC allows in a metadata key.
+	mdKeyRegex = regexp.MustCompile(`^[a-zA-Z0-9._-]+$`)
+
+	// promLabelRegex matches the Prometheus label name grammar.
+	promLabelRegex = regexp.MustCompile(`^[a-zA-Z_][a-zA-Z0-9_]*$`)
+)
+
+type (
+	// Metrics configures the Prometheus endpoint. HostPort is the address the
+	// /metrics handler listens on, and Namespace is the prefix stamped onto every
+	// collector: a Prometheus namespace, unrelated to a Temporal namespace. Tags
+	// name the request metadata to carry through as labels. Load defaults
+	// HostPort and Namespace, so neither is empty in a loaded config.
+	Metrics struct {
+		HostPort  string      `yaml:"hostPort"`
+		Namespace string      `yaml:"namespace"`
+		Tags      []MetricTag `yaml:"tags"`
+	}
+
+	// MetricTag pairs an inbound request metadata Header with the Prometheus
+	// Label it is reported under. It is written in YAML as "<header>:<label>".
+	//
+	// Choose the Header with care. Its value is published on the /metrics
+	// endpoint, which is served unauthenticated, so naming a header that carries
+	// a credential exposes that credential to anything able to reach the port.
+	// A tag also multiplies every request-scoped series rather than adding to
+	// them, so a header the caller varies freely multiplies cardinality with it.
+	MetricTag struct {
+		Header string
+		Label  string
+	}
+)
+
+// ParseMetricTag parses the "<header>:<label>" form, trimming whitespace around
+// each half. It splits on the first colon, so a label containing one survives
+// parsing and is rejected by Validate instead.
+func ParseMetricTag(v string) (MetricTag, error) {
+	tag := MetricTag{}
+	hdr, lbl, ok := strings.Cut(v, ":")
+	if !ok {
+		return tag, errors.New("invalid metric tag: format: <header>:<label>")
+	}
+
+	tag.Header = strings.TrimSpace(hdr)
+	tag.Label = strings.TrimSpace(lbl)
+	return tag, nil
 }
 
-// Validate requires a valid host:port and a non-empty namespace. Load defaults
-// both, so a namespace failure is only reachable for a Metrics built directly,
-// and a hostPort failure only for a config that sets one that will not parse.
+// Validate requires a valid host:port and a non-empty namespace, and checks
+// every tag. Load defaults the first two, so a namespace failure is only
+// reachable for a Metrics built directly, and a hostPort failure only for a
+// config that sets one that will not parse. Two tags may not share a label,
+// which Prometheus rejects as a duplicate label name.
 func (m *Metrics) Validate() error {
+	labels := make([]string, len(m.Tags))
+	for i, t := range m.Tags {
+		labels[i] = t.Label
+	}
+
 	return validation.Validate(
 		"",
 		validation.Field("hostPort", m.HostPort, validation.IsHostPort()),
 		validation.Field("namespace", m.Namespace, validation.Required[string]()),
+		validation.Field("tags[label]", labels, validation.Unique[string]()),
+		validation.Children("tags", m.Tags, func(t *MetricTag) error { return t.Validate() }),
 	)
 }
+
+// UnmarshalYAML decodes the scalar "<header>:<label>" form, so a tag is written
+// as a plain YAML string rather than a mapping.
+func (t *MetricTag) UnmarshalYAML(unmarshal func(any) error) error {
+	var decoded string
+	if err := unmarshal(&decoded); err != nil {
+		return err
+	}
+
+	tag, err := ParseMetricTag(decoded)
+	if err != nil {
+		return err
+	}
+
+	*t = tag
+	return nil
+}
+
+// Validate requires a header that is a legal, unreserved, non-binary gRPC
+// metadata key and a label that is a legal, unreserved Prometheus label name.
+func (t *MetricTag) Validate() error {
+	return validation.Validate(
+		"",
+		validation.Field(
+			"header",
+			t.Header,
+			validation.Required[string](),
+			match(mdKeyRegex),
+			unreservedHeader(),
+			textualHeader(),
+		),
+		validation.Field(
+			"label",
+			t.Label,
+			validation.Required[string](),
+			match(promLabelRegex),
+			unreservedLabel(),
+		),
+	)
+}
+
+// match rejects a value that does not match r. An empty value yields nothing so
+// the Required check on the same field owns that case and reports it once.
+func match(r *regexp.Regexp) validation.Check[string] {
+	return func(s string) error {
+		if s == "" || r.MatchString(s) {
+			return nil
+		}
+
+		return fmt.Errorf("is not valid, must match: %q", r.String())
+	}
+}
+
+// textualHeader rejects a metadata key ending in "-bin". gRPC uses that suffix
+// to mark a value as binary, and binary is not valid UTF-8, so such a value
+// cannot be reported as a Prometheus label. The comparison is case-insensitive
+// because metadata keys are.
+func textualHeader() validation.Check[string] {
+	return func(s string) error {
+		if strings.HasSuffix(strings.ToLower(s), binaryHeaderSuffix) {
+			return fmt.Errorf("must not end with %q, which gRPC uses to mark binary metadata", binaryHeaderSuffix)
+		}
+
+		return nil
+	}
+}
+
+// unreservedHeader rejects a metadata key beginning with "grpc-", which gRPC
+// keeps for itself. The comparison is case-insensitive because metadata keys
+// are, so "GRPC-Status" is rejected alongside "grpc-status".
+func unreservedHeader() validation.Check[string] {
+	return func(s string) error {
+		if strings.HasPrefix(strings.ToLower(s), reservedHeaderPrefix) {
+			return fmt.Errorf("must not begin with %q, which gRPC reserves", reservedHeaderPrefix)
+		}
+
+		return nil
+	}
+}
+
+// unreservedLabel rejects a label name Prometheus keeps for itself: one
+// beginning with "__", which it refuses to register, and "le" or "quantile",
+// which name a histogram's bucket bound and a summary's quantile. The reserved
+// names are refused here because a histogram panics on "le" only when it first
+// instantiates a series, which happens while serving a request rather than at
+// construction, so nothing downstream would catch it.
+func unreservedLabel() validation.Check[string] {
+	return func(s string) error {
+		if strings.HasPrefix(s, reservedLabelPrefix) {
+			return fmt.Errorf("must not begin with %q, which Prometheus reserves", reservedLabelPrefix)
+		}
+
+		if slices.Contains(promReservedLabels, s) {
+			return fmt.Errorf("must not be %q, which Prometheus reserves", s)
+		}
+
+		return nil
+	}
+}