Skip to content

feat(upstream): add HTTP/2 PING keepalive for long LLM generation streams #6281

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wuji-labs wants to merge 1 commit into
base: main
Choose a base branch
from
+41 −11
Open

feat(upstream): add HTTP/2 PING keepalive for long LLM generation streams #6281

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
24 changes: 14 additions & 10 deletions backend/internal/repository/http_upstream.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -76,13 +76,13 @@ const (
defaultOpenAIHTTP2FallbackErrorThreshold = 2
defaultOpenAIHTTP2FallbackWindow = 60 * time.Second
defaultOpenAIHTTP2FallbackTTL = 10 * time.Minute
// OpenAI HTTP/2 连接健康探测：Codex 上游改走 HTTP/2 后，池化连接被代理/NAT
// 长流 HTTP/2 连接健康探测：池化连接被代理/NAT
// 静默掐断会成为“死连接”（两端都以为存活），请求落上去会挂到 TCP 重传超时
// （分钟级）。Go 的 http2.Transport 默认 ReadIdleTimeout=0（不发健康 PING），
// 无法检测。启用主动 PING 探测：连接空闲 ReadIdleTimeout 后发 PING，PingTimeout
// 内无响应即判定死连接并关闭，从源头避免请求挂在死连接上。
openAIHTTP2ReadIdleTimeout = 15 * time.Second
openAIHTTP2PingTimeout = 15 * time.Second
longStreamHTTP2ReadIdleTimeout = 10 * time.Second
longStreamHTTP2PingTimeout = 5 * time.Second

// The Grok CLI proxy rejects requests that do not identify a supported
// client version. Host/env/version pins live in package xai so service,
Expand All @@ -96,6 +96,7 @@ const (

const (
upstreamProtocolModeDefault = "default"
upstreamProtocolModeLongStreamH2 = "long_stream_h2"
upstreamProtocolModeOpenAIH1 = "openai_h1"
upstreamProtocolModeOpenAIH2 = "openai_h2"
upstreamProtocolModeOpenAIH1Fallback = "openai_h1_fallback"
Expand Down Expand Up @@ -128,7 +129,7 @@ type upstreamClientEntry struct {
client *http.Client // HTTP 客户端实例
proxyKey string // 代理标识（用于检测代理变更）
poolKey string // 连接池配置标识（用于检测配置变更）
protocolMode string // 协议模式（default/openai_h1/openai_h2/openai_h1_fallback）
protocolMode string // 协议模式（default/long_stream_h2/openai_h1/openai_h2/openai_h1_fallback）
lastUsed int64 // 最后使用时间戳（纳秒），用于 LRU 淘汰
inFlight int64 // 当前进行中的请求数，>0 时不可淘汰
}
Expand Down Expand Up @@ -992,6 +993,9 @@ func (s *httpUpstreamService) resolveOpenAIHTTP2Settings() openAIHTTP2Settings {
}

func (s *httpUpstreamService) resolveProtocolMode(profile service.HTTPUpstreamProfile, proxyKey string, parsedProxy *url.URL) string {
if profile == service.HTTPUpstreamProfileLongStream {
return upstreamProtocolModeLongStreamH2
}
if profile == service.HTTPUpstreamProfileGrok {
return upstreamProtocolModeGrok
}
Expand Down Expand Up @@ -1317,11 +1321,11 @@ func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMo
ResponseHeaderTimeout: settings.responseHeaderTimeout,
}
switch protocolMode {
case upstreamProtocolModeOpenAIH2:
case upstreamProtocolModeLongStreamH2, upstreamProtocolModeOpenAIH2:
transport.ForceAttemptHTTP2 = true
// 显式配置 http2 并启用 PING 健康探测，剔除代理/NAT 静默掐断的死连接，
// 避免请求挂在死连接上直到 TCP 重传超时（分钟级）。
if _, err := enableOpenAIHTTP2KeepAlive(transport); err != nil {
if _, err := enableHTTP2KeepAlive(transport); err != nil {
return nil, err
}
case upstreamProtocolModeOpenAIH1:
Expand All @@ -1338,18 +1342,18 @@ func buildUpstreamTransport(settings poolSettings, proxyURL *url.URL, protocolMo
return transport, nil
}

// enableOpenAIHTTP2KeepAlive 在 http.Transport 上显式配置 HTTP/2 并启用连接健康探测。
// enableHTTP2KeepAlive 在 http.Transport 上显式配置 HTTP/2 并启用连接健康探测。
// Go 默认惰性配置 http2 且 ReadIdleTimeout=0（不发健康 PING），无法检测被代理/NAT
// 静默掐断的死连接。此处主动设置 ReadIdleTimeout/PingTimeout，让死连接被提前 PING
// 出并关闭，请求得以重建连接而非挂到 TCP 重传超时。返回底层 *http2.Transport 便于测试。
func enableOpenAIHTTP2KeepAlive(transport *http.Transport) (*http2.Transport, error) {
func enableHTTP2KeepAlive(transport *http.Transport) (*http2.Transport, error) {
h2, err := http2.ConfigureTransports(transport)
if err != nil {
return nil, err
}
if h2 != nil {
h2.ReadIdleTimeout = openAIHTTP2ReadIdleTimeout
h2.PingTimeout = openAIHTTP2PingTimeout
h2.ReadIdleTimeout = longStreamHTTP2ReadIdleTimeout
h2.PingTimeout = longStreamHTTP2PingTimeout
}
return h2, nil
}
Expand Down
18 changes: 18 additions & 0 deletions backend/internal/repository/http_upstream_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -629,6 +629,24 @@ func (s *HTTPUpstreamSuite) TestOpenAIProfileDefaultsToHTTP2AndNoHeaderTimeout()
require.Equal(s.T(), upstreamProtocolModeOpenAIH2, entry.protocolMode)
}

func (s *HTTPUpstreamSuite) TestLongStreamProfileUsesSharedHTTP2KeepAlive() {
s.cfg.Gateway = config.GatewayConfig{
ResponseHeaderTimeout: 600,
OpenAIHTTP2: config.GatewayOpenAIHTTP2Config{
Enabled: false,
},
}
svc := s.newService()
entry, err := svc.getClientEntry("", 1, 1, service.HTTPUpstreamProfileLongStream, false, false)
require.NoError(s.T(), err)
transport, ok := entry.client.Transport.(*http.Transport)
require.True(s.T(), ok, "expected *http.Transport")
require.Equal(s.T(), 600*time.Second, transport.ResponseHeaderTimeout, "long-stream profile should retain the generic header timeout")
require.True(s.T(), transport.ForceAttemptHTTP2, "long-stream profile must enable HTTP/2 independently of OpenAI settings")
require.NotNil(s.T(), transport.TLSNextProto["h2"], "long-stream profile must install HTTP/2 PING health checks")
require.Equal(s.T(), upstreamProtocolModeLongStreamH2, entry.protocolMode)
}

func (s *HTTPUpstreamSuite) TestOpenAIProfileCustomHeaderTimeout() {
s.cfg.Gateway = config.GatewayConfig{
ResponseHeaderTimeout: 600,
Expand Down
3 changes: 2 additions & 1 deletion backend/internal/service/http_upstream_profile.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ const (
HTTPUpstreamProfileDefault HTTPUpstreamProfile = ""
HTTPUpstreamProfileOpenAI HTTPUpstreamProfile = "openai"
HTTPUpstreamProfileGrok HTTPUpstreamProfile = "grok"
HTTPUpstreamProfileLongStream HTTPUpstreamProfile = "long_stream"
)

type httpUpstreamProfileContextKey struct{}
Expand All @@ -36,7 +37,7 @@ func HTTPUpstreamProfileFromContext(ctx context.Context) HTTPUpstreamProfile {
return HTTPUpstreamProfileDefault
}
switch profile {
case HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileGrok:
case HTTPUpstreamProfileOpenAI, HTTPUpstreamProfileGrok, HTTPUpstreamProfileLongStream:
return profile
default:
return HTTPUpstreamProfileDefault
Expand Down
7 changes: 7 additions & 0 deletions backend/internal/service/http_upstream_profile_test.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,13 @@ func TestWithHTTPUpstreamProfile_OpenAI(t *testing.T) {
}
}

func TestWithHTTPUpstreamProfile_LongStream(t *testing.T) {
ctx := WithHTTPUpstreamProfile(context.TODO(), HTTPUpstreamProfileLongStream)
if profile := HTTPUpstreamProfileFromContext(ctx); profile != HTTPUpstreamProfileLongStream {
t.Fatalf("expected profile %q, got %q", HTTPUpstreamProfileLongStream, profile)
}
}

func TestWithHTTPUpstreamRedirectsDisabled(t *testing.T) {
//nolint:staticcheck // Exercises the defensive nil-context fallback.
ctx := WithHTTPUpstreamRedirectsDisabled(nil)
Expand Down
Loading