From d88828095ae960c0a935f93c66a90c9d69eb0c3b Mon Sep 17 00:00:00 2001 From: akshaydeo Date: Tue, 19 May 2026 16:59:01 +0530 Subject: [PATCH] fix the race condition for remote stream close on context cancel --- .../{large_response.go => largeresponse.go} | 33 ++++++++++++++++--- ...ke_request_test.go => makerequest_test.go} | 0 ...client_test.go => streamingclient_test.go} | 0 core/providers/utils/utils.go | 16 +++++---- .../{utils_json_test.go => utilsjson_test.go} | 0 .../workspace/config/feature-flags/layout.tsx | 2 +- 6 files changed, 38 insertions(+), 13 deletions(-) rename core/providers/utils/{large_response.go => largeresponse.go} (90%) rename core/providers/utils/{make_request_test.go => makerequest_test.go} (100%) rename core/providers/utils/{streaming_client_test.go => streamingclient_test.go} (100%) rename core/providers/utils/{utils_json_test.go => utilsjson_test.go} (100%) diff --git a/core/providers/utils/large_response.go b/core/providers/utils/largeresponse.go similarity index 90% rename from core/providers/utils/large_response.go rename to core/providers/utils/largeresponse.go index c1c5da8a15b..bb6eb3cee0f 100644 --- a/core/providers/utils/large_response.go +++ b/core/providers/utils/largeresponse.go @@ -12,11 +12,14 @@ import ( // LargeResponseReader wraps an io.Reader and releases the fasthttp response on Close. // Used by providers to keep the response alive while the transport streams it to the client. +// ctx is held to check BifrostContextKeyConnectionClosed in Close, so a mid-stream +// cancellation that already tore down the underlying fasthttp conn does not double-release. type LargeResponseReader struct { io.Reader Resp *fasthttp.Response + ctx *schemas.BifrostContext cleanup func() - consumed bool // true after Read returns io.EOF — body fully consumed through Reader chain + consumed bool // true after Read returns io.EOF, body fully consumed through Reader chain } // Read delegates to the wrapped Reader and tracks EOF so Close() can skip @@ -42,6 +45,20 @@ func (r *LargeResponseReader) Close() error { if r == nil || r.Resp == nil { return nil } + // Run cleanup first so SetupStreamCancellation's goroutine settles (close(done); <-closed) + // before we read BifrostContextKeyConnectionClosed. The goroutine's done-branch can set the + // flag when ctx.Err() != nil, so checking it before cleanup would miss that interleaving and + // fall through to fasthttp.ReleaseResponse on an already-torn-down conn (nil-deref in connsCleaner). + if r.cleanup != nil { + r.cleanup() + r.cleanup = nil + } + if r.ctx != nil { + if closed, ok := r.ctx.Value(schemas.BifrostContextKeyConnectionClosed).(bool); ok && closed { + r.Resp = nil + return nil + } + } if !r.consumed { if bodyStream := r.Resp.BodyStream(); bodyStream != nil { _, _ = io.Copy(io.Discard, bodyStream) @@ -50,10 +67,6 @@ func (r *LargeResponseReader) Close() error { } } } - if r.cleanup != nil { - r.cleanup() - r.cleanup = nil - } fasthttp.ReleaseResponse(r.Resp) r.Resp = nil return nil @@ -186,6 +199,7 @@ func FinalizeResponseWithLargeDetection( closableReader := &LargeResponseReader{ Reader: combinedReader, Resp: resp, + ctx: ctx, cleanup: releaseGzip, } ctx.SetValue(schemas.BifrostContextKeyLargeResponseMode, true) @@ -246,6 +260,7 @@ func FinalizeResponseWithLargeDetection( closableReader := &LargeResponseReader{ Reader: combinedReader, Resp: resp, + ctx: ctx, cleanup: func() { if wasGzip { ReleaseGzipReader(gz) @@ -321,10 +336,18 @@ func SetupStreamingPassthrough(ctx *schemas.BifrostContext, resp *fasthttp.Respo // Wrap reader with idle timeout to detect stalled streams. reader, stopIdleTimeout := NewIdleTimeoutReader(reader, resp.BodyStream(), GetStreamIdleTimeout(ctx)) + // Wire cancellation to the raw fasthttp body. On a mid-stream client disconnect this fires + // wce.CloseWithError(ctx.Err()) to unblock the transport's Read and sets + // BifrostContextKeyConnectionClosed so LargeResponseReader.Close skips the double release. + // logger arg is unused inside SetupStreamCancellation (uses package getLogger), nil is safe. + stopCancellation := SetupStreamCancellation(ctx, resp.BodyStream(), nil) + closableReader := &LargeResponseReader{ Reader: reader, Resp: resp, + ctx: ctx, cleanup: func() { + stopCancellation() stopIdleTimeout() releaseGzip() }, diff --git a/core/providers/utils/make_request_test.go b/core/providers/utils/makerequest_test.go similarity index 100% rename from core/providers/utils/make_request_test.go rename to core/providers/utils/makerequest_test.go diff --git a/core/providers/utils/streaming_client_test.go b/core/providers/utils/streamingclient_test.go similarity index 100% rename from core/providers/utils/streaming_client_test.go rename to core/providers/utils/streamingclient_test.go diff --git a/core/providers/utils/utils.go b/core/providers/utils/utils.go index bfb2a39972e..5e08c69ba16 100644 --- a/core/providers/utils/utils.go +++ b/core/providers/utils/utils.go @@ -2062,22 +2062,24 @@ func SetupStreamCancellation(ctx *schemas.BifrostContext, bodyStream io.Reader, ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) } case <-done: - // If context was also cancelled (race between done and ctx.Done), - // still close the body stream to unblock the drain in ReleaseStreamingResponse. + // Race between done and ctx.Done: the streaming goroutine has reached its defer + // chain (Read has returned), and ctx is also cancelled. The body may already be + // at EOF and fasthttp may have released the underlying conn to the idle pool. + // We still attempt a close to unblock any pending drain in ReleaseStreamingResponse, + // but we set BifrostContextKeyConnectionClosed unconditionally (matching the + // ctx.Done branch above) so ReleaseStreamingResponse skips a second CloseWithError. + // A second close against an already-pooled conn nil-derefs in fasthttp's connsCleaner. if ctx.Err() != nil { if closer, ok := bodyStream.(io.Closer); ok { if err := closer.Close(); err != nil { getLogger().Debug(fmt.Sprintf("Error closing body stream on done with cancelled context: %v", err)) - } else { - ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) } + ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) } else if wce, ok := bodyStream.(streamCloserWithError); ok { if err := wce.CloseWithError(ctx.Err()); err != nil { getLogger().Debug(fmt.Sprintf("Error closing body stream on done with cancelled context: %v", err)) - } else { - ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) } - + ctx.SetValue(schemas.BifrostContextKeyConnectionClosed, true) } } } diff --git a/core/providers/utils/utils_json_test.go b/core/providers/utils/utilsjson_test.go similarity index 100% rename from core/providers/utils/utils_json_test.go rename to core/providers/utils/utilsjson_test.go diff --git a/ui/app/workspace/config/feature-flags/layout.tsx b/ui/app/workspace/config/feature-flags/layout.tsx index 4ec93465b07..d252b506d99 100644 --- a/ui/app/workspace/config/feature-flags/layout.tsx +++ b/ui/app/workspace/config/feature-flags/layout.tsx @@ -2,5 +2,5 @@ import { createFileRoute } from "@tanstack/react-router"; import FeatureFlagsPage from "./page"; export const Route = createFileRoute("/workspace/config/feature-flags")({ - component: FeatureFlagsPage, + component: FeatureFlagsPage, });