Take into account curried labels in promhttp

prometheus/promhttp/instrument_client.go

@@ -45,12 +45,11 @@ func InstrumentRoundTripperInFlight(gauge prometheus.Gauge, next http.RoundTripp
 
 // InstrumentRoundTripperCounter is a middleware that wraps the provided
 // http.RoundTripper to observe the request result with the provided CounterVec.
-// The CounterVec must have zero, one, or two labels. The only allowed label
+// The CounterVec must have zero, one, or two non-const non-curried labels. For
-// names are "code" and "method". The function panics if any other instance
+// those, the only allowed label names are "code" and "method". The function
-// labels are provided. Partitioning of the CounterVec happens by HTTP status
+// panics otherwise. Partitioning of the CounterVec happens by HTTP status code
-// code and/or HTTP method if the respective instance label names are present
+// and/or HTTP method if the respective instance label names are present in the
-// in the CounterVec. For unpartitioned counting, use a CounterVec with
+// CounterVec. For unpartitioned counting, use a CounterVec with zero labels.
-// zero labels.
 //
 // If the wrapped RoundTripper panics or returns a non-nil error, the Counter
 // is not incremented.

prometheus/promhttp/instrument_server.go

@@ -14,6 +14,7 @@
 package promhttp
 
 import (
+	"errors"
 	"net/http"
 	"strconv"
 	"strings"
@@ -42,10 +43,10 @@ func InstrumentHandlerInFlight(g prometheus.Gauge, next http.Handler) http.Handl
 
 // InstrumentHandlerDuration is a middleware that wraps the provided
 // http.Handler to observe the request duration with the provided ObserverVec.
-// The ObserverVec must have zero, one, or two labels. The only allowed label
+// The ObserverVec must have zero, one, or two non-const non-curried labels. For
-// names are "code" and "method". The function panics if any other instance
+// those, the only allowed label names are "code" and "method". The function
-// labels are provided. The Observe method of the Observer in the ObserverVec
+// panics otherwise. The Observe method of the Observer in the ObserverVec is
-// is called with the request duration in seconds. Partitioning happens by HTTP
+// called with the request duration in seconds. Partitioning happens by HTTP
 // status code and/or HTTP method if the respective instance label names are
 // present in the ObserverVec. For unpartitioned observations, use an
 // ObserverVec with zero labels. Note that partitioning of Histograms is
@@ -77,14 +78,13 @@ func InstrumentHandlerDuration(obs prometheus.ObserverVec, next http.Handler) ht
 	})
 }
 
-// InstrumentHandlerCounter is a middleware that wraps the provided
+// InstrumentHandlerCounter is a middleware that wraps the provided http.Handler
-// http.Handler to observe the request result with the provided CounterVec.
+// to observe the request result with the provided CounterVec.  The CounterVec
-// The CounterVec must have zero, one, or two labels. The only allowed label
+// must have zero, one, or two non-const non-curried labels. For those, the only
-// names are "code" and "method". The function panics if any other instance
+// allowed label names are "code" and "method". The function panics
-// labels are provided. Partitioning of the CounterVec happens by HTTP status
+// otherwise. Partitioning of the CounterVec happens by HTTP status code and/or
-// code and/or HTTP method if the respective instance label names are present
+// HTTP method if the respective instance label names are present in the
-// in the CounterVec. For unpartitioned counting, use a CounterVec with
+// CounterVec. For unpartitioned counting, use a CounterVec with zero labels.
-// zero labels.
 //
 // If the wrapped Handler does not set a status code, a status code of 200 is assumed.
 //
@@ -111,14 +111,13 @@ func InstrumentHandlerCounter(counter *prometheus.CounterVec, next http.Handler)
 // InstrumentHandlerTimeToWriteHeader is a middleware that wraps the provided
 // http.Handler to observe with the provided ObserverVec the request duration
 // until the response headers are written. The ObserverVec must have zero, one,
