fix(distributed): drop a departed node's cached HTTP clients

A frontend keeps two caches of one http.Client per worker: the control
client's, built on the first verb issued to a node, and the HTTP file stager's,
built on the first file staged to it. Both are keyed by node ID and neither was
ever pruned. Their own comments said so and named what a fix would need, a
signal that a node has left, which did not exist when they were written and
does now.

The map slot is the smaller half. Each entry holds an http.Transport whose idle
connections are streams on that worker's tunnel, kept until IdleConnTimeout
even after the tunnel is gone, so ForgetNode closes them rather than leaving
them to the collector. Both are registered on the deployment's one departure
notifier, as two subscribers and not one: a node can be in either cache without
being in the other, and a single hook would say only that some client was kept.

ForgetNode is on the FileStager interface rather than on the one implementation
that has state to drop, so a stager that grows a per-node map later cannot be
added without answering the question, and so registerDepartureEvictions can
take a FileStager and still fail to compile if the registration is deleted. The
S3 stager's is a documented no-op that deliberately does not forward to the
control client, which registers itself: forwarding would evict a cache it does
not own, twice per departure, and the second drop would not appear in the
subscriber names the wiring spec reads.

Assisted-by: Claude Opus 5 [claude-code]
Signed-off-by: Ettore Di Giacinto <mudler@localai.io>

core/application/distributed.go

@@ -673,7 +673,7 @@ func initDistributed(cfg *config.ApplicationConfig, authDB *gorm.DB, configLoade
 
 	// Every per-node cache a departure leaves stale, onto the one notification
 	// point, after the router that owns two of them exists.
-	if err := registerDepartureEvictions(departures, prefixDrop, router, galleryProgress); err != nil {
+	if err := registerDepartureEvictions(departures, prefixDrop, router, galleryProgress, controlClient, fileStager); err != nil {
 		return nil, err
 	}
 
@@ -755,6 +755,8 @@ const (
 	departureProbeCache     = "probe-cache"
 	departureStagingTracker = "staging-tracker"
 	departureGalleryNodes   = "gallery-node-progress"
+	departureControlClients = "control-http-clients"
+	departureStagerClients  = "file-stager-http-clients"
 )
 
 // registerDepartureEvictions registers every per-node cache that a node's
@@ -778,7 +780,7 @@ const (
 // *prefixcache.Sync for that decision to be safe, since a nil provider inside
 // an interface would compare non-nil here and dereference on the first
 // departure.
-func registerDepartureEvictions(departures *nodes.DepartureNotifier, prefix *prefixcache.Sync, router *nodes.SmartRouter, gallery nodeProgressDropper) error {
+func registerDepartureEvictions(departures *nodes.DepartureNotifier, prefix *prefixcache.Sync, router *nodes.SmartRouter, gallery nodeProgressDropper, control *nodes.ControlClient, stager nodes.FileStager) error {
 	if departures == nil {
 		return fmt.Errorf("wiring departure evictions: no departure notifier, so a departed node would keep every per-node cache entry it has for the life of the process")
 	}
@@ -788,6 +790,12 @@ func registerDepartureEvictions(departures *nodes.DepartureNotifier, prefix *pre
 	if gallery == nil {
 		return fmt.Errorf("wiring departure evictions: no gallery service, so a departed node would stay in every open operation's per-node breakdown")
 	}
+	if control == nil {
+		return fmt.Errorf("wiring departure evictions: no control client, so a departed node would keep its cached HTTP client and that client's idle streams on a tunnel that is gone")
+	}
+	if stager == nil {
+		return fmt.Errorf("wiring departure evictions: no file stager, so a departed node would keep the cached HTTP client its transfers ran on")
+	}
 	// S1. Inside a nil check and not inside the prefix-cache-enabled block, so
 	// that "the disabled deployment registers nothing" is a fact a spec can
 	// hold rather than a property of where a line was written.
@@ -809,6 +817,22 @@ func registerDepartureEvictions(departures *nodes.DepartureNotifier, prefix *pre
 	departures.OnDeparture(departureGalleryNodes, func(node nodes.DepartedNode) {
 		gallery.DropNodeProgress(node.ID)
 	})
+	// S5 and S6, the two per-node http.Client caches. Two registrations again,
+	// because they are two caches with two owners: the control client's entry
+	// is built on the first verb issued to a node and the stager's on the first
+	// file staged to it, so a node can be in either without being in the other,
+	// and one hook doing both would say only that some client was kept.
+	//
+	// Both are keyed by node ID and both are DROPPED rather than emptied. A
+	// worker that comes back builds a fresh client on its next verb, over
+	// whatever tunnel it has by then; keeping the old one would keep a
+	// transport whose idle streams belong to a session that has ended.
+	departures.OnDeparture(departureControlClients, func(node nodes.DepartedNode) {
+		control.ForgetNode(node.ID)
+	})
+	departures.OnDeparture(departureStagerClients, func(node nodes.DepartedNode) {
+		stager.ForgetNode(node.ID)
+	})
 	return nil
 }
 

core/application/distributed_test.go

@@ -222,6 +222,8 @@ var _ = Describe("wiring the per-node caches a departure evicts", func() {
 			departureProbeCache,
 			departureStagingTracker,
 			departureGalleryNodes,
+			departureControlClients,
+			departureStagerClients,
 		))
 	})
 
@@ -239,26 +241,54 @@ var _ = Describe("wiring the per-node caches a departure evicts", func() {
 			departureProbeCache,
 			departureStagingTracker,
 			departureGalleryNodes,
+			departureControlClients,
+			departureStagerClients,
 		))
 	})
 
 	It("refuses a deployment with no router, naming what its departed nodes would keep", func() {
-		err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nil, galleryop.NewGalleryService(&config.ApplicationConfig{}, nil))
+		err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nil, galleryop.NewGalleryService(&config.ApplicationConfig{}, nil), specControlClient(), specFileStager())
 
 		Expect(err).To(HaveOccurred())
 		Expect(err.Error()).To(ContainSubstring("probe-freshness"))
 	})
 
 	It("refuses a deployment with no gallery service", func() {
-		err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), nil)
+		err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), nil, specControlClient(), specFileStager())
 
 		Expect(err).To(HaveOccurred())
 		Expect(err.Error()).To(ContainSubstring("per-node breakdown"))
 	})
 
