From 409a89240c53bf998a048d80fab32654d587549f Mon Sep 17 00:00:00 2001 From: srstack Date: Fri, 31 Jul 2026 21:40:07 +0800 Subject: [PATCH] fix: reclaim tenant metrics on pool close --- pkg/metrics/operations.go | 12 +++++++ pkg/metrics/registry.go | 15 +++++++++ pkg/server/auth.go | 2 ++ pkg/server/instrumentation.go | 51 ++++++++++++++++++++++++++++-- pkg/server/instrumentation_test.go | 24 ++++++++++++++ pkg/tenant/pool.go | 1 + pkg/tenant/pool_test.go | 42 +++++++++++++++++------- 7 files changed, 133 insertions(+), 14 deletions(-) diff --git a/pkg/metrics/operations.go b/pkg/metrics/operations.go index 05259a7f..31d0438a 100644 --- a/pkg/metrics/operations.go +++ b/pkg/metrics/operations.go @@ -290,6 +290,18 @@ func DeleteTenantCounters(tenantID string) { globalRegistry.DeleteGaugesByLabel("tenant_id", tenantID) } +// DeleteTenantRequestCounters removes only request-counter series for a +// tenant. Cache eviction intentionally resets request counters without +// removing tenant gauges or histograms whose values do not share counter-reset +// semantics. +func DeleteTenantRequestCounters(tenantID string) { + tenantID = cleanMetricValue(tenantID, "unknown") + if tenantID == "unknown" { + return + } + globalRegistry.DeleteCounterByLabel("drive9_tenant_requests_total", "tenant_id", tenantID) +} + func RecordTiDBCloudRBACCacheRequest(path, scope, result string) { RegisterModule("tidbcloud_quota") tidbCloudRBACCacheRequestsTotal.Add(1, diff --git a/pkg/metrics/registry.go b/pkg/metrics/registry.go index 488056c3..be082582 100644 --- a/pkg/metrics/registry.go +++ b/pkg/metrics/registry.go @@ -342,6 +342,21 @@ func (r *Registry) DeleteCountersByLabel(labelKey, labelValue string) { } } +func (r *Registry) DeleteCounterByLabel(name, labelKey, labelValue string) { + r.mu.Lock() + defer r.mu.Unlock() + inst := r.counters[name] + if inst == nil { + return + } + match := labelKey + `="` + EscapePromLabel(labelValue) + `"` + for labels := range inst.values { + if labelHasKeyValue(labels, match) { + delete(inst.values, labels) + } + } +} + func (r *Registry) DeleteHistogramsByLabel(labelKey, labelValue string) { r.mu.Lock() defer r.mu.Unlock() diff --git a/pkg/server/auth.go b/pkg/server/auth.go index ca172195..9b4611b8 100644 --- a/pkg/server/auth.go +++ b/pkg/server/auth.go @@ -388,6 +388,7 @@ func tenantAuthMiddlewareWithFSScopeLoader(metaStore *meta.Store, pool *tenant.P return } defer release() + defer func() { recordTenantHTTPRequestAtCompletion(r.Context()) }() metricEvent(r.Context(), "auth", "result", "ok") totalDuration := time.Since(authStart) logger.InfoOpenPoolTiming(r.Context(), "tenant_auth_timing", totalDuration, @@ -530,6 +531,7 @@ func (s *Server) capabilityAuthMiddleware(metaStore *meta.Store, pool *tenant.Po return } defer release() + defer func() { recordTenantHTTPRequestAtCompletion(r.Context()) }() scope := &TenantScope{TenantID: tenantID, Provider: tenant.Provider, TiDBCloudOrgID: b.TiDBCloudOrgID(), Backend: b} setRequestMetricScope(r.Context(), scope, classifyTenantRequest(r)) diff --git a/pkg/server/instrumentation.go b/pkg/server/instrumentation.go index 67080ed3..81e269af 100644 --- a/pkg/server/instrumentation.go +++ b/pkg/server/instrumentation.go @@ -75,6 +75,9 @@ type requestMetricState struct { action string tenantInFlight bool + tenantHTTPRequestRecorded bool + tenantHTTPRequestRecorder func() + bodyReadObserved bool bodyReadMethod string bodyReadRoute string @@ -117,6 +120,37 @@ func requestMetricStateFromContext(ctx context.Context) *requestMetricState { return v } +func setRequestTenantHTTPRecorder(ctx context.Context, recorder func()) { + state := requestMetricStateFromContext(ctx) + if state == nil { + return + } + state.mu.Lock() + state.tenantHTTPRequestRecorder = recorder + state.mu.Unlock() +} + +// recordTenantHTTPRequestAtCompletion records tenant HTTP usage at most once. +// Auth middleware invokes it before releasing the final backend reference so a +// closeEntry cleanup cannot be followed by an outer observe re-creation of the +// same tenant request counter. observe invokes it as a fallback for requests +// that do not acquire a tenant backend through auth middleware. +func recordTenantHTTPRequestAtCompletion(ctx context.Context) { + state := requestMetricStateFromContext(ctx) + if state == nil { + return + } + state.mu.Lock() + if state.tenantHTTPRequestRecorded || state.tenantHTTPRequestRecorder == nil { + state.mu.Unlock() + return + } + state.tenantHTTPRequestRecorded = true + recorder := state.tenantHTTPRequestRecorder + state.mu.Unlock() + recorder() +} + // markSSEStreamEstablished signals that the SSE handler has written the // streaming response headers and entered its long-lived loop. It records // the establishment timestamp so observe can measure the time-to-first-byte @@ -750,6 +784,19 @@ func (s *Server) observe(next http.Handler, w http.ResponseWriter, r *http.Reque if strings.HasPrefix(r.URL.Path, "/s3/") { setRequestMetricTenant(r.Context(), requestTenantID(r), "", "", "", classifyTenantRequest(r)) } + setRequestTenantHTTPRecorder(r.Context(), func() { + status := ow.status + if status == 0 { + status = http.StatusOK + } + dur := time.Since(start) + if route == sseEventsRoute && sseStreamEstablished(r.Context()) { + if estDur, ok := sseEstablishmentDuration(r.Context(), start); ok { + dur = estDur + } + } + recordTenantHTTPRequest(r, status, dur, ow.size) + }) next.ServeHTTP(rw, r) if ow.status == 0 { @@ -776,7 +823,7 @@ func (s *Server) observe(next http.Handler, w http.ResponseWriter, r *http.Reque dur = estDur } s.metrics.record(r.Method, route, ow.status, dur) - recordTenantHTTPRequest(r, ow.status, dur, ow.size) + recordTenantHTTPRequestAtCompletion(r.Context()) } else { s.metrics.record(r.Method, route, ow.status, dur) if method, bodyRoute, bodyBytes, bodyReadDuration, ok := requestBodyReadMetric(r.Context()); ok { @@ -788,7 +835,7 @@ func (s *Server) observe(next http.Handler, w http.ResponseWriter, r *http.Reque } s.metrics.recordBodyRead(method, bodyRoute, ow.status, bodyBytes, bodyReadDuration) } - recordTenantHTTPRequest(r, ow.status, dur, ow.size) + recordTenantHTTPRequestAtCompletion(r.Context()) } tenantID, apiKeyID, _, tidbCloudOrgID := requestMetricScope(r.Context()) diff --git a/pkg/server/instrumentation_test.go b/pkg/server/instrumentation_test.go index cbd283a0..bcc7ad56 100644 --- a/pkg/server/instrumentation_test.go +++ b/pkg/server/instrumentation_test.go @@ -186,6 +186,30 @@ func TestSetRequestMetricTenantForAuthStatusOnlyScopesActiveTenants(t *testing.T } } +func TestRecordTenantHTTPRequestAtCompletionDoesNotRecreateCleanedCounter(t *testing.T) { + const tenantID = "tenant-request-completion-cleanup" + metrics.DeleteTenantCounters(tenantID) + t.Cleanup(func() { metrics.DeleteTenantCounters(tenantID) }) + + ctx := withRequestMetricState(context.Background(), &requestMetricState{}) + setRequestTenantHTTPRecorder(ctx, func() { + metrics.RecordTenantRequestCountWithOrg(tenantID, "org-completion-cleanup", "fs", "read", http.StatusOK) + }) + + // The auth middleware runs this before it releases the final pool reference. + recordTenantHTTPRequestAtCompletion(ctx) + metrics.DeleteTenantRequestCounters(tenantID) + + // The outer observe fallback must not recreate a counter after closeEntry. + recordTenantHTTPRequestAtCompletion(ctx) + recorder := httptest.NewRecorder() + metrics.WritePrometheus(recorder) + if strings.Contains(recorder.Body.String(), `drive9_tenant_requests_total{`) && + strings.Contains(recorder.Body.String(), `tenant_id="`+tenantID+`"`) { + t.Fatalf("tenant request counter recreated after cleanup:\n%s", recorder.Body.String()) + } +} + func TestAdjustTenantInFlightAggregatesActionsForSameSurface(t *testing.T) { m := newServerMetrics() if got := m.adjustTenantInFlight("tenant-inflight-actions", "org-inflight-actions", "fs", 1); got != 1 { diff --git a/pkg/tenant/pool.go b/pkg/tenant/pool.go index 94201135..3ac1e8cf 100644 --- a/pkg/tenant/pool.go +++ b/pkg/tenant/pool.go @@ -1781,6 +1781,7 @@ func (p *Pool) closeEntry(e *entry) { _ = e.store.Close() } e.sharedDBLease.Release() + metrics.DeleteTenantRequestCounters(e.tenantID) } type tenantPoolResult string diff --git a/pkg/tenant/pool_test.go b/pkg/tenant/pool_test.go index eae12a11..2b6e2545 100644 --- a/pkg/tenant/pool_test.go +++ b/pkg/tenant/pool_test.go @@ -1862,27 +1862,45 @@ func TestAcquireRefreshesStandaloneIdleTTL(t *testing.T) { assertStoreOpen(t, b.Store()) } -func TestCloseEntryPreservesLiveTenantCounters(t *testing.T) { - const tenantID = "tenant-live-counter-close-entry" +func TestCloseEntryCleansTenantRequestCounters(t *testing.T) { + const tenantID = "tenant-clean-counter-close-entry" metrics.DeleteTenantCounters(tenantID) t.Cleanup(func() { metrics.DeleteTenantCounters(tenantID) }) metrics.RecordTenantRequestCountWithOrg(tenantID, "org-live-counter", "fs", "read", 200) - assertTenantMetricPresent := func() { - t.Helper() - recorder := httptest.NewRecorder() - metrics.WritePrometheus(recorder) - if !strings.Contains(recorder.Body.String(), `drive9_tenant_requests_total{`) || - !strings.Contains(recorder.Body.String(), `tenant_id="`+tenantID+`"`) { - t.Fatalf("live tenant counter disappeared from metrics:\n%s", recorder.Body.String()) - } + metrics.RecordTenantInFlightWithOrg(tenantID, "org-live-counter", "fs", 2) + recorder := httptest.NewRecorder() + metrics.WritePrometheus(recorder) + if !strings.Contains(recorder.Body.String(), `drive9_tenant_requests_total{`) || + !strings.Contains(recorder.Body.String(), `tenant_id="`+tenantID+`"`) { + t.Fatalf("tenant request counter missing before close:\n%s", recorder.Body.String()) + } + if !strings.Contains(recorder.Body.String(), `drive9_tenant_inflight_requests{`) || + !strings.Contains(recorder.Body.String(), `tenant_id="`+tenantID+`"`) { + t.Fatalf("tenant inflight gauge missing before close:\n%s", recorder.Body.String()) } - assertTenantMetricPresent() pool := NewPool(PoolConfig{}, nil) pool.closeEntry(&entry{tenantID: tenantID}) - assertTenantMetricPresent() + recorder = httptest.NewRecorder() + metrics.WritePrometheus(recorder) + if metricHasTenantID(t, recorder.Body.String(), "drive9_tenant_requests_total", tenantID) { + t.Fatalf("tenant request counter remained after close:\n%s", recorder.Body.String()) + } + if !metricHasTenantID(t, recorder.Body.String(), "drive9_tenant_inflight_requests", tenantID) { + t.Fatalf("tenant inflight gauge disappeared after close:\n%s", recorder.Body.String()) + } +} + +func metricHasTenantID(t *testing.T, text, metric, tenantID string) bool { + t.Helper() + for _, line := range strings.Split(text, "\n") { + if strings.HasPrefix(line, metric+"{") && strings.Contains(line, `tenant_id="`+tenantID+`"`) { + return true + } + } + return false } func TestIdleEvictionSkippedByRecentAcquire(t *testing.T) {