-// or two labels. The only allowed label names are "code" and "method". The
+// or two non-const non-curried labels. For those, the only allowed label names
-// function panics if any other instance labels are provided. The Observe
+// are "code" and "method". The function panics otherwise. The Observe method of
-// method of the Observer in the ObserverVec is called with the request
+// the Observer in the ObserverVec is called with the request duration in
-// duration in seconds. Partitioning happens by HTTP status code and/or HTTP
+// seconds. Partitioning happens by HTTP status code and/or HTTP method if the
-// method if the respective instance label names are present in the
+// respective instance label names are present in the ObserverVec. For
-// ObserverVec. For unpartitioned observations, use an ObserverVec with zero
+// unpartitioned observations, use an ObserverVec with zero labels. Note that
-// labels. Note that partitioning of Histograms is expensive and should be used
+// partitioning of Histograms is expensive and should be used judiciously.
-// judiciously.
 //
 // If the wrapped Handler panics before calling WriteHeader, no value is
 // reported.
@@ -140,15 +139,15 @@ func InstrumentHandlerTimeToWriteHeader(obs prometheus.ObserverVec, next http.Ha
 }
 
 // InstrumentHandlerRequestSize is a middleware that wraps the provided
-// http.Handler to observe the request size with the provided ObserverVec.
+// http.Handler to observe the request size with the provided ObserverVec.  The
-// The ObserverVec must have zero, one, or two labels. The only allowed label
+// ObserverVec must have zero, one, or two non-const non-curried labels. For
-// names are "code" and "method". The function panics if any other instance
+// those, the only allowed label names are "code" and "method". The function
-// labels are provided. The Observe method of the Observer in the ObserverVec
+// panics otherwise. The Observe method of the Observer in the ObserverVec is
-// is called with the request size in bytes. Partitioning happens by HTTP
+// called with the request size in bytes. Partitioning happens by HTTP status
-// status code and/or HTTP method if the respective instance label names are
+// code and/or HTTP method if the respective instance label names are present in
-// present in the ObserverVec. For unpartitioned observations, use an
+// the ObserverVec. For unpartitioned observations, use an ObserverVec with zero
-// ObserverVec with zero labels. Note that partitioning of Histograms is
+// labels. Note that partitioning of Histograms is expensive and should be used
-// expensive and should be used judiciously.
+// judiciously.
 //
 // If the wrapped Handler does not set a status code, a status code of 200 is assumed.
 //
@@ -175,15 +174,15 @@ func InstrumentHandlerRequestSize(obs prometheus.ObserverVec, next http.Handler)
 }
 
 // InstrumentHandlerResponseSize is a middleware that wraps the provided
-// http.Handler to observe the response size with the provided ObserverVec.
+// http.Handler to observe the response size with the provided ObserverVec.  The
-// The ObserverVec must have zero, one, or two labels. The only allowed label
+// ObserverVec must have zero, one, or two non-const non-curried labels. For
-// names are "code" and "method". The function panics if any other instance
+// those, the only allowed label names are "code" and "method". The function
-// labels are provided. The Observe method of the Observer in the ObserverVec
+// panics otherwise. The Observe method of the Observer in the ObserverVec is
-// is called with the response size in bytes. Partitioning happens by HTTP
+// called with the response size in bytes. Partitioning happens by HTTP status
-// status code and/or HTTP method if the respective instance label names are
+// code and/or HTTP method if the respective instance label names are present in
-// present in the ObserverVec. For unpartitioned observations, use an
+// the ObserverVec. For unpartitioned observations, use an ObserverVec with zero
-// ObserverVec with zero labels. Note that partitioning of Histograms is
+// labels. Note that partitioning of Histograms is expensive and should be used
-// expensive and should be used judiciously.
+// judiciously.
 //
 // If the wrapped Handler does not set a status code, a status code of 200 is assumed.
 //
