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

Long-running LLM inferences (such as high-thinking Gemini models) may experience
silent TCP disconnections across intermediate proxies, NAT gateways, or firewalls
when no data frames are emitted for tens of seconds to minutes.

Go's default http2.Transport sets ReadIdleTimeout=0 (disabling active PING checks).
This commit introduces HTTPUpstreamProfileLongStream with:
- ReadIdleTimeout = 10s
- PingTimeout = 5s

When a pooled H2 connection stays idle for 10s during long streams, active PING frames
are dispatched. Dead connections failing to respond within 5s are proactively evicted,
preventing client generation requests from hanging on dead sockets or failing with EOF/502.

Signed-off-by: WUJI-Labs <zoujincool@gmail.com>
This commit is contained in:
WUJI-Labs
2026-09-02 12:11:17 +08:00
committed by 无极
parent 5097b31457
commit 81fd853005
4 changed files with 41 additions and 11 deletions
+14 -10
View File
@@ -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,
@@ -96,6 +96,7 @@ const (
const (
upstreamProtocolModeDefault = "default"
upstreamProtocolModeLongStreamH2 = "long_stream_h2"
upstreamProtocolModeOpenAIH1 = "openai_h1"
upstreamProtocolModeOpenAIH2 = "openai_h2"
upstreamProtocolModeOpenAIH1Fallback = "openai_h1_fallback"
@@ -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 时不可淘汰
}
@@ -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
}
@@ -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:
@@ -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
}
@@ -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,
@@ -10,6 +10,7 @@ const (
HTTPUpstreamProfileDefault HTTPUpstreamProfile = ""
HTTPUpstreamProfileOpenAI HTTPUpstreamProfile = "openai"
HTTPUpstreamProfileGrok HTTPUpstreamProfile = "grok"
HTTPUpstreamProfileLongStream HTTPUpstreamProfile = "long_stream"
)
type httpUpstreamProfileContextKey struct{}
@@ -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
@@ -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)