From f67608ae2196aff308508b06c28d2420e5122804 Mon Sep 17 00:00:00 2001 From: Psych0h3ad <41975091+Psych0h3ad@users.noreply.github.com> Date: Sun, 6 Sep 2026 00:43:20 +0000 Subject: [PATCH 1/2] Fix HTTP connection reuse in engine health probes Drain bounded HTTP response bodies before closing health checks so periodic polling can reuse connections. Preserve probe deadlines and status-based results, and cover HTTP/1 reuse, size limits, and timeouts. Signed-off-by: Psych0h3ad <41975091+Psych0h3ad@users.noreply.github.com> --- .../health_connections_test.go | 158 ++++++++++++++++ services/nvpair-engine-manager/lifecycle.go | 6 + services/nvpair-ui-broker/advertiser.go | 6 + .../health_connections_test.go | 171 ++++++++++++++++++ 4 files changed, 341 insertions(+) create mode 100644 services/nvpair-engine-manager/health_connections_test.go create mode 100644 services/nvpair-ui-broker/health_connections_test.go diff --git a/services/nvpair-engine-manager/health_connections_test.go b/services/nvpair-engine-manager/health_connections_test.go new file mode 100644 index 00000000..16a1d156 --- /dev/null +++ b/services/nvpair-engine-manager/health_connections_test.go @@ -0,0 +1,158 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "io" + "net" + "net/http" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +func TestProbeHTTPReusesConnections(t *testing.T) { + for _, chunked := range []bool{false, true} { + framing := "content-length" + if chunked { + framing = "chunked" + } + t.Run(framing, func(t *testing.T) { + var requests atomic.Int32 + client, connections := healthProbePipeClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.Header.Get(engineIdentityProbeHeader) != "1" { + t.Error("missing engine identity header") + } + status := http.StatusOK + if requests.Add(1)%2 == 0 { + status = http.StatusServiceUnavailable + } + w.WriteHeader(status) + if chunked { + _ = http.NewResponseController(w).Flush() + } + _, _ = io.WriteString(w, `{"models":[],"status":"responding"}`) + })) + ex := &Executor{client: client} + probe := &Probe{HTTP: "http://127.0.0.1:{port}/api/version"} + const rounds = 32 + for i := 0; i < rounds; i++ { + if got, want := ex.probe(context.Background(), probe, 1), i%2 == 0; got != want { + t.Fatalf("poll %d: healthy = %v, want %v", i, got, want) + } + } + if got := connections.Load(); got != 1 { + t.Fatalf("accepted %d HTTP/1 connections for %d polls, want 1", got, rounds) + } + }) + } +} + +func TestProbeHTTPBoundsBodyDrain(t *testing.T) { + body := &healthProbeBody{reader: strings.NewReader(strings.Repeat("x", 4<<20))} + ex := &Executor{client: &http.Client{Transport: healthProbeTransport(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusAccepted, Body: body, Header: make(http.Header)}, nil + })}} + probe := &Probe{HTTP: "http://127.0.0.1:{port}/", Status: http.StatusAccepted} + if !ex.probe(context.Background(), probe, 1) { + t.Fatal("body drain changed the configured HTTP status health result") + } + if body.read == 0 || body.read > 1<<20 || !body.closed { + t.Fatalf("body read = %d, closed = %v; want bounded drain and close", body.read, body.closed) + } +} + +func TestProbeHTTPBodyDrainHonorsDeadline(t *testing.T) { + var body *healthProbeBody + ex := &Executor{client: &http.Client{Transport: healthProbeTransport(func(req *http.Request) (*http.Response, error) { + body = &healthProbeBody{ctx: req.Context()} + return &http.Response{StatusCode: http.StatusOK, Body: body, Header: make(http.Header)}, nil + })}} + probe := &Probe{HTTP: "http://127.0.0.1:{port}/"} + ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond) + defer cancel() + start := time.Now() + if !ex.probe(ctx, probe, 1) { + t.Fatal("body drain changed the HTTP status health result") + } + if !body.attempted || !body.closed || time.Since(start) > time.Second { + t.Fatalf("stalled body did not drain and close within deadline: %+v", body) + } +} + +type healthProbeTransport func(*http.Request) (*http.Response, error) + +func (f healthProbeTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +type healthProbeBody struct { + reader io.Reader + ctx context.Context + read int + attempted bool + closed bool +} + +func (b *healthProbeBody) Read(p []byte) (int, error) { + b.attempted = true + if b.ctx != nil { + <-b.ctx.Done() + return 0, b.ctx.Err() + } + n, err := b.reader.Read(p) + b.read += n + return n, err +} + +func (b *healthProbeBody) Close() error { b.closed = true; return nil } + +// Exercise the real HTTP/1 transport over net.Pipe so this regression can run +// even on a machine whose TCP source ports have already been exhausted. +func healthProbePipeClient(t *testing.T, handler http.Handler) (*http.Client, *atomic.Int32) { + t.Helper() + listener := &healthProbeListener{conns: make(chan net.Conn), done: make(chan struct{})} + server := &http.Server{Handler: handler} + go func() { _ = server.Serve(listener) }() + t.Cleanup(func() { _ = server.Close() }) + connections := &atomic.Int32{} + transport := &http.Transport{DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) { + clientConn, serverConn := net.Pipe() + select { + case listener.conns <- serverConn: + connections.Add(1) + return clientConn, nil + case <-ctx.Done(): + _ = clientConn.Close() + _ = serverConn.Close() + return nil, ctx.Err() + } + }} + client := &http.Client{Transport: transport, Timeout: 2 * time.Second} + t.Cleanup(client.CloseIdleConnections) + return client, connections +} + +type healthProbeListener struct { + conns chan net.Conn + done chan struct{} + once sync.Once +} + +func (l *healthProbeListener) Accept() (net.Conn, error) { + select { + case conn := <-l.conns: + return conn, nil + case <-l.done: + return nil, net.ErrClosed + } +} + +func (l *healthProbeListener) Close() error { + l.once.Do(func() { close(l.done) }) + return nil +} + +func (l *healthProbeListener) Addr() net.Addr { return &net.TCPAddr{Port: 1} } diff --git a/services/nvpair-engine-manager/lifecycle.go b/services/nvpair-engine-manager/lifecycle.go index cea7aa16..e6ff20a4 100644 --- a/services/nvpair-engine-manager/lifecycle.go +++ b/services/nvpair-engine-manager/lifecycle.go @@ -7,6 +7,7 @@ import ( "context" "errors" "fmt" + "io" "log/slog" "net" "net/http" @@ -24,6 +25,7 @@ import ( const ( unavailableConfirmations = 3 engineIdentityProbeHeader = "X-NVPAIR-Engine-Identity-Probe" + maxHealthProbeBodyBytes = 1 << 20 ) type listenerProbeResult uint8 @@ -871,6 +873,10 @@ func (e *Executor) probe(ctx context.Context, p *Probe, port int) bool { if err != nil { return false } + // A body closed before EOF discards its HTTP/1 connection on every + // health poll. Drain normal responses for reuse, bounded by both size + // and the existing probe deadline for oversized or stalled bodies. + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, maxHealthProbeBodyBytes)) resp.Body.Close() want := p.Status if want == 0 { diff --git a/services/nvpair-ui-broker/advertiser.go b/services/nvpair-ui-broker/advertiser.go index 05d85fdf..ab4f6331 100644 --- a/services/nvpair-ui-broker/advertiser.go +++ b/services/nvpair-ui-broker/advertiser.go @@ -7,6 +7,7 @@ import ( "context" "encoding/json" "fmt" + "io" "net/http" "time" @@ -43,6 +44,10 @@ const ( // decide whether to register it with (or unregister it from) the discovery // daemon (a 5s cadence). autoAdvertiseInterval = 5 * time.Second + + // Health bodies normally contain a short status or model list. Bound the + // discard in case an engine returns an unexpectedly large response. + maxHealthProbeBodyBytes = 1 << 20 ) // runAutoAdvertise is the broker's ollama engine-registration loop. It polls @@ -274,6 +279,7 @@ func checkEngineHealth(profile engineProxyProfile, client *http.Client, port int if err != nil { return false } + _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, maxHealthProbeBodyBytes)) resp.Body.Close() return resp.StatusCode == http.StatusOK } diff --git a/services/nvpair-ui-broker/health_connections_test.go b/services/nvpair-ui-broker/health_connections_test.go new file mode 100644 index 00000000..c9830bdb --- /dev/null +++ b/services/nvpair-ui-broker/health_connections_test.go @@ -0,0 +1,171 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "io" + "net" + "net/http" + "strings" + "sync" + "sync/atomic" + "testing" + "time" +) + +var healthProbeChecks = []struct { + name string + path string + profile engineProxyProfile +}{ + {"ollama", "/", ollamaProxyProfile}, + {"lmstudio", "/v1/models", lmstudioProxyProfile}, +} + +func TestHealthChecksReuseConnections(t *testing.T) { + for _, check := range healthProbeChecks { + for _, chunked := range []bool{false, true} { + framing := "content-length" + if chunked { + framing = "chunked" + } + t.Run(check.name+"/"+framing, func(t *testing.T) { + var requests atomic.Int32 + client, connections := healthProbePipeClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != check.path { + t.Errorf("path = %q, want %q", r.URL.Path, check.path) + } + status := http.StatusOK + if requests.Add(1)%2 == 0 { + status = http.StatusServiceUnavailable + } + w.WriteHeader(status) + if chunked { + _ = http.NewResponseController(w).Flush() + } + _, _ = io.WriteString(w, `{"models":[],"status":"responding"}`) + })) + const rounds = 32 + for i := 0; i < rounds; i++ { + if got, want := checkEngineHealth(check.profile, client, 1), i%2 == 0; got != want { + t.Fatalf("poll %d: healthy = %v, want %v", i, got, want) + } + } + if got := connections.Load(); got != 1 { + t.Fatalf("accepted %d HTTP/1 connections for %d polls, want 1", got, rounds) + } + }) + } + } +} + +func TestHealthChecksBoundBodyDrain(t *testing.T) { + for _, check := range healthProbeChecks { + t.Run(check.name, func(t *testing.T) { + body := &healthProbeBody{reader: strings.NewReader(strings.Repeat("x", 4<<20))} + client := &http.Client{Transport: healthProbeTransport(func(*http.Request) (*http.Response, error) { + return &http.Response{StatusCode: http.StatusOK, Body: body, Header: make(http.Header)}, nil + })} + if !checkEngineHealth(check.profile, client, 1) { + t.Fatal("body drain changed the HTTP status health result") + } + if body.read == 0 || body.read > 1<<20 || !body.closed { + t.Fatalf("body read = %d, closed = %v; want bounded drain and close", body.read, body.closed) + } + }) + } +} + +func TestHealthChecksBodyDrainHonorsTimeout(t *testing.T) { + for _, check := range healthProbeChecks { + t.Run(check.name, func(t *testing.T) { + var body *healthProbeBody + client := &http.Client{Timeout: 50 * time.Millisecond, Transport: healthProbeTransport(func(req *http.Request) (*http.Response, error) { + body = &healthProbeBody{ctx: req.Context()} + return &http.Response{StatusCode: http.StatusOK, Body: body, Header: make(http.Header)}, nil + })} + start := time.Now() + if !checkEngineHealth(check.profile, client, 1) { + t.Fatal("body drain changed the HTTP status health result") + } + if !body.attempted || !body.closed || time.Since(start) > time.Second { + t.Fatalf("stalled body did not drain and close within timeout: %+v", body) + } + }) + } +} + +type healthProbeTransport func(*http.Request) (*http.Response, error) + +func (f healthProbeTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } + +type healthProbeBody struct { + reader io.Reader + ctx context.Context + read int + attempted bool + closed bool +} + +func (b *healthProbeBody) Read(p []byte) (int, error) { + b.attempted = true + if b.ctx != nil { + <-b.ctx.Done() + return 0, b.ctx.Err() + } + n, err := b.reader.Read(p) + b.read += n + return n, err +} + +func (b *healthProbeBody) Close() error { b.closed = true; return nil } + +// Exercise the real HTTP/1 transport over net.Pipe so this regression can run +// even on a machine whose TCP source ports have already been exhausted. +func healthProbePipeClient(t *testing.T, handler http.Handler) (*http.Client, *atomic.Int32) { + t.Helper() + listener := &healthProbeListener{conns: make(chan net.Conn), done: make(chan struct{})} + server := &http.Server{Handler: handler} + go func() { _ = server.Serve(listener) }() + t.Cleanup(func() { _ = server.Close() }) + connections := &atomic.Int32{} + transport := &http.Transport{DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) { + clientConn, serverConn := net.Pipe() + select { + case listener.conns <- serverConn: + connections.Add(1) + return clientConn, nil + case <-ctx.Done(): + _ = clientConn.Close() + _ = serverConn.Close() + return nil, ctx.Err() + } + }} + client := &http.Client{Transport: transport, Timeout: 2 * time.Second} + t.Cleanup(client.CloseIdleConnections) + return client, connections +} + +type healthProbeListener struct { + conns chan net.Conn + done chan struct{} + once sync.Once +} + +func (l *healthProbeListener) Accept() (net.Conn, error) { + select { + case conn := <-l.conns: + return conn, nil + case <-l.done: + return nil, net.ErrClosed + } +} + +func (l *healthProbeListener) Close() error { + l.once.Do(func() { close(l.done) }) + return nil +} + +func (l *healthProbeListener) Addr() net.Addr { return &net.TCPAddr{Port: 1} } From 6ee0b40b5c7dcff8c95a70c79a073b5838977d81 Mon Sep 17 00:00:00 2001 From: Kaylee Lubick Date: Tue, 15 Sep 2026 16:54:16 -0400 Subject: [PATCH 2/2] Extract helper code Signed-off-by: Kaylee Lubick --- .../health_connections_test.go | 67 +------ services/nvpair-engine-manager/lifecycle.go | 11 +- services/nvpair-ui-broker/advertiser.go | 9 +- .../health_connections_test.go | 183 ++++-------------- services/shared/httpcon/httpcon.go | 19 ++ services/shared/httpcon/testclient/client.go | 75 +++++++ .../shared/httpcon/testclient/client_test.go | 39 ++++ 7 files changed, 181 insertions(+), 222 deletions(-) create mode 100644 services/shared/httpcon/httpcon.go create mode 100644 services/shared/httpcon/testclient/client.go create mode 100644 services/shared/httpcon/testclient/client_test.go diff --git a/services/nvpair-engine-manager/health_connections_test.go b/services/nvpair-engine-manager/health_connections_test.go index 16a1d156..44e13a3a 100644 --- a/services/nvpair-engine-manager/health_connections_test.go +++ b/services/nvpair-engine-manager/health_connections_test.go @@ -6,24 +6,20 @@ package main import ( "context" "io" - "net" "net/http" "strings" - "sync" "sync/atomic" "testing" "time" + + "nvpair-shared/httpcon/testclient" ) func TestProbeHTTPReusesConnections(t *testing.T) { - for _, chunked := range []bool{false, true} { - framing := "content-length" - if chunked { - framing = "chunked" - } - t.Run(framing, func(t *testing.T) { + test := func(name string, chunked bool) { + t.Run(name, func(t *testing.T) { var requests atomic.Int32 - client, connections := healthProbePipeClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + client, connections := testclient.New(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if r.Header.Get(engineIdentityProbeHeader) != "1" { t.Error("missing engine identity header") } @@ -45,11 +41,14 @@ func TestProbeHTTPReusesConnections(t *testing.T) { t.Fatalf("poll %d: healthy = %v, want %v", i, got, want) } } - if got := connections.Load(); got != 1 { + if got := connections.Count(); got != 1 { t.Fatalf("accepted %d HTTP/1 connections for %d polls, want 1", got, rounds) } }) } + + test("content-length", false) + test("chunked", true) } func TestProbeHTTPBoundsBodyDrain(t *testing.T) { @@ -108,51 +107,3 @@ func (b *healthProbeBody) Read(p []byte) (int, error) { } func (b *healthProbeBody) Close() error { b.closed = true; return nil } - -// Exercise the real HTTP/1 transport over net.Pipe so this regression can run -// even on a machine whose TCP source ports have already been exhausted. -func healthProbePipeClient(t *testing.T, handler http.Handler) (*http.Client, *atomic.Int32) { - t.Helper() - listener := &healthProbeListener{conns: make(chan net.Conn), done: make(chan struct{})} - server := &http.Server{Handler: handler} - go func() { _ = server.Serve(listener) }() - t.Cleanup(func() { _ = server.Close() }) - connections := &atomic.Int32{} - transport := &http.Transport{DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) { - clientConn, serverConn := net.Pipe() - select { - case listener.conns <- serverConn: - connections.Add(1) - return clientConn, nil - case <-ctx.Done(): - _ = clientConn.Close() - _ = serverConn.Close() - return nil, ctx.Err() - } - }} - client := &http.Client{Transport: transport, Timeout: 2 * time.Second} - t.Cleanup(client.CloseIdleConnections) - return client, connections -} - -type healthProbeListener struct { - conns chan net.Conn - done chan struct{} - once sync.Once -} - -func (l *healthProbeListener) Accept() (net.Conn, error) { - select { - case conn := <-l.conns: - return conn, nil - case <-l.done: - return nil, net.ErrClosed - } -} - -func (l *healthProbeListener) Close() error { - l.once.Do(func() { close(l.done) }) - return nil -} - -func (l *healthProbeListener) Addr() net.Addr { return &net.TCPAddr{Port: 1} } diff --git a/services/nvpair-engine-manager/lifecycle.go b/services/nvpair-engine-manager/lifecycle.go index e6ff20a4..aa599903 100644 --- a/services/nvpair-engine-manager/lifecycle.go +++ b/services/nvpair-engine-manager/lifecycle.go @@ -7,7 +7,6 @@ import ( "context" "errors" "fmt" - "io" "log/slog" "net" "net/http" @@ -20,12 +19,13 @@ import ( "sync" "syscall" "time" + + "nvpair-shared/httpcon" ) const ( unavailableConfirmations = 3 engineIdentityProbeHeader = "X-NVPAIR-Engine-Identity-Probe" - maxHealthProbeBodyBytes = 1 << 20 ) type listenerProbeResult uint8 @@ -873,11 +873,8 @@ func (e *Executor) probe(ctx context.Context, p *Probe, port int) bool { if err != nil { return false } - // A body closed before EOF discards its HTTP/1 connection on every - // health poll. Drain normal responses for reuse, bounded by both size - // and the existing probe deadline for oversized or stalled bodies. - _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, maxHealthProbeBodyBytes)) - resp.Body.Close() + + httpcon.DrainAndClose(resp.Body) want := p.Status if want == 0 { want = 200 diff --git a/services/nvpair-ui-broker/advertiser.go b/services/nvpair-ui-broker/advertiser.go index ab4f6331..d48097e7 100644 --- a/services/nvpair-ui-broker/advertiser.go +++ b/services/nvpair-ui-broker/advertiser.go @@ -7,10 +7,10 @@ import ( "context" "encoding/json" "fmt" - "io" "net/http" "time" + "nvpair-shared/httpcon" "nvpair-shared/noderec" ) @@ -44,10 +44,6 @@ const ( // decide whether to register it with (or unregister it from) the discovery // daemon (a 5s cadence). autoAdvertiseInterval = 5 * time.Second - - // Health bodies normally contain a short status or model list. Bound the - // discard in case an engine returns an unexpectedly large response. - maxHealthProbeBodyBytes = 1 << 20 ) // runAutoAdvertise is the broker's ollama engine-registration loop. It polls @@ -279,7 +275,6 @@ func checkEngineHealth(profile engineProxyProfile, client *http.Client, port int if err != nil { return false } - _, _ = io.Copy(io.Discard, io.LimitReader(resp.Body, maxHealthProbeBodyBytes)) - resp.Body.Close() + httpcon.DrainAndClose(resp.Body) return resp.StatusCode == http.StatusOK } diff --git a/services/nvpair-ui-broker/health_connections_test.go b/services/nvpair-ui-broker/health_connections_test.go index c9830bdb..74e7231a 100644 --- a/services/nvpair-ui-broker/health_connections_test.go +++ b/services/nvpair-ui-broker/health_connections_test.go @@ -4,168 +4,51 @@ package main import ( - "context" "io" - "net" "net/http" - "strings" - "sync" "sync/atomic" "testing" - "time" -) -var healthProbeChecks = []struct { - name string - path string - profile engineProxyProfile -}{ - {"ollama", "/", ollamaProxyProfile}, - {"lmstudio", "/v1/models", lmstudioProxyProfile}, -} + "nvpair-shared/httpcon/testclient" +) func TestHealthChecksReuseConnections(t *testing.T) { - for _, check := range healthProbeChecks { - for _, chunked := range []bool{false, true} { - framing := "content-length" - if chunked { - framing = "chunked" - } - t.Run(check.name+"/"+framing, func(t *testing.T) { - var requests atomic.Int32 - client, connections := healthProbePipeClient(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { - if r.URL.Path != check.path { - t.Errorf("path = %q, want %q", r.URL.Path, check.path) + test := func(name, path string, profile engineProxyProfile) { + t.Run(name, func(t *testing.T) { + testFraming := func(name string, chunked bool) { + t.Run(name, func(t *testing.T) { + var requests atomic.Int32 + client, connections := testclient.New(t, http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + if r.URL.Path != path { + t.Errorf("path = %q, want %q", r.URL.Path, path) + } + status := http.StatusOK + if requests.Add(1)%2 == 0 { + status = http.StatusServiceUnavailable + } + w.WriteHeader(status) + if chunked { + _ = http.NewResponseController(w).Flush() + } + _, _ = io.WriteString(w, `{"models":[],"status":"responding"}`) + })) + const rounds = 32 + for i := 0; i < rounds; i++ { + if got, want := checkEngineHealth(profile, client, 1), i%2 == 0; got != want { + t.Fatalf("poll %d: healthy = %v, want %v", i, got, want) + } } - status := http.StatusOK - if requests.Add(1)%2 == 0 { - status = http.StatusServiceUnavailable + if got := connections.Count(); got != 1 { + t.Fatalf("accepted %d HTTP/1 connections for %d polls, want 1", got, rounds) } - w.WriteHeader(status) - if chunked { - _ = http.NewResponseController(w).Flush() - } - _, _ = io.WriteString(w, `{"models":[],"status":"responding"}`) - })) - const rounds = 32 - for i := 0; i < rounds; i++ { - if got, want := checkEngineHealth(check.profile, client, 1), i%2 == 0; got != want { - t.Fatalf("poll %d: healthy = %v, want %v", i, got, want) - } - } - if got := connections.Load(); got != 1 { - t.Fatalf("accepted %d HTTP/1 connections for %d polls, want 1", got, rounds) - } - }) - } - } -} - -func TestHealthChecksBoundBodyDrain(t *testing.T) { - for _, check := range healthProbeChecks { - t.Run(check.name, func(t *testing.T) { - body := &healthProbeBody{reader: strings.NewReader(strings.Repeat("x", 4<<20))} - client := &http.Client{Transport: healthProbeTransport(func(*http.Request) (*http.Response, error) { - return &http.Response{StatusCode: http.StatusOK, Body: body, Header: make(http.Header)}, nil - })} - if !checkEngineHealth(check.profile, client, 1) { - t.Fatal("body drain changed the HTTP status health result") - } - if body.read == 0 || body.read > 1<<20 || !body.closed { - t.Fatalf("body read = %d, closed = %v; want bounded drain and close", body.read, body.closed) + }) } - }) - } -} -func TestHealthChecksBodyDrainHonorsTimeout(t *testing.T) { - for _, check := range healthProbeChecks { - t.Run(check.name, func(t *testing.T) { - var body *healthProbeBody - client := &http.Client{Timeout: 50 * time.Millisecond, Transport: healthProbeTransport(func(req *http.Request) (*http.Response, error) { - body = &healthProbeBody{ctx: req.Context()} - return &http.Response{StatusCode: http.StatusOK, Body: body, Header: make(http.Header)}, nil - })} - start := time.Now() - if !checkEngineHealth(check.profile, client, 1) { - t.Fatal("body drain changed the HTTP status health result") - } - if !body.attempted || !body.closed || time.Since(start) > time.Second { - t.Fatalf("stalled body did not drain and close within timeout: %+v", body) - } + testFraming("content-length", false) + testFraming("chunked", true) }) } -} - -type healthProbeTransport func(*http.Request) (*http.Response, error) - -func (f healthProbeTransport) RoundTrip(r *http.Request) (*http.Response, error) { return f(r) } - -type healthProbeBody struct { - reader io.Reader - ctx context.Context - read int - attempted bool - closed bool -} - -func (b *healthProbeBody) Read(p []byte) (int, error) { - b.attempted = true - if b.ctx != nil { - <-b.ctx.Done() - return 0, b.ctx.Err() - } - n, err := b.reader.Read(p) - b.read += n - return n, err -} - -func (b *healthProbeBody) Close() error { b.closed = true; return nil } - -// Exercise the real HTTP/1 transport over net.Pipe so this regression can run -// even on a machine whose TCP source ports have already been exhausted. -func healthProbePipeClient(t *testing.T, handler http.Handler) (*http.Client, *atomic.Int32) { - t.Helper() - listener := &healthProbeListener{conns: make(chan net.Conn), done: make(chan struct{})} - server := &http.Server{Handler: handler} - go func() { _ = server.Serve(listener) }() - t.Cleanup(func() { _ = server.Close() }) - connections := &atomic.Int32{} - transport := &http.Transport{DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) { - clientConn, serverConn := net.Pipe() - select { - case listener.conns <- serverConn: - connections.Add(1) - return clientConn, nil - case <-ctx.Done(): - _ = clientConn.Close() - _ = serverConn.Close() - return nil, ctx.Err() - } - }} - client := &http.Client{Transport: transport, Timeout: 2 * time.Second} - t.Cleanup(client.CloseIdleConnections) - return client, connections -} - -type healthProbeListener struct { - conns chan net.Conn - done chan struct{} - once sync.Once -} - -func (l *healthProbeListener) Accept() (net.Conn, error) { - select { - case conn := <-l.conns: - return conn, nil - case <-l.done: - return nil, net.ErrClosed - } -} -func (l *healthProbeListener) Close() error { - l.once.Do(func() { close(l.done) }) - return nil + test("ollama", "/", ollamaProxyProfile) + test("lmstudio", "/v1/models", lmstudioProxyProfile) } - -func (l *healthProbeListener) Addr() net.Addr { return &net.TCPAddr{Port: 1} } diff --git a/services/shared/httpcon/httpcon.go b/services/shared/httpcon/httpcon.go new file mode 100644 index 00000000..7e928d90 --- /dev/null +++ b/services/shared/httpcon/httpcon.go @@ -0,0 +1,19 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Package httpcon provides helpers for HTTP connections. +package httpcon + +import "io" + +const maxDrainBytes = 1 << 20 + +// DrainAndClose discards a bounded response body before closing it so normal +// HTTP/1 responses reach EOF and their connections can return to the transport +// pool. Without this, connections cannot be reused and this can lead to socket +// exhaustion. The caller is responsible for bounding stalled reads with a +// request context or client timeout. +func DrainAndClose(body io.ReadCloser) { + _, _ = io.Copy(io.Discard, io.LimitReader(body, maxDrainBytes)) + _ = body.Close() +} diff --git a/services/shared/httpcon/testclient/client.go b/services/shared/httpcon/testclient/client.go new file mode 100644 index 00000000..67c9aedf --- /dev/null +++ b/services/shared/httpcon/testclient/client.go @@ -0,0 +1,75 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +// Package testclient provides in-memory HTTP clients for connection tests. +package testclient + +import ( + "context" + "net" + "net/http" + "sync" + "sync/atomic" + "testing" + "time" +) + +// ConnectionCounter counts connections accepted by a client returned by New. +type ConnectionCounter struct { + accepted atomic.Int32 +} + +// Count returns the number of HTTP/1 connections accepted by the test server. +func (c *ConnectionCounter) Count() int32 { + return c.accepted.Load() +} + +// New returns an HTTP client backed by an in-memory server and a counter for +// the connections that server accepts. It avoids consuming TCP source ports, +// so connection-reuse regressions remain testable after port exhaustion. +func New(t *testing.T, handler http.Handler) (*http.Client, *ConnectionCounter) { + t.Helper() + listener := &pipeListener{conns: make(chan net.Conn), done: make(chan struct{})} + server := &http.Server{Handler: handler} + go func() { _ = server.Serve(listener) }() + t.Cleanup(func() { _ = server.Close() }) + + connections := &ConnectionCounter{} + transport := &http.Transport{DialContext: func(ctx context.Context, _, _ string) (net.Conn, error) { + clientConn, serverConn := net.Pipe() + select { + case listener.conns <- serverConn: + connections.accepted.Add(1) + return clientConn, nil + case <-ctx.Done(): + _ = clientConn.Close() + _ = serverConn.Close() + return nil, ctx.Err() + } + }} + client := &http.Client{Transport: transport, Timeout: 2 * time.Second} + t.Cleanup(client.CloseIdleConnections) + return client, connections +} + +type pipeListener struct { + conns chan net.Conn + done chan struct{} + once sync.Once +} + +func (l *pipeListener) Accept() (net.Conn, error) { + select { + case conn := <-l.conns: + return conn, nil + case <-l.done: + return nil, net.ErrClosed + } +} + +func (l *pipeListener) Close() error { + l.once.Do(func() { close(l.done) }) + return nil +} + +func (l *pipeListener) Addr() net.Addr { return &net.TCPAddr{Port: 1} } diff --git a/services/shared/httpcon/testclient/client_test.go b/services/shared/httpcon/testclient/client_test.go new file mode 100644 index 00000000..522b6c4e --- /dev/null +++ b/services/shared/httpcon/testclient/client_test.go @@ -0,0 +1,39 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package testclient + +import ( + "io" + "net/http" + "testing" + + "nvpair-shared/httpcon" +) + +func TestConnectionCounterCountsDistinctConnections(t *testing.T) { + client, connections := New(t, http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + _, _ = io.WriteString(w, "ok") + })) + + const requests = 2 + for i := 0; i < requests; i++ { + req, err := http.NewRequest(http.MethodGet, "http://example.test/", nil) + if err != nil { + t.Fatal(err) + } + // Request.Close makes the transport reject connection reuse before it + // handles the response body. This forces a new connection so the + // counter proves it can count two of them. + req.Close = true + resp, err := client.Do(req) + if err != nil { + t.Fatal(err) + } + httpcon.DrainAndClose(resp.Body) + } + + if got := connections.Count(); got != requests { + t.Fatalf("accepted %d HTTP/1 connections, want %d", got, requests) + } +}