perf: avoid header policy join allocations

This commit is contained in:
wjqserver 2026-04-10 21:55:21 +08:00
parent 7c37d4c38c
commit 02861b5537
2 changed files with 137 additions and 1 deletions

View file

@ -7,6 +7,7 @@ import (
"io" "io"
"net" "net"
"net/http" "net/http"
"strings"
"testing" "testing"
"time" "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) { func TestReverseProxyCopyResponseWithoutBufferPool(t *testing.T) {
proxy := newReverseProxyHandler(ReverseProxyConfig{}) proxy := newReverseProxyHandler(ReverseProxyConfig{})
dst := newBenchmarkResponseWriter() 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) var _ io.Writer = (*benchmarkResponseWriter)(nil)

View file

@ -199,7 +199,7 @@ func (p *reverseProxyHandler) selectUpstreamWithPolicy(c *Context, candidates []
case reverseProxyLBPolicyHeader: case reverseProxyLBPolicyHeader:
if c.Request != nil && c.Request.Header != nil { if c.Request != nil && c.Request.Header != nil {
if values, ok := c.Request.Header[policy.key]; ok { 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)) return p.selectUpstreamWithPolicy(c, candidates, reverseProxyFallbackPolicy(policy))
@ -277,6 +277,25 @@ func reverseProxySelectHRW(candidates []*reverseProxyUpstream, key string) *reve
return selected 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 { func reverseProxyHRWScore(key, upstreamKey string) uint64 {
const ( const (
offset64 = 14695981039346656037 offset64 = 14695981039346656037
@ -296,6 +315,31 @@ func reverseProxyHRWScore(key, upstreamKey string) uint64 {
return h 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 { func reverseProxyFallbackPolicy(policy ReverseProxyLBPolicy) ReverseProxyLBPolicy {
if policy.fallback != nil { if policy.fallback != nil {
return *policy.fallback return *policy.fallback