mirror of
https://github.com/Wei-Shaw/sub2api.git
synced 2026-10-31 08:25:35 +00:00
fix(upstream): cancel response attempt before closing body
This commit is contained in:
little-greenbean
committed by
xusang1011
xusang1011
parent
efe9aab1e4
commit
3fdd54ca12
@@ -216,7 +216,7 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i
|
||||
// 执行请求
|
||||
client := s.httpClientForUpstreamRequest(entry.client, req)
|
||||
client = httpClientWithGrokAccessDeniedFallback(client)
|
||||
resp, err := servertiming.Do(client, req)
|
||||
resp, err := doUpstreamRequest(client, req)
|
||||
if err != nil {
|
||||
s.recordOpenAIHTTP2Failure(profile, entry.protocolMode, entry.proxyKey, err)
|
||||
// 请求失败，立即减少计数
|
||||
@@ -226,9 +226,6 @@ func (s *httpUpstreamService) Do(req *http.Request, proxyURL string, accountID i
|
||||
}
|
||||
s.recordOpenAIHTTP2Success(profile, entry.protocolMode, entry.proxyKey)
|
||||
|
||||
// 如果上游返回了压缩内容，解压后再交给业务层
|
||||
decompressResponseBody(resp)
|
||||
|
||||
// 包装响应体，在关闭时自动减少计数并更新时间戳
|
||||
// 这确保了流式响应（如 SSE）在完全读取前不会被淘汰
|
||||
resp.Body = wrapTrackedBody(resp.Body, func() {
|
||||
@@ -280,7 +277,7 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
|
||||
|
||||
client := s.httpClientForUpstreamRequest(entry.client, req)
|
||||
client = httpClientWithGrokAccessDeniedFallback(client)
|
||||
resp, err := servertiming.Do(client, req)
|
||||
resp, err := doUpstreamRequest(client, req)
|
||||
if err != nil {
|
||||
atomic.AddInt64(&entry.inFlight, -1)
|
||||
atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
|
||||
@@ -288,8 +285,6 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
|
||||
return nil, err
|
||||
}
|
||||
|
||||
decompressResponseBody(resp)
|
||||
|
||||
resp.Body = wrapTrackedBody(resp.Body, func() {
|
||||
atomic.AddInt64(&entry.inFlight, -1)
|
||||
atomic.StoreInt64(&entry.lastUsed, time.Now().UnixNano())
|
||||
@@ -298,6 +293,56 @@ func (s *httpUpstreamService) DoWithTLS(req *http.Request, proxyURL string, acco
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
// doUpstreamRequest owns cancellation for one attempt, without cancelling the
|
||||
// caller's context (which may be detached for billing or reused for retries).
|
||||
func doUpstreamRequest(client *http.Client, req *http.Request) (*http.Response, error) {
|
||||
ctx, cancel := context.WithCancel(req.Context())
|
||||
resp, err := servertiming.Do(client, req.WithContext(ctx))
|
||||
if err != nil {
|
||||
cancel()
|
||||
return resp, err
|
||||
}
|
||||
decompressResponseBody(resp)
|
||||
resp.Body = &cancelOnCloseBody{ReadCloser: resp.Body, cancel: cancel}
|
||||
return resp, nil
|
||||
}
|
||||
|
||||
type cancelOnCloseBody struct {
|
||||
io.ReadCloser
|
||||
cancel context.CancelFunc
|
||||
once sync.Once
|
||||
err error
|
||||
readMu sync.Mutex
|
||||
closed bool
|
||||
}
|
||||
|
||||
func (b *cancelOnCloseBody) Read(p []byte) (int, error) {
|
||||
b.readMu.Lock()
|
||||
defer b.readMu.Unlock()
|
||||
if b.closed {
|
||||
return 0, http.ErrBodyReadAfterClose
|
||||
}
|
||||
return b.ReadCloser.Read(p)
|
||||
}
|
||||
|
||||
func (b *cancelOnCloseBody) Close() error {
|
||||
b.once.Do(func() {
|
||||
// Cancel before closing, including before closing a decompressor. In Go
|
||||
// 1.27 an early HTTP/1 close concurrent with Read can otherwise leave an
|
||||
// EOF waiter on a reused connection and stall subsequent responses.
|
||||
// Fully consumed responses have already released their transport request,
|
||||
// so cancelling here preserves normal keep-alive reuse.
|
||||
b.cancel()
|
||||
// Wait for an active read to observe cancellation before touching the
|
||||
// body or decompressor. Do not hold readMu while cancelling the request.
|
||||
b.readMu.Lock()
|
||||
defer b.readMu.Unlock()
|
||||
b.closed = true
|
||||
b.err = b.ReadCloser.Close()
|
||||
})
|
||||
return b.err
|
||||
}
|
||||
|
||||
// httpClientForUpstreamRequest 按请求上下文的标记派生客户端：禁用重定向，或对重定向的每一跳做主机校验。
|
||||
// 派生的克隆与缓存客户端共享 Transport；未打标记时原样返回。
|
||||
func (s *httpUpstreamService) httpClientForUpstreamRequest(client *http.Client, req *http.Request) *http.Client {
|
||||
|
||||
@@ -0,0 +1,148 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/config"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
type lifecycleUsageLogRepo struct {
|
||||
service.UsageLogRepository
|
||||
logs []*service.UsageLog
|
||||
}
|
||||
|
||||
func (r *lifecycleUsageLogRepo) Create(ctx context.Context, log *service.UsageLog) (bool, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return false, err
|
||||
}
|
||||
r.logs = append(r.logs, log)
|
||||
return true, nil
|
||||
}
|
||||
|
||||
type lifecycleBillingRepo struct {
|
||||
service.UsageBillingRepository
|
||||
commands []*service.UsageBillingCommand
|
||||
}
|
||||
|
||||
func (r *lifecycleBillingRepo) Apply(ctx context.Context, cmd *service.UsageBillingCommand) (*service.UsageBillingApplyResult, error) {
|
||||
if err := ctx.Err(); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
r.commands = append(r.commands, cmd)
|
||||
return &service.UsageBillingApplyResult{Applied: true}, nil
|
||||
}
|
||||
|
||||
type lifecycleDisconnectedWriter struct {
|
||||
gin.ResponseWriter
|
||||
disconnect context.CancelFunc
|
||||
}
|
||||
|
||||
func (w *lifecycleDisconnectedWriter) Write([]byte) (int, error) {
|
||||
w.disconnect()
|
||||
return 0, errors.New("client disconnected")
|
||||
}
|
||||
|
||||
// Exercise the real forwarding and RecordUsage services with the real HTTP
|
||||
// upstream. Only the persistence boundary is replaced: closing an attempt must
|
||||
// not cancel the detached stream before its final usage reaches billing.
|
||||
func TestHTTPUpstreamForwardDrainsUsageAfterClientDisconnect(t *testing.T) {
|
||||
gin.SetMode(gin.TestMode)
|
||||
clientCtx, disconnect := context.WithCancel(t.Context())
|
||||
defer disconnect()
|
||||
releaseUsage := make(chan struct{})
|
||||
var releaseOnce sync.Once
|
||||
release := func() { releaseOnce.Do(func() { close(releaseUsage) }) }
|
||||
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
w.Header().Set("X-Request-Id", "lifecycle-billing")
|
||||
_, _ = io.WriteString(w, "data: {\"type\":\"response.output_text.delta\",\"delta\":\"hello\"}\n\n")
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
select {
|
||||
case <-releaseUsage:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
_, _ = io.WriteString(w, "data: {\"type\":\"response.completed\",\"response\":{\"id\":\"resp_lifecycle\",\"status\":\"completed\",\"usage\":{\"input_tokens\":17,\"output_tokens\":8,\"total_tokens\":25}}}\n\n")
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
defer release()
|
||||
|
||||
cfg := &config.Config{}
|
||||
cfg.Default.RateMultiplier = 1
|
||||
cfg.Security.URLAllowlist.AllowPrivateHosts = true
|
||||
cfg.Security.URLAllowlist.AllowInsecureHTTP = true
|
||||
cfg.Gateway.StreamKeepaliveInterval = 1
|
||||
upstream := NewHTTPUpstream(cfg)
|
||||
usageRepo := &lifecycleUsageLogRepo{}
|
||||
billingRepo := &lifecycleBillingRepo{}
|
||||
svc := service.NewOpenAIGatewayService(
|
||||
nil, usageRepo, billingRepo, nil, nil, nil, nil, cfg, nil, nil,
|
||||
service.NewBillingService(cfg, nil), nil, nil, upstream,
|
||||
&service.DeferredService{}, nil, nil, nil, nil, nil, nil, nil,
|
||||
)
|
||||
account := &service.Account{
|
||||
ID: 1, Platform: service.PlatformOpenAI, Type: service.AccountTypeAPIKey,
|
||||
Concurrency: 1,
|
||||
Credentials: map[string]any{"base_url": srv.URL, "api_key": "test-key"},
|
||||
}
|
||||
body := []byte(`{"model":"gpt-5.1","stream":true,"input":"hello"}`)
|
||||
c, _ := gin.CreateTestContext(httptest.NewRecorder())
|
||||
c.Request = httptest.NewRequest(http.MethodPost, "/v1/responses", strings.NewReader(string(body))).WithContext(clientCtx)
|
||||
c.Writer = &lifecycleDisconnectedWriter{ResponseWriter: c.Writer, disconnect: disconnect}
|
||||
|
||||
type forwardResult struct {
|
||||
result *service.OpenAIForwardResult
|
||||
err error
|
||||
}
|
||||
done := make(chan forwardResult, 1)
|
||||
go func() {
|
||||
result, err := svc.Forward(clientCtx, c, account, body)
|
||||
done <- forwardResult{result, err}
|
||||
}()
|
||||
select {
|
||||
case <-clientCtx.Done():
|
||||
case got := <-done:
|
||||
t.Fatalf("forward returned before downstream disconnect: %v", got.err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("downstream did not receive the first event")
|
||||
}
|
||||
release()
|
||||
select {
|
||||
case got := <-done:
|
||||
require.NoError(t, got.err)
|
||||
require.NotNil(t, got.result)
|
||||
require.Equal(t, 17, got.result.Usage.InputTokens)
|
||||
require.Equal(t, 8, got.result.Usage.OutputTokens)
|
||||
// RecordUsage must still work with the cancelled downstream context.
|
||||
require.NoError(t, svc.RecordUsage(clientCtx, &service.OpenAIRecordUsageInput{
|
||||
Result: got.result, APIKey: &service.APIKey{ID: 2},
|
||||
User: &service.User{ID: 3}, Account: account,
|
||||
}))
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("forward did not finish draining usage")
|
||||
}
|
||||
require.Len(t, billingRepo.commands, 1)
|
||||
require.Len(t, usageRepo.logs, 1)
|
||||
require.Equal(t, 17, usageRepo.logs[0].InputTokens)
|
||||
require.Equal(t, 8, usageRepo.logs[0].OutputTokens)
|
||||
require.Greater(t, usageRepo.logs[0].ActualCost, 0.0)
|
||||
s, ok := upstream.(*httpUpstreamService)
|
||||
require.True(t, ok)
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
for _, entry := range s.clients {
|
||||
entry.client.CloseIdleConnections()
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,377 @@
|
||||
package repository
|
||||
|
||||
import (
|
||||
"compress/gzip"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net/http"
|
||||
"net/http/httptest"
|
||||
"net/http/httptrace"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/Wei-Shaw/sub2api/internal/pkg/tlsfingerprint"
|
||||
"github.com/Wei-Shaw/sub2api/internal/service"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// Exercise both public entry points. The TLS case uses the test server's trusted
|
||||
// transport so the test remains local and does not depend on system trust roots.
|
||||
func lifecycleClient(t *testing.T, srv *httptest.Server, withTLS bool) (*httpUpstreamService, func(*http.Request) (*http.Response, error)) {
|
||||
t.Helper()
|
||||
s, ok := NewHTTPUpstream(nil).(*httpUpstreamService)
|
||||
require.True(t, ok)
|
||||
profile := &tlsfingerprint.Profile{Name: "body-lifecycle-test"}
|
||||
var entry *upstreamClientEntry
|
||||
var err error
|
||||
if withTLS {
|
||||
entry, err = s.getClientEntryWithTLS("", 1, 1, profile, service.HTTPUpstreamProfileDefault, false, true)
|
||||
} else {
|
||||
entry, err = s.acquireClientWithProfile("", 1, 1, service.HTTPUpstreamProfileDefault)
|
||||
if err == nil {
|
||||
atomic.AddInt64(&entry.inFlight, -1)
|
||||
}
|
||||
}
|
||||
require.NoError(t, err)
|
||||
baseTransport, ok := srv.Client().Transport.(*http.Transport)
|
||||
require.True(t, ok)
|
||||
tr := baseTransport.Clone()
|
||||
tr.MaxConnsPerHost = 1
|
||||
tr.MaxIdleConnsPerHost = 1
|
||||
tr.IdleConnTimeout = 10 * time.Second
|
||||
entry.client = &http.Client{Transport: tr}
|
||||
t.Cleanup(tr.CloseIdleConnections)
|
||||
return s, func(req *http.Request) (*http.Response, error) {
|
||||
if withTLS {
|
||||
return s.DoWithTLS(req, "", 1, 1, profile)
|
||||
}
|
||||
return s.Do(req, "", 1, 1)
|
||||
}
|
||||
}
|
||||
|
||||
func requireNoUpstreamInFlight(t *testing.T, s *httpUpstreamService) {
|
||||
t.Helper()
|
||||
s.mu.RLock()
|
||||
defer s.mu.RUnlock()
|
||||
for _, entry := range s.clients {
|
||||
require.Zero(t, atomic.LoadInt64(&entry.inFlight))
|
||||
}
|
||||
}
|
||||
|
||||
type notifyReadCloser struct {
|
||||
io.ReadCloser
|
||||
started chan<- struct{}
|
||||
once sync.Once
|
||||
}
|
||||
|
||||
func (r *notifyReadCloser) Read(p []byte) (int, error) {
|
||||
r.once.Do(func() { close(r.started) })
|
||||
return r.ReadCloser.Read(p)
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamConcurrentEarlyCloseDoesNotPoisonNextResponse(t *testing.T) {
|
||||
for _, withTLS := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("tls=%v", withTLS), func(t *testing.T) {
|
||||
finishFirst := make(chan struct{})
|
||||
firstFinished := make(chan struct{})
|
||||
var requests atomic.Int32
|
||||
srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Type", "text/event-stream")
|
||||
_, _ = io.WriteString(w, "data: begin\n\n")
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
if requests.Add(1) == 1 {
|
||||
defer close(firstFinished)
|
||||
select {
|
||||
case <-finishFirst:
|
||||
case <-r.Context().Done():
|
||||
return
|
||||
}
|
||||
}
|
||||
_, _ = io.WriteString(w, "data: complete\n\n")
|
||||
}))
|
||||
if withTLS {
|
||||
srv.StartTLS()
|
||||
} else {
|
||||
srv.Start()
|
||||
}
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, withTLS)
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 10*time.Second)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
prefix := make([]byte, len("data: begin\n\n"))
|
||||
_, err = io.ReadFull(resp.Body, prefix)
|
||||
require.NoError(t, err)
|
||||
readerDone := make(chan struct{})
|
||||
readStarted := make(chan struct{})
|
||||
resp.Body = ¬ifyReadCloser{ReadCloser: resp.Body, started: readStarted}
|
||||
go func() { _, _ = io.Copy(io.Discard, resp.Body); close(readerDone) }()
|
||||
// Start Close only after the background reader has entered Read;
|
||||
// the server is still holding the response open at this point.
|
||||
<-readStarted
|
||||
closed := make(chan error, 1)
|
||||
go func() { closed <- resp.Body.Close() }()
|
||||
select {
|
||||
case err := <-closed:
|
||||
require.NoError(t, err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("Close blocked behind the active reader")
|
||||
}
|
||||
// Let the early-close drain finish only after Close has returned.
|
||||
// Go 1.27 can reuse that connection while its old reader awaits EOF.
|
||||
close(finishFirst)
|
||||
<-firstFinished
|
||||
for i := 0; i < 8; i++ {
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
done := make(chan error, 1)
|
||||
go func() { _, err := io.Copy(io.Discard, resp.Body); _ = resp.Body.Close(); done <- err }()
|
||||
select {
|
||||
case err := <-done:
|
||||
require.NoError(t, err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("response EOF waits for a subsequent request")
|
||||
}
|
||||
}
|
||||
select {
|
||||
case <-readerDone:
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("early-closed reader leaked")
|
||||
}
|
||||
require.NoError(t, req.Context().Err(), "closing an attempt must not cancel its caller or retries")
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamCompletedBodyPreservesKeepAlive(t *testing.T) {
|
||||
for _, encoding := range []string{"", "gzip"} {
|
||||
t.Run("encoding="+encoding, func(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if encoding == "gzip" {
|
||||
w.Header().Set("Content-Encoding", encoding)
|
||||
zw := gzip.NewWriter(w)
|
||||
_, _ = io.WriteString(zw, "complete response")
|
||||
_ = zw.Close()
|
||||
return
|
||||
}
|
||||
_, _ = io.WriteString(w, "complete response")
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
for i := 0; i < 20; i++ {
|
||||
var reused bool
|
||||
ctx := httptrace.WithClientTrace(t.Context(), &httptrace.ClientTrace{GotConn: func(info httptrace.GotConnInfo) { reused = info.Reused }})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("Accept-Encoding", encoding)
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "complete response", string(body))
|
||||
var wg sync.WaitGroup
|
||||
for j := 0; j < 4; j++ {
|
||||
wg.Add(1)
|
||||
go func() { defer wg.Done(); _ = resp.Body.Close() }()
|
||||
}
|
||||
wg.Wait()
|
||||
if i > 0 {
|
||||
require.True(t, reused, "fully consumed responses must keep their connection")
|
||||
}
|
||||
}
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamCancellationReleasesAttempt(t *testing.T) {
|
||||
for _, beforeHeaders := range []bool{false, true} {
|
||||
t.Run(fmt.Sprintf("before_headers=%v", beforeHeaders), func(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if !beforeHeaders {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
}
|
||||
<-r.Context().Done()
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
ctx, cancel := context.WithTimeout(t.Context(), 50*time.Millisecond)
|
||||
defer cancel()
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
resp, err := do(req)
|
||||
if beforeHeaders {
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
} else {
|
||||
require.NoError(t, err)
|
||||
_, err = io.ReadAll(resp.Body)
|
||||
require.ErrorIs(t, err, context.DeadlineExceeded)
|
||||
require.NoError(t, resp.Body.Close())
|
||||
}
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamCompressedEarlyCloseStopsReader(t *testing.T) {
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
w.Header().Set("Content-Encoding", "gzip")
|
||||
zw := gzip.NewWriter(w)
|
||||
_, _ = io.WriteString(zw, "data: begin\n\n")
|
||||
_ = zw.Flush()
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
<-r.Context().Done()
|
||||
_ = zw.Close()
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
req.Header.Set("Accept-Encoding", "gzip")
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
prefix := make([]byte, len("data: begin\n\n"))
|
||||
_, err = io.ReadFull(resp.Body, prefix)
|
||||
require.NoError(t, err)
|
||||
done := make(chan error, 1)
|
||||
readStarted := make(chan struct{})
|
||||
resp.Body = ¬ifyReadCloser{ReadCloser: resp.Body, started: readStarted}
|
||||
go func() { _, err := io.Copy(io.Discard, resp.Body); done <- err }()
|
||||
<-readStarted
|
||||
closed := make(chan error, 1)
|
||||
go func() { closed <- resp.Body.Close() }()
|
||||
select {
|
||||
case err := <-closed:
|
||||
require.NoError(t, err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("compressed body Close blocked")
|
||||
}
|
||||
select {
|
||||
case err := <-done:
|
||||
require.Error(t, err)
|
||||
case <-time.After(5 * time.Second):
|
||||
t.Fatal("compressed body reader leaked")
|
||||
}
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamHTTP2CloseDoesNotCancelOtherStream(t *testing.T) {
|
||||
srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
if r.URL.Path == "/held" {
|
||||
w.WriteHeader(http.StatusOK)
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
<-r.Context().Done()
|
||||
return
|
||||
}
|
||||
_, _ = io.WriteString(w, "complete response")
|
||||
}))
|
||||
srv.EnableHTTP2 = true
|
||||
srv.StartTLS()
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
req, err := http.NewRequestWithContext(t.Context(), http.MethodGet, srv.URL+"/held", nil)
|
||||
require.NoError(t, err)
|
||||
held, err := do(req)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, 2, held.ProtoMajor)
|
||||
for i := 0; i < 2; i++ {
|
||||
var reused bool
|
||||
ctx := httptrace.WithClientTrace(t.Context(), &httptrace.ClientTrace{GotConn: func(info httptrace.GotConnInfo) { reused = info.Reused }})
|
||||
req, err := http.NewRequestWithContext(ctx, http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
if i == 0 {
|
||||
require.NoError(t, held.Body.Close())
|
||||
}
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "complete response", string(body))
|
||||
require.NoError(t, resp.Body.Close())
|
||||
require.True(t, reused, "closing one HTTP/2 stream must retain the shared connection")
|
||||
}
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
}
|
||||
|
||||
func TestHTTPUpstreamDetachedRequestDrainsAfterClientDisconnect(t *testing.T) {
|
||||
disconnectCtx, disconnect := context.WithCancel(t.Context())
|
||||
defer disconnect()
|
||||
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
||||
_, _ = io.WriteString(w, "begin ")
|
||||
_ = http.NewResponseController(w).Flush()
|
||||
<-disconnectCtx.Done()
|
||||
_, _ = io.WriteString(w, "usage and completion")
|
||||
}))
|
||||
defer srv.Close()
|
||||
defer srv.CloseClientConnections()
|
||||
s, do := lifecycleClient(t, srv, false)
|
||||
req, err := http.NewRequestWithContext(context.WithoutCancel(disconnectCtx), http.MethodGet, srv.URL, nil)
|
||||
require.NoError(t, err)
|
||||
resp, err := do(req)
|
||||
require.NoError(t, err)
|
||||
disconnect()
|
||||
body, err := io.ReadAll(resp.Body)
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, "begin usage and completion", string(body))
|
||||
require.NoError(t, resp.Body.Close())
|
||||
requireNoUpstreamInFlight(t, s)
|
||||
}
|
||||
|
||||
// Keep the same benchmark runnable against the base revision: it measures a
|
||||
// complete public Do/read/Close lifecycle without network latency hiding the
|
||||
// per-attempt context and per-Read synchronization costs.
|
||||
func BenchmarkHTTPUpstreamBodyLifecycle(b *testing.B) {
|
||||
for _, size := range []int{256, 64 * 1024} {
|
||||
b.Run(fmt.Sprintf("bytes=%d", size), func(b *testing.B) {
|
||||
s, ok := NewHTTPUpstream(nil).(*httpUpstreamService)
|
||||
require.True(b, ok)
|
||||
entry, err := s.getOrCreateClient("", 1, 1)
|
||||
require.NoError(b, err)
|
||||
payload := strings.Repeat("x", size)
|
||||
entry.client = &http.Client{Transport: roundTripFunc(func(*http.Request) (*http.Response, error) {
|
||||
return &http.Response{
|
||||
StatusCode: http.StatusOK, Header: make(http.Header),
|
||||
Body: io.NopCloser(strings.NewReader(payload)),
|
||||
}, nil
|
||||
})}
|
||||
req, err := http.NewRequestWithContext(b.Context(), http.MethodGet, "https://example.com", nil)
|
||||
require.NoError(b, err)
|
||||
buf := make([]byte, 256)
|
||||
b.ReportAllocs()
|
||||
b.SetBytes(int64(size))
|
||||
b.ResetTimer()
|
||||
for b.Loop() {
|
||||
resp, err := s.Do(req, "", 1, 1)
|
||||
if err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
for {
|
||||
_, err = resp.Body.Read(buf)
|
||||
if err != nil {
|
||||
break
|
||||
}
|
||||
}
|
||||
if err != io.EOF {
|
||||
b.Fatal(err)
|
||||
}
|
||||
if err := resp.Body.Close(); err != nil {
|
||||
b.Fatal(err)
|
||||
}
|
||||
}
|
||||
})
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user