fix: both atelet and ate-apiserver leak gRPC connections

cmd/ateapi/internal/controlapi/dialer.go

@@ -17,6 +17,7 @@ package controlapi
 import (
 	"errors"
 	"fmt"
+	"log/slog"
 
 	"go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc"
 	"google.golang.org/grpc"
@@ -28,6 +29,17 @@ import (
 
 var ErrWorkerPodNotFound = errors.New("worker pod not found")
 
+var onEvict lru.EvictionFunc = func(key lru.Key, value any) {
+	// Close connection when evicting from cache.
+	conn, ok := value.(*grpc.ClientConn)
+	if ok {
+		err := conn.Close()
+		if err != nil {
+			slog.Debug("Failed to close evicted connection", slog.Any("error", err))
+		}
+	}
+}
+
 // AteletDialer handles gRPC connections to Atelet pods.
 type AteletDialer struct {
 	workerIndexer cache.Indexer
@@ -40,7 +52,7 @@ func NewAteletDialer(workerIndexer cache.Indexer, ateletIndexer cache.Indexer) *
 	return &AteletDialer{
 		workerIndexer: workerIndexer,
 		ateletIndexer: ateletIndexer,
-		ateletConns:   lru.New(1024),
+		ateletConns:   lru.NewWithEvictionFunc(1024, onEvict),
 	}
 }
 

cmd/ateapi/internal/controlapi/dialer_test.go

@@ -0,0 +1,92 @@
+// Copyright 2026 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package controlapi
+
+import (
+	"testing"
+
+	"google.golang.org/grpc/connectivity"
+	corev1 "k8s.io/api/core/v1"
+	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
+	"k8s.io/client-go/tools/cache"
+	"k8s.io/utils/lru"
+)
+
+func TestAteletDialerClosesEvictedConn(t *testing.T) {
+	workerIndexer := cache.NewIndexer(cache.MetaNamespaceKeyFunc, cache.Indexers{
+		byNamespaceAndName: func(obj any) ([]string, error) {
+			pod := obj.(*corev1.Pod)
+			return []string{pod.Namespace + "/" + pod.Name}, nil
+		},
+	})
+	ateletIndexer := cache.NewIndexer(cache.MetaNamespaceKeyFunc, cache.Indexers{
+		byNode: func(obj any) ([]string, error) {
+			return []string{obj.(*corev1.Pod).Spec.NodeName}, nil
+		},
+	})
+
+	for _, pod := range []*corev1.Pod{
+		{
+			ObjectMeta: metav1.ObjectMeta{Namespace: "worker-ns", Name: "worker-1"},
+			Spec:       corev1.PodSpec{NodeName: "node-1"},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{Namespace: "worker-ns", Name: "worker-2"},
+			Spec:       corev1.PodSpec{NodeName: "node-2"},
+		},
+	} {
+		if err := workerIndexer.Add(pod); err != nil {
+			t.Fatalf("adding worker pod: %v", err)
+		}
+	}
+
+	for _, pod := range []*corev1.Pod{
+		{
+			ObjectMeta: metav1.ObjectMeta{Namespace: "ate-system", Name: "atelet-1"},
+			Spec:       corev1.PodSpec{NodeName: "node-1"},
+			Status:     corev1.PodStatus{PodIPs: []corev1.PodIP{{IP: "127.0.0.1"}}},
+		},
+		{
+			ObjectMeta: metav1.ObjectMeta{Namespace: "ate-system", Name: "atelet-2"},
+			Spec:       corev1.PodSpec{NodeName: "node-2"},
+			Status:     corev1.PodStatus{PodIPs: []corev1.PodIP{{IP: "127.0.0.2"}}},
+		},
+	} {
+		if err := ateletIndexer.Add(pod); err != nil {
+			t.Fatalf("adding atelet pod: %v", err)
+		}
+	}
+
+	d := NewAteletDialer(workerIndexer, ateletIndexer)
+	d.ateletConns = lru.NewWithEvictionFunc(1, onEvict)
+
+	firstConn, err := d.DialForWorker("worker-ns", "worker-1")
+	if err != nil {
+		t.Fatalf("DialForWorker(worker-1): %v", err)
+	}
+
+	secondConn, err := d.DialForWorker("worker-ns", "worker-2")
+	if err != nil {
+		t.Fatalf("DialForWorker(worker-2): %v", err)
+	}
+	t.Cleanup(func() { _ = secondConn.Close() })
+
+	if got := firstConn.GetState(); got != connectivity.Shutdown {
+		t.Errorf("evicted connection state = %v, want %v", got, connectivity.Shutdown)
+	}
+	if got := secondConn.GetState(); got == connectivity.Shutdown {
+		t.Error("cached connection unexpectedly closed")
+	}
+}

cmd/atelet/main.go

@@ -69,6 +69,17 @@ var (
 	localhostRegistryReplacement = pflag.String("localhost-registry-replacement", "", "The replacement registry endpoint for localhost and/or loopback IP addresses, useful for local development. for example kind-registry:5000")
 
 	showVersion = pflag.Bool("version", false, "Print version and exit.")
+
+	onEvict lru.EvictionFunc = func(key lru.Key, value any) {
+		// Close connection when evicting from cache.
+		conn, ok := value.(*grpc.ClientConn)
+		if ok {
+			err := conn.Close()
+			if err != nil {
+				slog.Debug("Failed to close evicted connection", slog.Any("error", err))
+			}
+		}
+	}
 )
 
 func main() {
@@ -101,9 +112,7 @@ func main() {
 
 	go serverboot.StartMetricsServer(ctx, serverboot.MetricsServerOptions{Addr: *metricsListenAddr})
 
-	ateomDialer := &AteomDialer{
-		conns: lru.New(256),
-	}
+	ateomDialer := &AteomDialer{conns: lru.NewWithEvictionFunc(256, onEvict)}
 
 	var gcpRegistryAuthn authn.Authenticator
 	if *gcpAuthForImagePulls {
@@ -806,6 +815,7 @@ func toAteomReadyz(in *ateletpb.Readyz) *ateompb.Readyz {
 	return out
 }
 
+// AteomDialer handles gRPC connections to Ateom pods.
 type AteomDialer struct {
 	conns *lru.Cache
 }

cmd/atelet/main_test.go

@@ -33,8 +33,10 @@ import (
 	"github.com/agent-substrate/substrate/internal/proto/ateompb"
 	"github.com/google/go-cmp/cmp"
 	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/connectivity"
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/testing/protocmp"
+	"k8s.io/utils/lru"
 )
 
 func TestWriteFileAtomic(t *testing.T) {
@@ -255,6 +257,25 @@ func TestValidateRestoreRequest(t *testing.T) {
 	}
 }
 
+func TestAteomDialerClosesEvictedConn(t *testing.T) {
+	d := &AteomDialer{conns: lru.NewWithEvictionFunc(1, onEvict)}
+
+	first_conn, err := d.DialAteomPod(context.Background(), "pod-1")
+	if err != nil {
+		t.Fatalf("Failed to open connection for pod-1: %v", err)
+	}
+
+	// Should lead to the first connection being evicted, closing the connection.
+	_, err = d.DialAteomPod(context.Background(), "pod-2")
+	if err != nil {
+		t.Fatalf("Failed to open connection for pod-2: %v", err)
+	}
+
+	if state := first_conn.GetState(); state != connectivity.Shutdown {
+		t.Errorf("Expected connection state to be %v, got %v", connectivity.Shutdown, state)
+	}
+}
+
 // TestFetchAssetRejectsBadHash confirms fetchAsset validates the asset hash
 // before the cache-hit os.Stat/early-return, not merely "at some point". To
 // prove the ordering, it plants a real file at the exact path an invalid hash