+	It("refuses a deployment with no control client, naming the streams a departed node would keep", func() {
+		err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), galleryop.NewGalleryService(&config.ApplicationConfig{}, nil), nil, specFileStager())
+
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("idle streams"))
+	})
+
+	It("refuses a deployment with no file stager", func() {
+		err := registerDepartureEvictions(nodes.NewDepartureNotifier(), nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), galleryop.NewGalleryService(&config.ApplicationConfig{}, nil), specControlClient(), nil)
+
+		Expect(err).To(HaveOccurred())
+		Expect(err.Error()).To(ContainSubstring("cached HTTP client"))
+	})
+
 	It("refuses a deployment with no notifier at all", func() {
-		err := registerDepartureEvictions(nil, nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), galleryop.NewGalleryService(&config.ApplicationConfig{}, nil))
+		err := registerDepartureEvictions(nil, nil, nodes.NewSmartRouter(nil, nodes.SmartRouterOptions{}), galleryop.NewGalleryService(&config.ApplicationConfig{}, nil), specControlClient(), specFileStager())
 
 		Expect(err).To(HaveOccurred())
 	})
 })
+
+// specControlClient and specFileStager build the two per-node client caches a
+// registration needs. Neither is dialled: what the refusal specs assert is that
+// the wiring refuses a MISSING one, and what the registration specs assert is
+// that a subscriber was registered for it.
+func specControlClient() *nodes.ControlClient {
+	return nodes.NewControlClient(nil, "")
+}
+
+func specFileStager() nodes.FileStager {
+	return nodes.NewHTTPFileStager(func(string) (string, error) { return "", nil }, "", nil)
+}

core/services/nodes/control_client.go

@@ -56,20 +56,31 @@ type ControlClient struct {
 	// row reuses the tunnel stream its transport already holds instead of
 	// opening a new one.
 	//
-	// Entries are never pruned, and that is judged rather than overlooked: the
-	// map is bounded by the number of distinct workers this frontend has ever
-	// commanded, which is bounded by the fleet, and a departed worker's entry
-	// holds a map slot plus a transport whose idle connections IdleConnTimeout
-	// reclaims. It is the same shape, and would need the same missing
-	// node-departure signal to fix, as HTTPFileStager.clients.
+	// Entries are dropped by ForgetNode, on the deployment's one departure
+	// notification, the same way HTTPFileStager.clients is. Without that the
+	// map grew once per distinct worker for the life of the process, which is
+	// bounded by the fleet only in a deployment whose nodes never change
+	// identity.
 	clientsMu sync.Mutex
-	clients   map[string]*http.Client
+	clients   map[string]*nodeHTTPClient
+}
+
+// nodeHTTPClient is one cached per-node client together with the transport it
+// was built on.
+//
+// The transport is kept beside the client rather than read back off
+// http.Client.Transport, because what is stored there is whatever the shared
+// httpclient constructor decided to wrap it in, and a per-node cache that
+// cannot close its own idle connections has only half a ForgetNode.
+type nodeHTTPClient struct {
+	client    *http.Client
+	transport *http.Transport
 }
 
 // NewControlClient returns the control client for the workers dialFor can
 // reach, authenticating with the deployment's registration token.
 func NewControlClient(dialFor WorkerNetDialerFor, token string) *ControlClient {
-	return &ControlClient{dialFor: dialFor, token: token, clients: map[string]*http.Client{}}
+	return &ControlClient{dialFor: dialFor, token: token, clients: map[string]*nodeHTTPClient{}}
 }
 
 // clientFor returns the HTTP client that reaches one worker, building it on
