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"
"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)