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