fix(relay): keep long JSON generation requests alive - #4506
Closed
gaoren002 wants to merge 2 commits into
Closed
Conversation
Contributor
WalkthroughIntroduces an unexported JSON keepalive that emits periodic HTTP 102 Processing responses and integrates it into multiple relay handlers (image, mjproxy, relay task) to run during outbound requests, with explicit and deferred stop semantics to ensure cleanup. Changes
Sequence DiagramsequenceDiagram
participant Client
participant Handler as HTTP Handler
participant Keepalive as JSON Keepalive
participant Remote as Remote Service
participant RW as ResponseWriter
Client->>Handler: Long HTTP request
Handler->>Keepalive: startJSONKeepalive(ctx, initialDelay, interval)
Keepalive->>RW: snapshot headers, set processing headers
Handler->>Remote: DoRequest / DoMidjourneyHttpRequest
par Keepalive ticker vs Remote request
Keepalive->>RW: Write 102 Processing, Flush
Remote-->>Remote: remote processing
end
Remote-->>Handler: Response (or error)
Handler->>Keepalive: keepalive.stop()
Keepalive->>Keepalive: signal stop, wait done
Handler->>RW: Write final status/body
Handler-->>Client: Final HTTP response
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~25 minutes Possibly related PRs
Suggested reviewers
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
relay/json_keepalive_test.go (1)
71-108: Cover the repeated-tick path and the post-stop()quiet period.The current tests prove that one
102 Processingcan be emitted and that the final JSON status survives, but they do not lock in the two behaviors this loop depends on most: sending another 102 on the next interval and emitting nothing afterstop()returns. A regression in either branch would still pass here.🧪 Suggested additions
+func TestJSONKeepaliveRepeatsUntilStopped(t *testing.T) { + rec := newJSONKeepaliveRecorder() + c := newJSONKeepaliveTestContext(rec) + + keepalive := startJSONKeepalive(c, 5*time.Millisecond, 5*time.Millisecond) + require.NotNil(t, keepalive) + require.Eventually(t, func() bool { + informational, _, _, _ := rec.snapshot() + return len(informational) >= 2 + }, time.Second, 5*time.Millisecond) + keepalive.stop() +} + +func TestJSONKeepaliveStopPreventsFurtherWrites(t *testing.T) { + rec := newJSONKeepaliveRecorder() + c := newJSONKeepaliveTestContext(rec) + + keepalive := startJSONKeepalive(c, 5*time.Millisecond, 5*time.Millisecond) + require.NotNil(t, keepalive) + require.Eventually(t, func() bool { + return keepalive.wasWritten() + }, time.Second, 5*time.Millisecond) + + before, _, _, _ := rec.snapshot() + keepalive.stop() + time.Sleep(20 * time.Millisecond) + after, _, _, _ := rec.snapshot() + + require.Len(t, after, len(before)) +}🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@relay/json_keepalive_test.go` around lines 71 - 108, Extend the existing tests (TestJSONKeepaliveSendsInformationalProcessing and/or TestJSONKeepalivePreservesFinalJSONStatus) to exercise the repeated-tick and post-stop quiet-period paths: after confirming the first 102 was written (using startJSONKeepalive and rec.snapshot / keepalive.wasWritten), wait for one more tick interval plus a small margin and assert a second informational 102 was emitted (e.g., len(informational) increased by 1 or informational[1] == http.StatusProcessing), then call keepalive.stop(), wait for a short quiet period (greater than the write interval) and assert no further writes occur (snapshot shows same flush/body/status counts and no new informational codes after stop returned). Use the existing helpers newJSONKeepaliveRecorder, newJSONKeepaliveTestContext, startJSONKeepalive, keepalive.stop, rec.snapshot and wasWritten to locate and implement these assertions.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@relay/json_keepalive.go`:
- Around line 30-32: The three headers set by the c.Header calls (Content-Type,
Cache-Control, X-Accel-Buffering) must only be attached when you actually emit
the 102 keepalive; move those c.Header(...) calls out of the fast-path and into
the branch or helper that writes the interim 102 response (the code that
performs the keepalive write/emit), e.g., the function or block that sends the
102 keepalive (where you currently call WriteHeader or send the keepalive byte).
Remove the global header-setting from the initial path so fast responses (and
other handlers like relay/image_handler.go) do not inherit application/json or
X-Accel-Buffering: no. Ensure headers are set immediately before sending the 102
so they accompany only the keepalive.
---
Nitpick comments:
In `@relay/json_keepalive_test.go`:
- Around line 71-108: Extend the existing tests
(TestJSONKeepaliveSendsInformationalProcessing and/or
TestJSONKeepalivePreservesFinalJSONStatus) to exercise the repeated-tick and
post-stop quiet-period paths: after confirming the first 102 was written (using
startJSONKeepalive and rec.snapshot / keepalive.wasWritten), wait for one more
tick interval plus a small margin and assert a second informational 102 was
emitted (e.g., len(informational) increased by 1 or informational[1] ==
http.StatusProcessing), then call keepalive.stop(), wait for a short quiet
period (greater than the write interval) and assert no further writes occur
(snapshot shows same flush/body/status counts and no new informational codes
after stop returned). Use the existing helpers newJSONKeepaliveRecorder,
newJSONKeepaliveTestContext, startJSONKeepalive, keepalive.stop, rec.snapshot
and wasWritten to locate and implement these assertions.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: 0b1c2f57-09f7-45e2-8794-3e1e2107bb33
📒 Files selected for processing (5)
relay/image_handler.gorelay/json_keepalive.gorelay/json_keepalive_test.gorelay/mjproxy_handler.gorelay/relay_task.go
Contributor
Author
gaoren002
force-pushed
the
pr/image-json-keepalive
branch
from
April 28, 2026 08:06
3a7f316 to
d55f048
Compare
Contributor
There was a problem hiding this comment.
🧹 Nitpick comments (2)
relay/json_keepalive_test.go (1)
28-30: Consider adding mutex protection toHeader()for race-free testing.The
Header()method returnsr.headerwithout acquiring the lock. During theEventuallypolling loops, the keepalive goroutine may be modifying the header map concurrently withsnapshot()cloning it. While the final assertions occur afterstop()(blocking until goroutine exit), the race detector may flag this during intermediate polling.🔧 Suggested fix
func (r *jsonKeepaliveRecorder) Header() http.Header { + r.mu.Lock() + defer r.mu.Unlock() return r.header }Alternatively, since
Header()is called by the production code via the unwrapped writer and needs to return the same map instance for modifications, you could leave it as-is and accept the benign race in test code.🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@relay/json_keepalive_test.go` around lines 28 - 30, The Header() method of jsonKeepaliveRecorder returns r.header without locking, causing a data race during Eventually polling; fix this by acquiring the recorder's mutex before returning the header (i.e., lock the same mutex used by snapshot()/stop() inside jsonKeepaliveRecorder.Header()), then return r.header and unlock, ensuring Header() is race-free during concurrent header cloning and modifications in the keepalive goroutine.relay/json_keepalive.go (1)
12-15: Consider documenting the timeout constants' rationale.The 75-second initial delay and 25-second interval appear well-chosen (75s is under typical proxy timeouts like nginx's 60s default plus buffer; 25s provides multiple keepalives before most 60s idle timeouts). A brief comment explaining this would aid future maintainability.
📝 Suggested documentation
const ( + // jsonKeepaliveInitialDelay: wait before first 102; set below common proxy + // read-timeout defaults (e.g., nginx 60s) plus margin for slow cold-starts. jsonKeepaliveInitialDelay = 75 * time.Second + // jsonKeepaliveInterval: subsequent 102s to refresh proxy/CDN idle timers; + // 25s keeps connections alive under most 60s idle-timeout configurations. jsonKeepaliveInterval = 25 * time.Second )🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@relay/json_keepalive.go` around lines 12 - 15, Add a brief comment above the jsonKeepaliveInitialDelay and jsonKeepaliveInterval constants explaining the rationale for their values (e.g., 75s initial delay chosen to sit just under/around common proxy/idle timeouts with buffer, and 25s interval to ensure multiple keepalives within a typical 60s idle window); update the comment near the constants jsonKeepaliveInitialDelay and jsonKeepaliveInterval in json_keepalive.go so future readers understand why these specific timeouts were chosen.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Nitpick comments:
In `@relay/json_keepalive_test.go`:
- Around line 28-30: The Header() method of jsonKeepaliveRecorder returns
r.header without locking, causing a data race during Eventually polling; fix
this by acquiring the recorder's mutex before returning the header (i.e., lock
the same mutex used by snapshot()/stop() inside jsonKeepaliveRecorder.Header()),
then return r.header and unlock, ensuring Header() is race-free during
concurrent header cloning and modifications in the keepalive goroutine.
In `@relay/json_keepalive.go`:
- Around line 12-15: Add a brief comment above the jsonKeepaliveInitialDelay and
jsonKeepaliveInterval constants explaining the rationale for their values (e.g.,
75s initial delay chosen to sit just under/around common proxy/idle timeouts
with buffer, and 25s interval to ensure multiple keepalives within a typical 60s
idle window); update the comment near the constants jsonKeepaliveInitialDelay
and jsonKeepaliveInterval in json_keepalive.go so future readers understand why
these specific timeouts were chosen.
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: edc7d3ae-2fa5-46ba-b220-c954c4ce9950
📒 Files selected for processing (5)
relay/image_handler.gorelay/json_keepalive.gorelay/json_keepalive_test.gorelay/mjproxy_handler.gorelay/relay_task.go
🚧 Files skipped from review as they are similar to previous changes (2)
- relay/image_handler.go
- relay/mjproxy_handler.go
Contributor
Author
|
需要指出，这个不一定总有效，我只在cf直接代理时测试生效 cf 隧道好像不行 |
Contributor
Author
|
属于不稳定的解决方案，应该用流来解决此问题 |
11 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
背景
图片、视频、Midjourney 等生成类接口可能在上游处理阶段超过 CDN 或反向代理的空闲超时时间。此时客户端连接可能先被断开，表现为 524、超时或响应体断流，但上游请求实际上可能仍在执行并产生费用。
改动
102 Processing信息状态码。102 Processing，刷新连接空闲时间。200、4xx或5xx响应。适用范围
/v1/images/generations、/v1/images/edits。/v1/video/generations、/v1/videos、/v1/videos/:video_id/remix、/suno/submit/*、/kling/v1/videos/*、/jimeng等走RelayTaskSubmit的 JSON 提交路径。/mj/submit/*、/mj/insight-face/swap、/mj/task/:id/image-seed。1xx信息响应，最终响应仍保持兼容，只是该代理层可能无法获得保活效果。测试
测试覆盖：
102 Processing并 flush。102 Processing不会覆盖最终 JSON 响应状态码。Summary by CodeRabbit
New Features
Bug Fixes
Tests