wjqserver/touka
1
0
Fork 0
mirror of https://github.com/infinite-iroha/touka.git synced 2026-06-13 15:47:38 +08:00
Code Issues Projects Releases Packages Wiki Activity Actions 13

Merge pull request #86 from infinite-iroha/perf/go126-memory-pass

Browse source 
Perf/go126 memory pass
This commit is contained in:
WJQSERVER 2026-04-21 16:29:12 +08:00 • committed by GitHub
commit 0f7cf23abb
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 531 additions and 29 deletions
Download patch file Download diff file Expand all files Collapse all files

33
context.go
View file

@ -128,6 +128,19 @@ func (c *Context) reset(w http.ResponseWriter, req *http.Request) {
}
}
func (c *Context) writeResponseBody(data []byte, contextMsg string) {
if len(data) == 0 {
return
}
if _, err := c.Writer.Write(data); err != nil {
wrapped := fmt.Errorf("%s: %w", contextMsg, err)
c.AddError(wrapped)
if c != nil && c.engine != nil && c.engine.LogReco != nil {
c.engine.LogReco.Errorf("%s: %v", contextMsg, err)
}
}
}
// Next 在处理链中执行下一个处理函数
// 这是中间件模式的核心，允许请求依次经过多个处理函数
func (c *Context) Next() {
@ -344,20 +357,20 @@ func (c *Context) Param(key string) string {
func (c *Context) Raw(code int, contentType string, data []byte) {
c.Writer.Header().Set("Content-Type", contentType)
c.Writer.WriteHeader(code)
c.Writer.Write(data)
c.writeResponseBody(data, "failed to write raw response")
}
// String 向响应写入格式化的字符串
func (c *Context) String(code int, format string, values ...any) {
c.Writer.WriteHeader(code)
c.Writer.Write(fmt.Appendf(nil, format, values...))
c.writeResponseBody(fmt.Appendf(nil, format, values...), "failed to write string response")
}
// Text 向响应写入无需格式化的string
func (c *Context) Text(code int, text string) {
c.Writer.Header().Set("Content-Type", "text/plain; charset=utf-8")
c.Writer.WriteHeader(code)
c.Writer.Write([]byte(text))
c.writeResponseBody([]byte(text), "failed to write text response")
}
// FileText
@ -495,7 +508,7 @@ func (c *Context) JSONBuf(code int, obj any) {
c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8")
c.Writer.WriteHeader(code)
c.Writer.Write(buf.Bytes())
c.writeResponseBody(buf.Bytes(), "failed to write buffered JSON response")
}
// GOB 向响应写入GOB数据
@ -524,7 +537,7 @@ func (c *Context) GOBBuf(code int, obj any) {
}
c.Writer.Header().Set("Content-Type", "application/octet-stream")
c.Writer.WriteHeader(code)
c.Writer.Write(buf.Bytes())
c.writeResponseBody(buf.Bytes(), "failed to write buffered GOB response")
}
// WANF向响应写入WANF数据
@ -553,7 +566,7 @@ func (c *Context) WANFBuf(code int, obj any) {
}
c.Writer.Header().Set("Content-Type", "application/vnd.wjqserver.wanf; charset=utf-8")
c.Writer.WriteHeader(code)
c.Writer.Write(buf.Bytes())
c.writeResponseBody(buf.Bytes(), "failed to write buffered WANF response")
}
// HTML 渲染 HTML 模板
@ -577,7 +590,7 @@ func (c *Context) HTML(code int, name string, obj any) {
// 可以扩展支持其他渲染器接口
}
// 默认简单输出，用于未配置 HTMLRender 的情况
c.Writer.Write(fmt.Appendf(nil, "<!-- HTML rendered for %s -->\n<pre>%v</pre>", name, obj))
c.writeResponseBody(fmt.Appendf(nil, "<!-- HTML rendered for %s -->\n<pre>%v</pre>", name, obj), "failed to write HTML response")
}
// HTMLBuf 先将 HTML 模板渲染到 buffer, 成功后再写入状态码和响应体.
@ -602,7 +615,7 @@ func (c *Context) HTMLBuf(code int, name string, obj any) {
// 渲染成功，写入响应
c.Writer.Header().Set("Content-Type", "text/html; charset=utf-8")
c.Writer.WriteHeader(code)
c.Writer.Write(buf.Bytes())
c.writeResponseBody(buf.Bytes(), "failed to write buffered HTML response")
return
}
@ -938,7 +951,7 @@ func (c *Context) GetReqBodyFull() ([]byte, error) {
}
}()
data, err := iox.ReadAll(body)
data, err := io.ReadAll(body)
if err != nil {
c.AddError(fmt.Errorf("failed to read request body: %w", err))
return nil, fmt.Errorf("failed to read request body: %w", err)
@ -959,7 +972,7 @@ func (c *Context) GetReqBodyBuffer() (*bytes.Buffer, error) {
}
}()
data, err := iox.ReadAll(body)
data, err := io.ReadAll(body)
if err != nil {
c.AddError(fmt.Errorf("failed to read request body: %w", err))
return nil, fmt.Errorf("failed to read request body: %w", err)

37
engine.go
View file

@ -19,6 +19,7 @@ import (
"github.com/WJQSERVER-STUDIO/httpc"
"github.com/fenthope/reco"
"github.com/go-json-experiment/json"
)
// Last 返回链中的最后一个处理函数
@ -132,6 +133,32 @@ type defaultErrorResponse struct {
Error string `json:"error"`
}
var defaultNotFoundBody = mustMarshalDefaultErrorBody(http.StatusNotFound, errNotFound.Error())
var defaultMethodNotAllowedBody = mustMarshalDefaultErrorBody(http.StatusMethodNotAllowed, errMethodNotAllowed.Error())
func mustMarshalDefaultErrorBody(code int, errMsg string) []byte {
body, err := json.Marshal(defaultErrorResponse{
Code: code,
Message: http.StatusText(code),
Error: errMsg,
})
if err != nil {
panic(err)
}
return body
}
func writeDefaultErrorJSON(c *Context, code int, body []byte) {
if c == nil || c.Writer == nil {
return
}
c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8")
c.Writer.WriteHeader(code)
c.writeResponseBody(body, "failed to write default error response")
c.Writer.Flush()
c.Abort()
}
var methodNotAllowedHandler HandlerFunc = func(c *Context) {
httpMethod := c.Request.Method
requestPath := routeLookupPath(c.Request)
@ -191,6 +218,16 @@ func defaultErrorHandle(c *Context, code int, err error) { // 检查客户端是
if c.Writer.Written() {
return
}
if len(c.Errors) == 0 {
switch {
case code == http.StatusNotFound && errors.Is(err, errNotFound):
writeDefaultErrorJSON(c, code, defaultNotFoundBody)
return
case code == http.StatusMethodNotAllowed && errors.Is(err, errMethodNotAllowed):
writeDefaultErrorJSON(c, code, defaultMethodNotAllowedBody)
return
}
}
// 输出json 状态码与状态码对应描述
var errMsg string
if err != nil {

165
engine_test.go
View file

@ -1,11 +1,66 @@
package touka
import (
"bufio"
"encoding/json"
"errors"
"html/template"
"net"
"net/http"
"testing"
)
type failingResponseWriter struct {
header http.Header
status int
err error
}
func (w *failingResponseWriter) Header() http.Header {
if w.header == nil {
w.header = make(http.Header)
}
return w.header
}
func (w *failingResponseWriter) WriteHeader(statusCode int) {
if w.status == 0 {
w.status = statusCode
}
}
func (w *failingResponseWriter) Write(p []byte) (int, error) {
if w.status == 0 {
w.status = http.StatusOK
}
if w.err != nil {
return 0, w.err
}
return len(p), nil
}
func (w *failingResponseWriter) Flush() {}
func (w *failingResponseWriter) Status() int {
return w.status
}
func (w *failingResponseWriter) Size() int {
return 0
}
func (w *failingResponseWriter) Written() bool {
return w.status != 0
}
func (w *failingResponseWriter) IsHijacked() bool {
return false
}
func (w *failingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
return nil, nil, http.ErrNotSupported
}
func TestHandleRequestRedirectFixedPath(t *testing.T) {
engine := New()
engine.GET("/api/v1/users/:id/settings", func(c *Context) {
@ -139,3 +194,113 @@ func TestDefaultErrorHandleJSONShape(t *testing.T) {
t.Fatalf("unexpected error payload: %+v", body)
}
}
func TestDefaultMethodNotAllowedJSONShape(t *testing.T) {
engine := New()
engine.GET("/users", func(c *Context) {
c.Status(http.StatusNoContent)
})
rr := PerformRequest(engine, http.MethodDelete, "/users", nil, nil)
if rr.Code != http.StatusMethodNotAllowed {
t.Fatalf("expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code)
}
var body struct {
Code int `json:"code"`
Message string `json:"message"`
Error string `json:"error"`
}
if err := json.Unmarshal(rr.Body.Bytes(), &body); err != nil {
t.Fatalf("expected JSON error body, got %q: %v", rr.Body.String(), err)
}
if body.Code != http.StatusMethodNotAllowed || body.Message != http.StatusText(http.StatusMethodNotAllowed) || body.Error != "method not allowed" {
t.Fatalf("unexpected error payload: %+v", body)
}
}
func TestCustomErrorHandlerStillOverridesDefaultFastPath(t *testing.T) {
engine := New()
engine.SetErrorHandler(func(c *Context, code int, err error) {
c.Writer.Header().Set("X-Custom-Error", "1")
c.String(code, "custom:%v", err)
})
engine.GET("/users", func(c *Context) {
c.Status(http.StatusNoContent)
})
rr := PerformRequest(engine, http.MethodDelete, "/users", nil, nil)
if rr.Code != http.StatusMethodNotAllowed {
t.Fatalf("expected status %d, got %d", http.StatusMethodNotAllowed, rr.Code)
}
if got := rr.Header().Get("X-Custom-Error"); got != "1" {
t.Fatalf("expected custom error header, got %q", got)
}
if rr.Body.String() != "custom:method not allowed" {
t.Fatalf("expected custom error body, got %q", rr.Body.String())
}
}
func TestResponseHelpersCaptureWriteErrors(t *testing.T) {
testCases := []struct {
name string
run func(*Context)
}{
{name: "Raw", run: func(c *Context) { c.Raw(http.StatusOK, "application/octet-stream", []byte("payload")) }},
{name: "String", run: func(c *Context) { c.String(http.StatusOK, "value=%d", 1) }},
{name: "Text", run: func(c *Context) { c.Text(http.StatusOK, "payload") }},
{name: "JSONBuf", run: func(c *Context) { c.JSONBuf(http.StatusOK, map[string]string{"a": "b"}) }},
{name: "GOBBuf", run: func(c *Context) { c.GOBBuf(http.StatusOK, struct{ A string }{A: "b"}) }},
{name: "WANFBuf", run: func(c *Context) { c.WANFBuf(http.StatusOK, map[string]string{"a": "b"}) }},
{name: "HTMLFallback", run: func(c *Context) { c.HTML(http.StatusOK, "page", map[string]string{"a": "b"}) }},
{name: "HTMLBuf", run: func(c *Context) {
c.engine.HTMLRender = template.Must(template.New("page").Parse(`{{.a}}`))
c.HTMLBuf(http.StatusOK, "page", map[string]string{"a": "b"})
}},
}
for _, tc := range testCases {
t.Run(tc.name, func(t *testing.T) {
writerErr := errors.New("write failed")
w := &failingResponseWriter{err: writerErr}
c, _ := CreateTestContext(w)
tc.run(c)
if got := len(c.Errors); got != 1 {
t.Fatalf("expected exactly one captured error, got %d", got)
}
if !errors.Is(c.Errors[len(c.Errors)-1], writerErr) {
t.Fatalf("expected captured error to wrap write failure, got %v", c.Errors[len(c.Errors)-1])
}
})
}
}
func TestDefaultErrorFastPathCapturesWriteErrors(t *testing.T) {
writerErr := errors.New("write failed")
w := &failingResponseWriter{err: writerErr}
engine := New()
c, _ := CreateTestContext(w)
c.engine = engine
req, err := http.NewRequest(http.MethodGet, "/missing", nil)
if err != nil {
t.Fatalf("failed to build request: %v", err)
}
c.reset(w, req)
defaultErrorHandle(c, http.StatusNotFound, errNotFound)
if len(c.Errors) == 0 {
t.Fatal("expected write error to be captured")
}
if !errors.Is(c.Errors[len(c.Errors)-1], writerErr) {
t.Fatalf("expected captured error to wrap write failure, got %v", c.Errors[len(c.Errors)-1])
}
if c.Writer.Status() != http.StatusNotFound {
t.Fatalf("expected status %d, got %d", http.StatusNotFound, c.Writer.Status())
}
if !c.IsAborted() {
t.Fatal("expected fast path to abort context")
}
}

2
go.mod
View file

@ -3,7 +3,7 @@ module github.com/infinite-iroha/touka
go 1.26
require (
github.com/WJQSERVER-STUDIO/go-utils/iox v0.0.2
github.com/WJQSERVER-STUDIO/go-utils/iox v0.0.3
github.com/WJQSERVER-STUDIO/httpc v0.9.0
github.com/WJQSERVER/wanf v0.0.8
github.com/fenthope/reco v0.0.5

2
go.sum
View file

@ -1,5 +1,7 @@
github.com/WJQSERVER-STUDIO/go-utils/iox v0.0.2 h1:AiIHXP21LpK7pFfqUlUstgQEWzjbekZgxOuvVwiMfyM=
github.com/WJQSERVER-STUDIO/go-utils/iox v0.0.2/go.mod h1:mCLqYU32bTmEE6dpj37MKKiZgz70Jh/xyK9vVbq6pok=
github.com/WJQSERVER-STUDIO/go-utils/iox v0.0.3 h1:Hc1O6D50U3URkdSzfQ/SgeUU750wUBCYhefdvAbE2Ck=
github.com/WJQSERVER-STUDIO/go-utils/iox v0.0.3/go.mod h1:nFQzepAwwdj5Hp5U+X19l4FVvsaOSBTW41BzfI/CkMA=
github.com/WJQSERVER-STUDIO/httpc v0.9.0 h1:MpXcQQqukrSLHH/2tTfnXrhqD6nEDHB/gbzehXaS8o4=
github.com/WJQSERVER-STUDIO/httpc v0.9.0/go.mod h1:filzryrl4eAtFVyl4oVHcJqx1SpNFbrCn+ddQPLlCSg=
github.com/WJQSERVER/wanf v0.0.8 h1:1Ri9d7nKhu22hGxP8O9B9rXnYym6DYGKgi6WRVx3VF8=

150
iox_benchmark_test.go Normal file
View file

@ -0,0 +1,150 @@
package touka
import (
"bytes"
"io"
"testing"
"github.com/WJQSERVER-STUDIO/go-utils/iox"
)
type benchmarkResetReader struct {
data []byte
off int
}
func (r *benchmarkResetReader) Read(p []byte) (int, error) {
if r.off >= len(r.data) {
return 0, io.EOF
}
n := copy(p, r.data[r.off:])
r.off += n
return n, nil
}
func (r *benchmarkResetReader) Reset() {
r.off = 0
}
type benchmarkDiscardWriter struct{}
func (benchmarkDiscardWriter) Write(p []byte) (int, error) {
return len(p), nil
}
var benchmarkIOXResult int64
var benchmarkIOXBytes []byte
func BenchmarkIOXCopyComparison(b *testing.B) {
payload := bytes.Repeat([]byte("0123456789abcdef"), 4096)
b.Run("io.Copy", func(b *testing.B) {
r := &benchmarkResetReader{data: payload}
w := benchmarkDiscardWriter{}
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
r.Reset()
n, err := io.Copy(w, r)
if err != nil {
b.Fatalf("io.Copy failed: %v", err)
}
benchmarkIOXResult = n
}
})
b.Run("iox.Copy", func(b *testing.B) {
r := &benchmarkResetReader{data: payload}
w := benchmarkDiscardWriter{}
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
r.Reset()
n, err := iox.Copy(w, r)
if err != nil {
b.Fatalf("iox.Copy failed: %v", err)
}
benchmarkIOXResult = n
}
})
}
func BenchmarkIOXCopyBufferComparison(b *testing.B) {
payload := bytes.Repeat([]byte("0123456789abcdef"), 4096)
b.Run("io.CopyBuffer", func(b *testing.B) {
r := &benchmarkResetReader{data: payload}
w := benchmarkDiscardWriter{}
buf := make([]byte, 32*1024)
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
r.Reset()
n, err := io.CopyBuffer(w, r, buf)
if err != nil {
b.Fatalf("io.CopyBuffer failed: %v", err)
}
benchmarkIOXResult = n
}
})
b.Run("iox.CopyBuffer", func(b *testing.B) {
r := &benchmarkResetReader{data: payload}
w := benchmarkDiscardWriter{}
buf := make([]byte, 32*1024)
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
r.Reset()
n, err := iox.CopyBuffer(w, r, buf)
if err != nil {
b.Fatalf("iox.CopyBuffer failed: %v", err)
}
benchmarkIOXResult = n
}
})
}
func BenchmarkIOXReadAllComparison(b *testing.B) {
payload := bytes.Repeat([]byte("0123456789abcdef"), 4096)
b.Run("io.ReadAll", func(b *testing.B) {
r := &benchmarkResetReader{data: payload}
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
r.Reset()
data, err := io.ReadAll(r)
if err != nil {
b.Fatalf("io.ReadAll failed: %v", err)
}
benchmarkIOXBytes = data
}
})
b.Run("iox.ReadAll", func(b *testing.B) {
r := &benchmarkResetReader{data: payload}
b.ReportAllocs()
b.ResetTimer()
for i := 0; i < b.N; i++ {
r.Reset()
data, err := io.ReadAll(r)
if err != nil {
b.Fatalf("iox.ReadAll failed: %v", err)
}
benchmarkIOXBytes = data
}
})
}

12
reverseproxy.go
View file

@ -1041,7 +1041,7 @@ func (p *reverseProxyHandler) handleBridgedExtendedConnectResponse(c *Context, r
go copyer.copyFromBackend(errc)
var firstErr error
for i := 0; i < 2; i++ {
for range 2 {
err := <-errc
if reverseProxyIsBenignTunnelError(err) {
continue
@ -1123,7 +1123,7 @@ func (p *reverseProxyHandler) handleExtendedConnectResponse(c *Context, req *htt
}()
var firstErr error
for i := 0; i < 2; i++ {
for range 2 {
err := <-errc
if reverseProxyIsBenignTunnelError(err) {
continue
@ -1587,8 +1587,8 @@ func reverseProxyViaProtocol(major, minor int, raw string) string {
if major > 0 {
return strconv.Itoa(major) + "." + strconv.Itoa(minor)
}
if strings.HasPrefix(raw, "HTTP/") {
return strings.TrimPrefix(raw, "HTTP/")
if after, ok := strings.CutPrefix(raw, "HTTP/"); ok {
return after
}
return raw
}
@ -1702,7 +1702,7 @@ var reverseProxyHopHeaders = []string{
func removeHopByHopHeaders(header http.Header) {
for _, connectionValue := range header["Connection"] {
for _, token := range strings.Split(connectionValue, ",") {
for token := range strings.SplitSeq(connectionValue, ",") {
trimmed := textproto.TrimString(token)
if trimmed != "" {
header.Del(trimmed)
@ -1726,7 +1726,7 @@ func headerValuesContainToken(values []string, token string) bool {
return false
}
for _, value := range values {
for _, part := range strings.Split(value, ",") {
for part := range strings.SplitSeq(value, ",") {
if strings.EqualFold(textproto.TrimString(part), token) {
return true
}

94
reverseproxy_benchmark_test.go
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()
@ -207,7 +235,7 @@ func (r *recordingReader) Read(p []byte) (int, error) {
if n == 0 {
return 0, errors.New("reader received zero-length buffer")
}
for i := 0; i < n; i++ {
for i := range n {
p[i] = 'x'
}
r.left -= n
@ -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)

54
reverseproxy_lb.go
View file

@ -10,6 +10,7 @@ import (
"net/http"
"net/textproto"
"net/url"
"slices"
"strings"
"sync"
"sync/atomic"
@ -199,7 +200,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 +278,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 +316,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
@ -360,10 +405,5 @@ func reverseProxyStatusIsUnhealthy(config ReverseProxyPassiveHealthConfig, statu
if status <= 0 {
return false
}
for _, unhealthyStatus := range config.UnhealthyStatus {
if status == unhealthyStatus {
return true
}
}
return false
return slices.Contains(config.UnhealthyStatus, status)
}

8
reverseproxy_test.go
View file

@ -1866,7 +1866,9 @@ func TestReverseProxyHTTP2ExtendedConnectForcesHTTP1ToTLSUpstream(t *testing.T)
if message != "echo:ping\n" {
t.Fatalf("unexpected tunneled response body: %q", message)
}
_ = pw.Close()
if err := pw.Close(); err != nil {
t.Fatalf("close tunneled request body: %v", err)
}
select {
case err := <-errCh:
@ -2215,7 +2217,9 @@ func TestReverseProxyHTTP2ExtendedConnectCancelDoesNotTriggerProxyError(t *testi
time.Sleep(50 * time.Millisecond)
cancel()
_ = pw.CloseWithError(context.Canceled)
if err := pw.CloseWithError(context.Canceled); err != nil {
t.Fatalf("close request body with cancellation: %v", err)
}
select {
case <-writeErrCh:
case <-time.After(2 * time.Second):

3
serve_test.go
View file

@ -182,8 +182,7 @@ func TestBuildRedirectServerRejectsHTTPSAddrWithoutPort(t *testing.T) {
func TestValidateRunConfigRejectsShutdownContextWithoutGraceful(t *testing.T) {
cfg := defaultRunConfig()
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
ctx := t.Context()
if err := WithShutdownContext(ctx).apply(&cfg); err != nil {
t.Fatalf("apply shutdown context option: %v", err)
}