Fix HTTP connection reuse in engine health probes

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 }

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

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
 }

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)
+}

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()
+}

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} }

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)
+	}
+}