[codex] Add latency phase metrics (#634)

* Add latency phase metrics

* Rename upstream latency metric

* Fix upstream timing boundary

* Align request metrics with latency counts

* Clarify unknown request metrics

Author: winhowes

app/integration.go

@@ -234,11 +234,16 @@ func prepareIntegration(i *Integration) error {
 	}
 
 	i.proxy.ModifyResponse = func(resp *http.Response) error {
+		start := time.Now()
+		if t, ok := metrics.UpstreamRoundtripStart(resp.Request.Context()); ok {
+			metrics.RecordUpstreamRoundtripDuration(i.Name, time.Since(t))
+		}
 		caller := metrics.Caller(resp.Request.Context())
 		metrics.OnResponse(i.Name, caller, resp.Request, resp)
 		if resp.StatusCode < http.StatusOK || resp.StatusCode >= 300 {
 			resp.Header.Set("X-AT-Upstream-Error", "true")
 		}
+		metrics.RecordResponseProcessingDuration(i.Name, time.Since(start))
 		return nil
 	}
 

app/main.go

@@ -1296,6 +1296,16 @@ const (
 
 // proxyHandler handles incoming requests and proxies them according to the integration.
 func proxyHandler(w http.ResponseWriter, r *http.Request) {
+	start := time.Now()
+	integrationName := "unknown"
+	preProxyRecorded := false
+	defer func() {
+		if !preProxyRecorded {
+			metrics.RecordPreProxyDuration(integrationName, time.Since(start))
+		}
+		metrics.RecordEndToEndDuration(integrationName, time.Since(start))
+	}()
+
 	host := r.Host
 	if !*disableXATInt {
 		hdr := r.Header.Get("X-AT-Int")
@@ -1315,6 +1325,8 @@ func proxyHandler(w http.ResponseWriter, r *http.Request) {
 		http.Error(w, fmt.Sprintf("integration for host %s not found", host), http.StatusNotFound)
 		return
 	}
+	integrationName = integ.Name
+	metrics.IncRequest(integ.Name)
 
 	clientIP, _, err := net.SplitHostPort(r.RemoteAddr)
 	if err != nil {
@@ -1463,6 +1475,10 @@ func proxyHandler(w http.ResponseWriter, r *http.Request) {
 	}
 
 	metrics.OnRequest(integ.Name, r)
+	handoffStart := time.Now()
+	r = r.WithContext(metrics.WithUpstreamRoundtripStart(r.Context(), handoffStart))
+	metrics.RecordPreProxyDuration(integ.Name, handoffStart.Sub(start))
+	preProxyRecorded = true
 	rec := &statusRecorder{ResponseWriter: w}
 	integ.proxy.ServeHTTP(rec, r)
 	if rec.status == 0 {

app/metrics/builtin.go

@@ -20,16 +20,17 @@ import (
 const metricKeySeparator = "|"
 
 var (
-	requestCounts          = expvar.NewMap("authtranslator_requests_total")
-	rateLimitCounts        = expvar.NewMap("authtranslator_rate_limit_events_total")
-	authFailureCounts      = expvar.NewMap("authtranslator_auth_failures_total")
-	internalResponseCounts = expvar.NewMap("authtranslator_internal_responses_total")
-	upstreamStatusCounts   = expvar.NewMap("authtranslator_upstream_responses_total")
-	requestDurations       = expvar.NewMap("authtranslator_request_duration_seconds")
-	LastReloadTime         = expvar.NewString("authtranslator_last_reload")
-	durationHistsMu        sync.Mutex
-	durationHists          = make(map[string]*histogram)
-	durationBuckets        = []float64{0.1, 0.25, 0.5, 1, 2.5, 5, 10}
+	requestCounts               = expvar.NewMap("authtranslator_requests_total")
+	rateLimitCounts             = expvar.NewMap("authtranslator_rate_limit_events_total")
+	authFailureCounts           = expvar.NewMap("authtranslator_auth_failures_total")
+	internalResponseCounts      = expvar.NewMap("authtranslator_internal_responses_total")
+	upstreamStatusCounts        = expvar.NewMap("authtranslator_upstream_responses_total")
+	upstreamRoundtripDurations  = newDurationMetric("authtranslator_upstream_roundtrip_duration_seconds")
+	endToEndDurations           = newDurationMetric("authtranslator_end_to_end_duration_seconds")
+	preProxyDurations           = newDurationMetric("authtranslator_pre_proxy_duration_seconds")
+	responseProcessingDurations = newDurationMetric("authtranslator_response_processing_duration_seconds")
+	LastReloadTime              = expvar.NewString("authtranslator_last_reload")
+	durationBuckets             = []float64{0.1, 0.25, 0.5, 1, 2.5, 5, 10}
 )
 
 type histogram struct {
@@ -39,13 +40,29 @@ type histogram struct {
 	sum     float64
 }
 
+type durationMetric struct {
+	name string
+	exp  *expvar.Map
+
+	mu    sync.Mutex
+	hists map[string]*histogram
+}
+
 func newHistogram() *histogram {
 	return &histogram{
 		buckets: durationBuckets,
 		counts:  make([]uint64, len(durationBuckets)+1),
 	}
 }
 
+func newDurationMetric(name string) *durationMetric {
+	return &durationMetric{
+		name:  name,
+		exp:   expvar.NewMap(name),
+		hists: make(map[string]*histogram),
+	}
+}
+
 func init() {
 	LastReloadTime.Set(time.Now().Format(time.RFC3339))
 }
@@ -89,7 +106,7 @@ func (h *histogram) String() string {
 	return string(buf)
 }
 
-func (h *histogram) writeProm(w http.ResponseWriter, integ string) {
+func (h *histogram) writeProm(w http.ResponseWriter, name, integ string) {
 	h.mu.Lock()
 	buckets := append([]float64(nil), h.buckets...)
 	counts := append([]uint64(nil), h.counts...)
@@ -98,12 +115,51 @@ func (h *histogram) writeProm(w http.ResponseWriter, integ string) {
 	var cum uint64
 	for i, b := range buckets {
 		cum += counts[i]
-		fmt.Fprintf(w, "authtranslator_request_duration_seconds_bucket{integration=%q,le=%q} %d\n", integ, strconv.FormatFloat(b, 'f', -1, 64), cum)
+		fmt.Fprintf(w, "%s_bucket{integration=%q,le=%q} %d\n", name, integ, strconv.FormatFloat(b, 'f', -1, 64), cum)
 	}
 	cum += counts[len(buckets)]
-	fmt.Fprintf(w, "authtranslator_request_duration_seconds_bucket{integration=%q,le=\"+Inf\"} %d\n", integ, cum)
-	fmt.Fprintf(w, "authtranslator_request_duration_seconds_sum{integration=%q} %f\n", integ, sum)
-	fmt.Fprintf(w, "authtranslator_request_duration_seconds_count{integration=%q} %d\n", integ, cum)
+	fmt.Fprintf(w, "%s_bucket{integration=%q,le=\"+Inf\"} %d\n", name, integ, cum)
+	fmt.Fprintf(w, "%s_sum{integration=%q} %f\n", name, integ, sum)
+	fmt.Fprintf(w, "%s_count{integration=%q} %d\n", name, integ, cum)
+}
+
+func (m *durationMetric) Record(integration string, d time.Duration) {
+	m.mu.Lock()
+	h, ok := m.hists[integration]
+	if !ok {
+		h = newHistogram()
+		m.hists[integration] = h
+		m.exp.Set(integration, h)
+	}
+	m.mu.Unlock()
+	h.Observe(d.Seconds())
+}
+
+func (m *durationMetric) Reset() {
+	m.exp.Init()
+	m.mu.Lock()
+	m.hists = make(map[string]*histogram)
+	m.mu.Unlock()
+}
+
+func (m *durationMetric) WriteProm(w http.ResponseWriter) {
+	writePromType(w, m.name, "histogram")
+
+	type histSnapshot struct {
+		name string
+		h    *histogram
+	}
+
+	m.mu.Lock()
+	hists := make([]histSnapshot, 0, len(m.hists))
+	for name, h := range m.hists {
+		hists = append(hists, histSnapshot{name: name, h: h})
+	}
+	m.mu.Unlock()
+
+	for _, hs := range hists {
+		hs.h.writeProm(w, m.name, hs.name)
+	}
 }
 
 // IncRequest increments the request counter for the integration.
@@ -128,17 +184,29 @@ func RecordStatus(integration string, status int) {
 	upstreamStatusCounts.Add(key, 1)
 }
 
-// RecordDuration records the upstream request duration.
-func RecordDuration(integration string, d time.Duration) {
-	durationHistsMu.Lock()
-	h, ok := durationHists[integration]
-	if !ok {
-		h = newHistogram()
-		durationHists[integration] = h
-		requestDurations.Set(integration, h)
-	}
-	durationHistsMu.Unlock()
-	h.Observe(d.Seconds())
+// RecordUpstreamRoundtripDuration records the duration from proxy handoff until
+// AuthTranslator receives the upstream response.
+func RecordUpstreamRoundtripDuration(integration string, d time.Duration) {
+	upstreamRoundtripDurations.Record(integration, d)
+}
+
+// RecordEndToEndDuration records full request latency as observed by the
+// AuthTranslator handler.
+func RecordEndToEndDuration(integration string, d time.Duration) {
+	endToEndDurations.Record(integration, d)
+}
+
+// RecordPreProxyDuration records request-side processing time before
+// AuthTranslator proxies upstream or returns a local response.
+func RecordPreProxyDuration(integration string, d time.Duration) {
+	preProxyDurations.Record(integration, d)
+}
+
+// RecordResponseProcessingDuration records response-side processing time inside
+// AuthTranslator after an upstream response is received and before streaming the
+// body to the client begins.
+func RecordResponseProcessingDuration(integration string, d time.Duration) {
+	responseProcessingDurations.Record(integration, d)
 }
 
 func writePromType(w http.ResponseWriter, name, metricType string) {
@@ -152,20 +220,10 @@ func WriteProm(w http.ResponseWriter) {
 	requestCounts.Do(func(kv expvar.KeyValue) {
 		fmt.Fprintf(w, "authtranslator_requests_total{integration=%q} %s\n", kv.Key, kv.Value.String())
 	})
-	writePromType(w, "authtranslator_request_duration_seconds", "histogram")
-	durationHistsMu.Lock()
-	type histSnapshot struct {
-		name string
-		h    *histogram
-	}
-	hists := make([]histSnapshot, 0, len(durationHists))
-	for name, h := range durationHists {
-		hists = append(hists, histSnapshot{name: name, h: h})
-	}
-	durationHistsMu.Unlock()
-	for _, hs := range hists {
-		hs.h.writeProm(w, hs.name)
-	}
+	upstreamRoundtripDurations.WriteProm(w)
+	endToEndDurations.WriteProm(w)
+	preProxyDurations.WriteProm(w)
+	responseProcessingDurations.WriteProm(w)
 	writePromType(w, "authtranslator_rate_limit_events_total", "counter")
 	rateLimitCounts.Do(func(kv expvar.KeyValue) {
 		fmt.Fprintf(w, "authtranslator_rate_limit_events_total{integration=%q} %s\n", kv.Key, kv.Value.String())

app/metrics/metrics_histogram_test.go

@@ -50,7 +50,7 @@ func TestHistogramWriteProm(t *testing.T) {
 	h.Observe(1)
 
 	rr := httptest.NewRecorder()
-	h.writeProm(rr, "foo")
+	h.writeProm(rr, "authtranslator_upstream_roundtrip_duration_seconds", "foo")
 
 	lines := strings.Split(strings.TrimSpace(rr.Body.String()), "\n")
 	if len(lines) != 10 {
@@ -59,7 +59,7 @@ func TestHistogramWriteProm(t *testing.T) {
 	if !strings.Contains(lines[0], "le=\"0.1\"") || !strings.HasSuffix(lines[0], " 0") {
 		t.Fatalf("unexpected first line: %s", lines[0])
 	}
-	if !strings.HasPrefix(lines[8], "authtranslator_request_duration_seconds_sum{integration=\"foo\"}") {
+	if !strings.HasPrefix(lines[8], "authtranslator_upstream_roundtrip_duration_seconds_sum{integration=\"foo\"}") {
 		t.Fatalf("missing sum line")
 	}
 	if !strings.HasSuffix(lines[9], " 2") {

app/metrics/metrics_test.go

@@ -49,7 +49,10 @@ func TestMetricsHandlerEmpty(t *testing.T) {
 	body := rr.Body.String()
 	for _, line := range []string{
 		"# TYPE authtranslator_requests_total counter",
-		"# TYPE authtranslator_request_duration_seconds histogram",
+		"# TYPE authtranslator_upstream_roundtrip_duration_seconds histogram",
+		"# TYPE authtranslator_end_to_end_duration_seconds histogram",
+		"# TYPE authtranslator_pre_proxy_duration_seconds histogram",
+		"# TYPE authtranslator_response_processing_duration_seconds histogram",
 		"# TYPE authtranslator_rate_limit_events_total counter",
 		"# TYPE authtranslator_auth_failures_total counter",
 		"# TYPE authtranslator_internal_responses_total counter",
@@ -74,9 +77,15 @@ func TestMetricsHandlerOutput(t *testing.T) {
 	RecordStatus("foo", http.StatusOK)
 	RecordStatus("bar", http.StatusTeapot)
 	IncRequest("bar")
-	RecordDuration("foo", 100*time.Millisecond)
-	RecordDuration("foo", 200*time.Millisecond)
-	RecordDuration("bar", 50*time.Millisecond)
+	RecordUpstreamRoundtripDuration("foo", 100*time.Millisecond)
+	RecordUpstreamRoundtripDuration("foo", 200*time.Millisecond)
+	RecordUpstreamRoundtripDuration("bar", 50*time.Millisecond)
+	RecordEndToEndDuration("foo", 250*time.Millisecond)
+	RecordEndToEndDuration("bar", 75*time.Millisecond)
+	RecordPreProxyDuration("foo", 20*time.Millisecond)
+	RecordPreProxyDuration("bar", 10*time.Millisecond)
+	RecordResponseProcessingDuration("foo", 15*time.Millisecond)
+	RecordResponseProcessingDuration("bar", 5*time.Millisecond)
 
 	req := httptest.NewRequest(http.MethodGet, "/_at_internal/metrics", nil)
 	rr := httptest.NewRecorder()
@@ -89,7 +98,10 @@ func TestMetricsHandlerOutput(t *testing.T) {
 	body := rr.Body.String()
 	for _, line := range []string{
 		"# TYPE authtranslator_requests_total counter",
-		"# TYPE authtranslator_request_duration_seconds histogram",
+		"# TYPE authtranslator_upstream_roundtrip_duration_seconds histogram",
+		"# TYPE authtranslator_end_to_end_duration_seconds histogram",
+		"# TYPE authtranslator_pre_proxy_duration_seconds histogram",
+		"# TYPE authtranslator_response_processing_duration_seconds histogram",
 		"# TYPE authtranslator_rate_limit_events_total counter",
 		"# TYPE authtranslator_auth_failures_total counter",
 		"# TYPE authtranslator_internal_responses_total counter",
@@ -120,11 +132,20 @@ func TestMetricsHandlerOutput(t *testing.T) {
 	if !strings.Contains(body, `authtranslator_upstream_responses_total{integration="bar",code="418"} 1`) {
 		t.Fatal("missing bar status metric")
 	}
-	if !strings.Contains(body, `authtranslator_request_duration_seconds_sum{integration="foo"}`) {
-		t.Fatal("missing foo duration histogram")
+	if !strings.Contains(body, `authtranslator_upstream_roundtrip_duration_seconds_sum{integration="foo"}`) {
+		t.Fatal("missing foo upstream duration histogram")
 	}
-	if !strings.Contains(body, `authtranslator_request_duration_seconds_sum{integration="bar"}`) {
-		t.Fatal("missing bar duration histogram")
+	if !strings.Contains(body, `authtranslator_upstream_roundtrip_duration_seconds_sum{integration="bar"}`) {
+		t.Fatal("missing bar upstream duration histogram")
+	}
+	if !strings.Contains(body, `authtranslator_end_to_end_duration_seconds_sum{integration="foo"}`) {
+		t.Fatal("missing foo end-to-end duration histogram")
+	}
+	if !strings.Contains(body, `authtranslator_pre_proxy_duration_seconds_sum{integration="foo"}`) {
+		t.Fatal("missing foo pre-proxy duration histogram")
+	}
+	if !strings.Contains(body, `authtranslator_response_processing_duration_seconds_sum{integration="foo"}`) {
+		t.Fatal("missing foo response processing duration histogram")
 	}
 }
 
@@ -135,7 +156,7 @@ func TestMetricsHandlerOutputWithPunctuation(t *testing.T) {
 	underscoreName := "with_underscore"
 
 	IncRequest(dotName)
-	RecordDuration(dotName, 150*time.Millisecond)
+	RecordUpstreamRoundtripDuration(dotName, 150*time.Millisecond)
 	RecordStatus(dotName, http.StatusAccepted)
 	IncInternalResponse(dotName, http.StatusBadRequest, "invalid_destination")
 
@@ -149,7 +170,7 @@ func TestMetricsHandlerOutputWithPunctuation(t *testing.T) {
 	if !strings.Contains(body, fmt.Sprintf(`authtranslator_requests_total{integration=%q} 1`, dotName)) {
 		t.Fatalf("missing request counter for %s: %s", dotName, body)
 	}
-	if !strings.Contains(body, fmt.Sprintf(`authtranslator_request_duration_seconds_sum{integration=%q}`, dotName)) {
+	if !strings.Contains(body, fmt.Sprintf(`authtranslator_upstream_roundtrip_duration_seconds_sum{integration=%q}`, dotName)) {
 		t.Fatalf("missing duration histogram for %s: %s", dotName, body)
 	}
 	if !strings.Contains(body, fmt.Sprintf(`authtranslator_upstream_responses_total{integration=%q,code=%q} 1`, dotName, "202")) {
@@ -211,6 +232,17 @@ func TestCallerContext(t *testing.T) {
 	}
 }
 
+func TestUpstreamRoundtripStartContext(t *testing.T) {
+	start := time.Now()
+	ctx := WithUpstreamRoundtripStart(context.Background(), start)
+	if got, ok := UpstreamRoundtripStart(ctx); !ok || !got.Equal(start) {
+		t.Fatalf("expected upstream roundtrip start %v, got %v ok=%v", start, got, ok)
+	}
+	if _, ok := UpstreamRoundtripStart(context.Background()); ok {
+		t.Fatal("expected missing upstream roundtrip start on background context")
+	}
+}
+
 func TestWritePromPlugins(t *testing.T) {
 	Reset()
 	mu.Lock()
@@ -260,7 +292,7 @@ func TestWritePromSkipsMalformedUpstreamKeys(t *testing.T) {
 
 func TestWritePromDoesNotBlockRecordDurationForOtherIntegrations(t *testing.T) {
 	Reset()
-	RecordDuration("one", 100*time.Millisecond)
+	RecordUpstreamRoundtripDuration("one", 100*time.Millisecond)
 
 	blockCh := make(chan struct{})
 	started := make(chan struct{})
@@ -284,14 +316,14 @@ func TestWritePromDoesNotBlockRecordDurationForOtherIntegrations(t *testing.T) {
 
 	recordDone := make(chan struct{})
 	go func() {
-		RecordDuration("two", 50*time.Millisecond)
+		RecordUpstreamRoundtripDuration("two", 50*time.Millisecond)
 		close(recordDone)
 	}()
 
 	select {
 	case <-recordDone:
 	case <-time.After(300 * time.Millisecond):
-		t.Fatal("RecordDuration blocked while metrics response writer was blocked")
+		t.Fatal("RecordUpstreamRoundtripDuration blocked while metrics response writer was blocked")
 	}
 
 	close(blockCh)