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..44e13a3a --- /dev/null +++ b/services/nvpair-engine-manager/health_connections_test.go @@ -0,0 +1,109 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "context" + "io" + "net/http" + "strings" + "sync/atomic" + "testing" + "time" + + "nvpair-shared/httpcon/testclient" +) + +func TestProbeHTTPReusesConnections(t *testing.T) { + test := 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.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.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) { + 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 } diff --git a/services/nvpair-engine-manager/lifecycle.go b/services/nvpair-engine-manager/lifecycle.go index cea7aa16..aa599903 100644 --- a/services/nvpair-engine-manager/lifecycle.go +++ b/services/nvpair-engine-manager/lifecycle.go @@ -19,6 +19,8 @@ import ( "sync" "syscall" "time" + + "nvpair-shared/httpcon" ) const ( @@ -871,7 +873,8 @@ func (e *Executor) probe(ctx context.Context, p *Probe, port int) bool { if err != nil { return false } - 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 05d85fdf..d48097e7 100644 --- a/services/nvpair-ui-broker/advertiser.go +++ b/services/nvpair-ui-broker/advertiser.go @@ -10,6 +10,7 @@ import ( "net/http" "time" + "nvpair-shared/httpcon" "nvpair-shared/noderec" ) @@ -274,6 +275,6 @@ func checkEngineHealth(profile engineProxyProfile, client *http.Client, port int if err != nil { return false } - 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 new file mode 100644 index 00000000..74e7231a --- /dev/null +++ b/services/nvpair-ui-broker/health_connections_test.go @@ -0,0 +1,54 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package main + +import ( + "io" + "net/http" + "sync/atomic" + "testing" + + "nvpair-shared/httpcon/testclient" +) + +func TestHealthChecksReuseConnections(t *testing.T) { + 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) + } + } + if got := connections.Count(); got != 1 { + t.Fatalf("accepted %d HTTP/1 connections for %d polls, want 1", got, rounds) + } + }) + } + + testFraming("content-length", false) + testFraming("chunked", true) + }) + } + + test("ollama", "/", ollamaProxyProfile) + test("lmstudio", "/v1/models", lmstudioProxyProfile) +} 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) + } +}