From 878251b73b66f062c0a5bbac722a78826d3ec37f Mon Sep 17 00:00:00 2001 From: Jared Ping <13906915+JaredP94@users.noreply.github.com> Date: Mon, 2 Nov 2026 16:32:00 +0200 Subject: [PATCH] fix(services): reuse connections in engine health probes The broker's Ollama, LM Studio and oMLX health checks and the engine manager's readiness probe closed each response body without reading it. Go's HTTP/1 transport only returns a connection to the pool once its body reaches EOF, so every poll opened a new connection to the engine. Drain a bounded amount of each body before closing it (shared httpcon.DrainAndClose), so one connection serves every poll. Ported from NVIDIA/Personal-AI-Router@5ae90de (NVIDIA/Personal-AI-Router#37), which fixed the Ollama path; extended to this fork's LM Studio and oMLX probes. Co-authored-by: Psych0h3ad <41975091+Psych0h3ad@users.noreply.github.com> Signed-off-by: Jared Ping <13906915+JaredP94@users.noreply.github.com> --- services/nvpair-engine-manager/lifecycle.go | 4 +- services/nvpair-ui-broker/advertiser.go | 9 ++-- services/shared/httpcon/httpcon.go | 19 +++++++++ services/shared/httpcon/httpcon_test.go | 47 +++++++++++++++++++++ 4 files changed, 74 insertions(+), 5 deletions(-) create mode 100644 services/shared/httpcon/httpcon.go create mode 100644 services/shared/httpcon/httpcon_test.go diff --git a/services/nvpair-engine-manager/lifecycle.go b/services/nvpair-engine-manager/lifecycle.go index 6ee8d7a21..ed2b48fa5 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 5c2296c17..3277403cf 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 000000000..7e928d90b --- /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 000000000..bb0d7c6f4 --- /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) + } + } +}