@@ -96,7 +107,7 @@ func (c *ControlClient) clientFor(nodeID string) (*http.Client, error) {
 	c.clientsMu.Lock()
 	defer c.clientsMu.Unlock()
 	if cl, ok := c.clients[nodeID]; ok {
-		return cl, nil
+		return cl.client, nil
 	}
 	dial := c.dialFor(nodeID)
 	if dial == nil {
@@ -111,10 +122,30 @@ func (c *ControlClient) clientFor(nodeID string) (*http.Client, error) {
 		ExpectContinueTimeout: 1 * time.Second,
 	}
 	cl := httpclient.New(httpclient.WithTransport(transport))
-	c.clients[nodeID] = cl
+	c.clients[nodeID] = &nodeHTTPClient{client: cl, transport: transport}
 	return cl, nil
 }
 
+// ForgetNode drops the cached client for a departed node and closes the idle
+// tunnel streams its transport is holding. See HTTPFileStager.ForgetNode: the
+// two are one rule applied to the two per-node client caches this frontend
+// keeps, and both are registered on the departure notifier.
+//
+// Safe on a nil receiver, because the notifier's subscribers are registered at
+// wiring time and a deployment can reach it with no control client.
+func (c *ControlClient) ForgetNode(nodeID string) {
+	if c == nil {
+		return
+	}
+	c.clientsMu.Lock()
+	entry, ok := c.clients[nodeID]
+	delete(c.clients, nodeID)
+	c.clientsMu.Unlock()
+	if ok {
+		entry.transport.CloseIdleConnections()
+	}
+}
+
 // Call issues one control RPC and decodes its reply. reply may be nil for the
 // verbs that answer 204.
 //

core/services/nodes/file_stager.go

@@ -35,4 +35,15 @@ type FileStager interface {
 	// ListRemoteDir returns relative file paths within a directory on the remote node.
 	// keyPrefix is a storage-style key prefix (e.g. "models/mymodel").
 	ListRemoteDir(ctx context.Context, nodeID, keyPrefix string) ([]string, error)
+
+	// ForgetNode drops whatever this stager holds for one node, and is called
+	// when the deployment decides that node has departed.
+	//
+	// On the INTERFACE rather than on the one implementation that has state to
+	// drop, so that a stager which grows a per-node map later cannot be added
+	// without answering this question, and so that the wiring in
+	// core/application can name a FileStager and still fail to compile if the
+	// registration is deleted. An implementation with nothing per-node is a
+	// documented no-op.
+	ForgetNode(nodeID string)
 }

core/services/nodes/file_stager_http.go

@@ -42,16 +42,13 @@ type HTTPFileStager struct {
 	// connection pool: a client built per request would open a fresh tunnel
 	// stream for every chunk of a multi-gigabyte upload.
 	//
-	// Entries are never pruned, and that is judged acceptable rather than
-	// overlooked. The map is bounded by the number of distinct workers this
-	// frontend has ever staged to, which is bounded by the fleet; each entry is
-	// a transport whose idle connections the 90s IdleConnTimeout above reclaims,
-	// so a departed worker's entry holds a map slot and nothing else. It is the
-	// same shape as PeerPool.links and would need the same thing to fix
-	// properly: a signal that a node has left, which the deregistration path
-	// does not publish today.
+	// Entries are dropped by ForgetNode, on the deployment's one departure
+	// notification. Without that the map grew once per distinct worker for the
+	// life of the process, which is bounded by the fleet only in a deployment
+	// whose nodes never change identity; a Kubernetes worker pool that
+	// re-registers under a new id on every rollout is not one.
 	clientsMu       sync.Mutex
-	clients         map[string]*http.Client
+	clients         map[string]*nodeHTTPClient
 	responseTimeout time.Duration // timeout waiting for server response after upload
 	maxRetries      int           // number of retry attempts for transient failures
 }
@@ -79,7 +76,7 @@ func NewHTTPFileStager(httpAddrFor func(nodeID string) (string, error), token st
 		httpAddrFor:     httpAddrFor,
 		token:           token,
 		dialFor:         dialFor,
-		clients:         map[string]*http.Client{},
+		clients:         map[string]*nodeHTTPClient{},
 		responseTimeout: responseTimeout,
 		maxRetries:      maxRetries,
 	}
@@ -110,7 +107,7 @@ func (h *HTTPFileStager) clientFor(nodeID string) (*http.Client, error) {
 	h.clientsMu.Lock()
 	defer h.clientsMu.Unlock()
 	if c, ok := h.clients[nodeID]; ok {
-		return c, nil
+		return c.client, nil
 	}
 	dial := h.dialFor(nodeID)
 	if dial == nil {
@@ -127,10 +124,27 @@ func (h *HTTPFileStager) clientFor(nodeID string) (*http.Client, error) {
 		ReadBufferSize:        256 << 10, // 256 KB
 	}
 	c := httpclient.New(httpclient.WithTransport(transport))
-	h.clients[nodeID] = c
+	h.clients[nodeID] = &nodeHTTPClient{client: c, transport: transport}
 	return c, nil
 }
 
+// ForgetNode drops the cached client for a departed node and closes the idle
+// tunnel streams its transport is holding.
+//
+// Closing and not merely deleting: the map slot is the smaller half. A
+// transport left only to the garbage collector keeps its idle connections
+// until IdleConnTimeout, and each of those is a stream on a tunnel that is
+// already gone.
+func (h *HTTPFileStager) ForgetNode(nodeID string) {
+	h.clientsMu.Lock()
+	entry, ok := h.clients[nodeID]
+	delete(h.clients, nodeID)
+	h.clientsMu.Unlock()
+	if ok {
+		entry.transport.CloseIdleConnections()
+	}
+}
+
 func (h *HTTPFileStager) EnsureRemote(ctx context.Context, nodeID, localPath, key string) (string, error) {
 	xlog.Debug("Staging file to remote node via HTTP", "node", nodeID, "localPath", localPath, "key", key)
 

core/services/nodes/file_stager_s3.go

@@ -29,6 +29,17 @@ func NewS3FileStager(fm *storage.FileManager, control *ControlClient) *S3FileSta
 	return &S3FileStager{fm: fm, control: control}
 }
 
+// ForgetNode has nothing of its own to drop: this stager keeps no per-node
+// state at all. Bytes travel through the object store, and the only per-node
+// client involved is the ControlClient's, which the deployment shares with
+// every other control verb and registers on the departure notifier itself.
+//
+// Deliberately NOT forwarded to that client. Dropping it from here as well
+// would close a cache this stager does not own, on a departure it would then
+// be reacting to twice, and the second drop would be invisible in the
+// subscriber names the wiring spec reads.
+func (s *S3FileStager) ForgetNode(string) {}
+
 // The two budgets a file-staging RPC gets. They are the ones the NATS
 // request-reply timeouts carried, kept verbatim: a transfer verb waits out a
 // multi-gigabyte copy, and a metadata verb does not.

core/services/nodes/file_staging_3d_test.go

@@ -44,3 +44,6 @@ var _ = Describe("FileStagingClient 3D output", func() {
 		Expect(backend.request.Dst).To(Equal("/remote/tmp"))
 	})
 })