@@ -204,9 +203,12 @@ func checkLabels(c prometheus.Collector) (code bool, method bool) {
 	// once Descriptors can have their dimensionality queried.
 	var (
 		desc *prometheus.Desc
+		m    prometheus.Metric
 		pm   dto.Metric
+		lvs  []string
 	)
 
+	// Get the Desc from the Collector.
 	descc := make(chan *prometheus.Desc, 1)
 	c.Describe(descc)
 
@@ -223,49 +225,54 @@ func checkLabels(c prometheus.Collector) (code bool, method bool) {
 
 	close(descc)
 
-	if _, err := prometheus.NewConstMetric(desc, prometheus.UntypedValue, 0); err == nil {
+	// Create a ConstMetric with the Desc. Since we don't know how many
-		return
+	// variable labels there are, try for as long as it needs.
+	for err := errors.New("dummy"); err != nil; lvs = append(lvs, magicString) {
+		m, err = prometheus.NewConstMetric(desc, prometheus.UntypedValue, 0, lvs...)
 	}
-	if m, err := prometheus.NewConstMetric(desc, prometheus.UntypedValue, 0, magicString); err == nil {
+
-		if err := m.Write(&pm); err != nil {
+	// Write out the metric into a proto message and look at the labels.
-			panic("error checking metric for labels")
+	// If the value is not the magicString, it is a constLabel, which doesn't interest us.
-		}
+	// If the label is curried, it doesn't interest us.
-		for _, label := range pm.Label {
+	// In all other cases, only "code" or "method" is allowed.
-			name, value := label.GetName(), label.GetValue()
+	if err := m.Write(&pm); err != nil {
-			if value != magicString {
+		panic("error checking metric for labels")
-				continue
-			}
-			switch name {
-			case "code":
-				code = true
-			case "method":
-				method = true
-			default:
-				panic("metric partitioned with non-supported labels")
-			}
-			return
-		}
-		panic("previously set label not found – this must never happen")
 	}
-	if m, err := prometheus.NewConstMetric(desc, prometheus.UntypedValue, 0, magicString, magicString); err == nil {
+	for _, label := range pm.Label {
-		if err := m.Write(&pm); err != nil {
+		name, value := label.GetName(), label.GetValue()
-			panic("error checking metric for labels")
+		if value != magicString || isLabelCurried(c, name) {
+			continue
 		}
-		for _, label := range pm.Label {
+		switch name {
-			name, value := label.GetName(), label.GetValue()
+		case "code":
-			if value != magicString {
+			code = true
-				continue
+		case "method":
-			}
+			method = true
-			if name == "code" || name == "method" {
+		default:
-				continue
-			}
 			panic("metric partitioned with non-supported labels")
 		}
-		code = true
-		method = true
-		return
 	}
-	panic("metric partitioned with non-supported labels")
+	return
+}
+
+func isLabelCurried(c prometheus.Collector, label string) bool {
+	// This is even hackier than the label test above.
+	// We essentially try to curry again and see if it works.
+	// But for that, we need to type-convert to the two
+	// types we use here, ObserverVec or *CounterVec.
+	switch v := c.(type) {
+	case *prometheus.CounterVec:
+		if _, err := v.CurryWith(prometheus.Labels{label: "dummy"}); err == nil {
+			return false
+		}
+	case prometheus.ObserverVec:
+		if _, err := v.CurryWith(prometheus.Labels{label: "dummy"}); err == nil {
+			return false
+		}
+	default:
+		panic("unsupported metric vec type")
+	}
+	return true
 }
 
 // emptyLabels is a one-time allocation for non-partitioned metrics to avoid

prometheus/promhttp/instrument_server_test.go

@@ -23,6 +23,162 @@ import (
 	"github.com/prometheus/client_golang/prometheus"
 )
 
+func TestLabelCheck(t *testing.T) {
+	scenarios := map[string]struct {
+		varLabels     []string
+		constLabels   []string
+		curriedLabels []string
+		ok            bool
+	}{
+		"empty": {
+			varLabels:     []string{},
+			constLabels:   []string{},
+			curriedLabels: []string{},
+			ok:            true,
+		},
+		"code as single var label": {
+			varLabels:     []string{"code"},
+			constLabels:   []string{},
+			curriedLabels: []string{},
+			ok:            true,
+		},
+		"method as single var label": {
+			varLabels:     []string{"method"},
+			constLabels:   []string{},
+			curriedLabels: []string{},
+			ok:            true,
+		},
+		"cade and method as var labels": {
+			varLabels:     []string{"method", "code"},
+			constLabels:   []string{},
+			curriedLabels: []string{},
+			ok:            true,
+		},
+		"valid case with all labels used": {
+			varLabels:     []string{"code", "method"},
+			constLabels:   []string{"foo", "bar"},
+			curriedLabels: []string{"dings", "bums"},
+			ok:            true,
+		},
+		"unsupported var label": {
+			varLabels:     []string{"foo"},
+			constLabels:   []string{},
+			curriedLabels: []string{},
+			ok:            false,
+		},
+		"mixed var labels": {
+			varLabels:     []string{"method", "foo", "code"},
+			constLabels:   []string{},
+			curriedLabels: []string{},
+			ok:            false,
+		},
+		"unsupported var label but curried": {
+			varLabels:     []string{},
+			constLabels:   []string{},
+			curriedLabels: []string{"foo"},
+			ok:            true,
+		},
+		"mixed var labels but unsupported curried": {
+			varLabels:     []string{"code", "method"},
+			constLabels:   []string{},
+			curriedLabels: []string{"foo"},
+			ok:            true,
+		},
+		"supported label as const and curry": {
+			varLabels:     []string{},
+			constLabels:   []string{"code"},
+			curriedLabels: []string{"method"},
+			ok:            true,
+		},
+		"supported label as const and curry with unsupported as var": {
+			varLabels:     []string{"foo"},
+			constLabels:   []string{"code"},
+			curriedLabels: []string{"method"},
+			ok:            false,
+		},
+	}
+
+	for name, sc := range scenarios {
+		t.Run(name, func(t *testing.T) {
+			constLabels := prometheus.Labels{}
+			for _, l := range sc.constLabels {
+				constLabels[l] = "dummy"
+			}
+			c := prometheus.NewCounterVec(
+				prometheus.CounterOpts{
+					Name:        "c",
+					Help:        "c help",
+					ConstLabels: constLabels,
+				},
+				append(sc.varLabels, sc.curriedLabels...),
+			)
+			o := prometheus.ObserverVec(prometheus.NewHistogramVec(
+				prometheus.HistogramOpts{
+					Name:        "c",
+					Help:        "c help",
+					ConstLabels: constLabels,
+				},
+				append(sc.varLabels, sc.curriedLabels...),
+			))
+			for _, l := range sc.curriedLabels {
+				c = c.MustCurryWith(prometheus.Labels{l: "dummy"})
+				o = o.MustCurryWith(prometheus.Labels{l: "dummy"})
+			}
+
+			func() {
+				defer func() {
+					if err := recover(); err != nil {
+						if sc.ok {
+							t.Error("unexpected panic:", err)
+						}
+					} else if !sc.ok {
+						t.Error("expected panic")
+					}
+				}()
+				InstrumentHandlerCounter(c, nil)
+			}()
+			func() {
+				defer func() {
+					if err := recover(); err != nil {
+						if sc.ok {
+							t.Error("unexpected panic:", err)
+						}
+					} else if !sc.ok {
+						t.Error("expected panic")
+					}
+				}()
+				InstrumentHandlerDuration(o, nil)
+			}()
+			if sc.ok {
+				// Test if wantCode and wantMethod were detected correctly.
+				var wantCode, wantMethod bool
+				for _, l := range sc.varLabels {
+					if l == "code" {
+						wantCode = true
+					}
+					if l == "method" {
+						wantMethod = true
+					}
+				}
+				gotCode, gotMethod := checkLabels(c)
+				if gotCode != wantCode {
+					t.Errorf("wanted code=%t for counter, got code=%t", wantCode, gotCode)
+				}
+				if gotMethod != wantMethod {
+					t.Errorf("wanted method=%t for counter, got method=%t", wantMethod, gotMethod)
+				}
+				gotCode, gotMethod = checkLabels(o)
+				if gotCode != wantCode {
+					t.Errorf("wanted code=%t for observer, got code=%t", wantCode, gotCode)
+				}
+				if gotMethod != wantMethod {
+					t.Errorf("wanted method=%t for observer, got method=%t", wantMethod, gotMethod)
+				}
+			}
+		})
+	}
+}
+
 func TestMiddlewareAPI(t *testing.T) {
 	reg := prometheus.NewRegistry()