From 02861b5537fe237d91e271cff0659835b7b90ceb Mon Sep 17 00:00:00 2001
From: wjqserver <114663932+WJQSERVER@users.noreply.github.com>
Date: Fri, 10 Apr 2026 21:55:21 +0800
Subject: [PATCH] perf: avoid header policy join allocations

---
 reverseproxy_benchmark_test.go | 92 ++++++++++++++++++++++++++++++++++
 reverseproxy_lb.go             | 46 ++++++++++++++++-
 2 files changed, 137 insertions(+), 1 deletion(-)

diff --git a/reverseproxy_benchmark_test.go b/reverseproxy_benchmark_test.go
index f55f5f0..7a03bd4 100644
--- a/reverseproxy_benchmark_test.go
+++ b/reverseproxy_benchmark_test.go
@@ -7,6 +7,7 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"strings"
 	"testing"
 	"time"
 )
@@ -167,6 +168,33 @@ func BenchmarkReverseProxySelectUpstream(b *testing.B) {
 	}
 }
 
+func BenchmarkReverseProxySelectUpstreamHeaderPolicy(b *testing.B) {
+	proxy := &reverseProxyHandler{
+		upstreams: []*reverseProxyUpstream{
+			{key: "a", index: 0},
+			{key: "b", index: 1},
+			{key: "c", index: 2},
+			{key: "d", index: 3},
+		},
+		config: ReverseProxyConfig{
+			LoadBalancing: ReverseProxyLoadBalancingConfig{Policy: LBHeader("X-Tenant", LBRandom())},
+		},
+	}
+	c, _ := CreateTestContext(nil)
+	c.Request.Header["X-Tenant"] = []string{"tenant-a", "tenant-b", "tenant-c"}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		selected, err := proxy.selectUpstream(c, nil)
+		if err != nil {
+			b.Fatalf("selectUpstream failed: %v", err)
+		}
+		benchmarkReverseProxySink = selected.index
+	}
+}
+
 func TestReverseProxyCopyResponseWithoutBufferPool(t *testing.T) {
 	proxy := newReverseProxyHandler(ReverseProxyConfig{})
 	dst := newBenchmarkResponseWriter()
@@ -260,4 +288,68 @@ func TestReverseProxyAvailableUpstreamsFiltersExcludedAndUnhealthy(t *testing.T)
 	}
 }
 
+func TestReverseProxyHeaderPolicyUsesAllHeaderValues(t *testing.T) {
+	proxy := &reverseProxyHandler{
+		upstreams: []*reverseProxyUpstream{
+			{key: "a", index: 0},
+			{key: "b", index: 1},
+			{key: "c", index: 2},
+		},
+		config: ReverseProxyConfig{
+			LoadBalancing: ReverseProxyLoadBalancingConfig{Policy: LBHeader("X-Tenant", LBRandom())},
+		},
+	}
+
+	c, _ := CreateTestContext(nil)
+	c.Request.Header["X-Tenant"] = []string{"tenant-a", "tenant-b"}
+
+	selectedA, err := proxy.selectUpstream(c, nil)
+	if err != nil {
+		t.Fatalf("selectUpstream failed: %v", err)
+	}
+	selectedB, err := proxy.selectUpstream(c, nil)
+	if err != nil {
+		t.Fatalf("selectUpstream failed: %v", err)
+	}
+	if selectedA.key != selectedB.key {
+		t.Fatalf("expected stable selection for identical multi-value header, got %q and %q", selectedA.key, selectedB.key)
+	}
+
+	c.Request.Header["X-Tenant"] = []string{"tenant-b", "tenant-a"}
+	selectedC, err := proxy.selectUpstream(c, nil)
+	if err != nil {
+		t.Fatalf("selectUpstream failed: %v", err)
+	}
+	if selectedC == nil {
+		t.Fatal("expected upstream for reordered multi-value header")
+	}
+}
+
+func TestReverseProxyHeaderPolicyMatchesJoinCompatibility(t *testing.T) {
+	candidates := []*reverseProxyUpstream{
+		{key: "a", index: 0},
+		{key: "b", index: 1},
+		{key: "c", index: 2},
+	}
+
+	testCases := [][]string{
+		{"tenant-a"},
+		{"tenant-a", "tenant-b"},
+		{"", "tenant-b"},
+		{"tenant-a", ""},
+		{"", ""},
+	}
+
+	for _, values := range testCases {
+		got := reverseProxySelectHRWValues(candidates, values)
+		want := reverseProxySelectHRW(candidates, strings.Join(values, ","))
+		if got == nil || want == nil {
+			t.Fatalf("expected non-nil upstreams for values %v", values)
+		}
+		if got.key != want.key {
+			t.Fatalf("expected joined compatibility for values %v, got %q want %q", values, got.key, want.key)
+		}
+	}
+}
+
 var _ io.Writer = (*benchmarkResponseWriter)(nil)
diff --git a/reverseproxy_lb.go b/reverseproxy_lb.go
index 02895fb..3be7234 100644
--- a/reverseproxy_lb.go
+++ b/reverseproxy_lb.go
@@ -199,7 +199,7 @@ func (p *reverseProxyHandler) selectUpstreamWithPolicy(c *Context, candidates []
 	case reverseProxyLBPolicyHeader:
 		if c.Request != nil && c.Request.Header != nil {
 			if values, ok := c.Request.Header[policy.key]; ok {
-				return reverseProxySelectHRW(candidates, strings.Join(values, ","))
+				return reverseProxySelectHRWValues(candidates, values)
 			}
 		}
 		return p.selectUpstreamWithPolicy(c, candidates, reverseProxyFallbackPolicy(policy))
@@ -277,6 +277,25 @@ func reverseProxySelectHRW(candidates []*reverseProxyUpstream, key string) *reve
 	return selected
 }
 
+func reverseProxySelectHRWValues(candidates []*reverseProxyUpstream, values []string) *reverseProxyUpstream {
+	if len(candidates) == 0 {
+		return nil
+	}
+	if len(values) == 0 {
+		return reverseProxySelectRandom(candidates)
+	}
+	selected := candidates[0]
+	bestScore := reverseProxyHRWValuesScore(values, selected.key)
+	for _, upstream := range candidates[1:] {
+		score := reverseProxyHRWValuesScore(values, upstream.key)
+		if score > bestScore {
+			selected = upstream
+			bestScore = score
+		}
+	}
+	return selected
+}
+
 func reverseProxyHRWScore(key, upstreamKey string) uint64 {
 	const (
 		offset64 = 14695981039346656037
@@ -296,6 +315,31 @@ func reverseProxyHRWScore(key, upstreamKey string) uint64 {
 	return h
 }
 
+func reverseProxyHRWValuesScore(values []string, upstreamKey string) uint64 {
+	const (
+		offset64 = 14695981039346656037
+		prime64  = 1099511628211
+	)
+	h := uint64(offset64)
+	for valueIndex, value := range values {
+		for i := 0; i < len(value); i++ {
+			h ^= uint64(value[i])
+			h *= prime64
+		}
+		if valueIndex+1 < len(values) {
+			h ^= ','
+			h *= prime64
+		}
+	}
+	h ^= 0xff
+	h *= prime64
+	for i := 0; i < len(upstreamKey); i++ {
+		h ^= uint64(upstreamKey[i])
+		h *= prime64
+	}
+	return h
+}
+
 func reverseProxyFallbackPolicy(policy ReverseProxyLBPolicy) ReverseProxyLBPolicy {
 	if policy.fallback != nil {
 		return *policy.fallback