diff --git a/services/nvpair-engine-manager/lifecycle.go b/services/nvpair-engine-manager/lifecycle.go index 6ee8d7a2..ed2b48fa 100644 --- a/services/nvpair-engine-manager/lifecycle.go +++ b/services/nvpair-engine-manager/lifecycle.go @@ -18,6 +18,8 @@ import ( "sync" "syscall" "time" + + "nvpair-shared/httpcon" ) const ( @@ -816,7 +818,7 @@ 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 5c2296c1..3277403c 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" ) @@ -271,7 +272,7 @@ func checkOllamaHealth(client *http.Client, port int) bool { if err != nil { return false } - resp.Body.Close() + httpcon.DrainAndClose(resp.Body) return resp.StatusCode == http.StatusOK } @@ -285,7 +286,7 @@ func checkLMStudioHealth(client *http.Client, port int) bool { if err != nil { return false } - resp.Body.Close() + httpcon.DrainAndClose(resp.Body) return resp.StatusCode == http.StatusOK } @@ -334,7 +335,7 @@ func (b *Broker) reconcileAdvertiseOMLX(client *http.Client) { func checkOMLXHealth(client *http.Client, port int) bool { resp, err := client.Get(fmt.Sprintf("http://localhost:%d/health", port)) if err == nil { - resp.Body.Close() + httpcon.DrainAndClose(resp.Body) if resp.StatusCode == http.StatusOK { return true } @@ -343,6 +344,6 @@ func checkOMLXHealth(client *http.Client, port int) bool { if err != nil { return false } - resp.Body.Close() + httpcon.DrainAndClose(resp.Body) return resp.StatusCode == http.StatusOK || resp.StatusCode == http.StatusUnauthorized } 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/httpcon_test.go b/services/shared/httpcon/httpcon_test.go new file mode 100644 index 00000000..bb0d7c6f --- /dev/null +++ b/services/shared/httpcon/httpcon_test.go @@ -0,0 +1,47 @@ +// SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +// SPDX-License-Identifier: Apache-2.0 + +package httpcon + +import ( + "io" + "net" + "net/http" + "net/http/httptest" + "strings" + "sync/atomic" + "testing" +) + +// Closing a body unread keeps its connection out of the pool, so a poller +// opens one connection per poll; draining it first lets one connection serve +// every poll. +func TestDrainAndCloseReusesTheConnection(t *testing.T) { + for _, chunked := range []bool{false, true} { + var conns atomic.Int32 + srv := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { + if chunked { + w.(http.Flusher).Flush() + } + _, _ = io.WriteString(w, strings.Repeat("x", 4096)) + })) + srv.Config.ConnState = func(_ net.Conn, state http.ConnState) { + if state == http.StateNew { + conns.Add(1) + } + } + srv.Start() + client := srv.Client() + for i := 0; i < 16; i++ { + resp, err := client.Get(srv.URL) + if err != nil { + t.Fatal(err) + } + DrainAndClose(resp.Body) + } + srv.Close() + if got := conns.Load(); got != 1 { + t.Errorf("chunked=%v: %d connections for 16 polls, want 1", chunked, got) + } + } +}