+
+// ForgetNode drops per-node state, which this double keeps none of.
+func (*failingFetchStager) ForgetNode(string) {}

core/services/nodes/per_node_clients_test.go

@@ -0,0 +1,135 @@
+package nodes
+
+import (
+	"context"
+	"net"
+
+	. "github.com/onsi/ginkgo/v2"
+	. "github.com/onsi/gomega"
+)
+
+// A frontend keeps two caches of one *http.Client per worker: the control
+// client's, built on the first verb issued to a node, and the HTTP file
+// stager's, built on the first file staged to it. Both are keyed by node ID and
+// neither used to be pruned, so each grew once per distinct worker for the life
+// of the process. A worker pool that re-registers under a new id on every
+// rollout is not bounded by the fleet.
+//
+// Each cached entry is more than a map slot: it holds an http.Transport whose
+// idle connections are streams on that worker's tunnel, kept until
+// IdleConnTimeout even after the tunnel has gone.
+var _ = Describe("per-node HTTP client caches", func() {
+	// aDialer stands in for a worker's tunnel dialler. It is never called: what
+	// these specs assert is which client the cache hands back, not what that
+	// client can reach.
+	aDialer := func(string) func(ctx context.Context, network, addr string) (net.Conn, error) {
+		return func(context.Context, string, string) (net.Conn, error) { return nil, nil }
+	}
+
+	Describe("ControlClient", func() {
+		var c *ControlClient
+
+		BeforeEach(func() { c = NewControlClient(aDialer, "token") })
+
+		It("reuses one client per node until that node departs", func() {
+			first, err := c.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+			again, err := c.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(again).To(BeIdenticalTo(first),
+				"a verb issued twice must reuse the tunnel stream the transport already holds")
+		})
+
+		It("drops a departed node's client rather than holding it for the life of the process", func() {
+			first, err := c.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+
+			c.ForgetNode("node-1")
+			Expect(c.clients).NotTo(HaveKey("node-1"))
+
+			rebuilt, err := c.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(rebuilt).NotTo(BeIdenticalTo(first),
+				"a worker that comes back gets a client over whatever tunnel it has now")
+		})
+
+		It("leaves the other nodes' clients alone", func() {
+			kept, err := c.clientFor("node-2")
+			Expect(err).NotTo(HaveOccurred())
+			_, err = c.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+
+			c.ForgetNode("node-1")
+
+			again, err := c.clientFor("node-2")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(again).To(BeIdenticalTo(kept))
+		})
+
+		It("is a no-op for a node it never built a client for", func() {
+			Expect(func() { c.ForgetNode("never-seen") }).NotTo(Panic())
+			Expect(c.clients).To(BeEmpty())
+		})
+	})
+
+	Describe("HTTPFileStager", func() {
+		var h *HTTPFileStager
+
+		BeforeEach(func() {
+			h = NewHTTPFileStager(func(string) (string, error) { return "worker:1", nil }, "token", aDialer)
+		})
+
+		It("reuses one client per node until that node departs", func() {
+			first, err := h.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+			again, err := h.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(again).To(BeIdenticalTo(first),
+				"a client built per request would open a fresh tunnel stream for every chunk of a multi-gigabyte upload")
+		})
+
+		It("drops a departed node's client rather than holding it for the life of the process", func() {
+			first, err := h.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+
+			h.ForgetNode("node-1")
+			Expect(h.clients).NotTo(HaveKey("node-1"))
+
+			rebuilt, err := h.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(rebuilt).NotTo(BeIdenticalTo(first))
+		})
+
+		It("leaves the other nodes' clients alone", func() {
+			kept, err := h.clientFor("node-2")
+			Expect(err).NotTo(HaveOccurred())
+			_, err = h.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+
+			h.ForgetNode("node-1")
+
+			again, err := h.clientFor("node-2")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(again).To(BeIdenticalTo(kept))
+		})
+	})
+
+	Describe("S3FileStager", func() {
+		It("does not drop the control client's entry on behalf of a stager that has no state", func() {
+			// The object-store stager keeps nothing per node; its only per-node
+			// client belongs to the ControlClient, which the deployment shares
+			// with every other control verb and registers separately. Dropping
+			// it from here as well would evict a cache this stager does not own
+			// and would do it twice per departure, invisibly.
+			control := NewControlClient(aDialer, "token")
+			kept, err := control.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+
+			NewS3FileStager(nil, control).ForgetNode("node-1")
+
+			again, err := control.clientFor("node-1")
+			Expect(err).NotTo(HaveOccurred())
+			Expect(again).To(BeIdenticalTo(kept))
+		})
+	})
+})

core/services/nodes/router_staging_context_test.go

@@ -78,3 +78,6 @@ var _ = Describe("Route cold-load staging context", func() {
 			"staging context must survive cancellation of the triggering request")
 	})
 })
+
+// ForgetNode drops per-node state, which this double keeps none of.
+func (*cancelOnStageStager) ForgetNode(string) {}

core/services/nodes/router_staging_deadline_test.go

@@ -219,3 +219,6 @@ var _ = Describe("cold-load staging deadline", func() {
 		Expect(time.Until(deadline)).To(BeNumerically("~", 3*time.Hour, time.Minute))
 	})
 })
+
+// ForgetNode drops per-node state, which this double keeps none of.
+func (*progressingStager) ForgetNode(string) {}

core/services/nodes/router_test.go

@@ -1712,3 +1712,6 @@ var _ = Describe("SmartRouter prefix-cache routing", func() {
 		})
 	})
 })
+
+// ForgetNode drops per-node state, which this double keeps none of.
+func (*fakeFileStager) ForgetNode(string) {}