diff --git a/cmd/ateapi/internal/controlapi/dialer.go b/cmd/ateapi/internal/controlapi/dialer.go index 89e8e37a6..a0f696f4d 100644 --- a/cmd/ateapi/internal/controlapi/dialer.go +++ b/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), } } diff --git a/cmd/ateapi/internal/controlapi/dialer_test.go b/cmd/ateapi/internal/controlapi/dialer_test.go new file mode 100644 index 000000000..e7d3f804a --- /dev/null +++ b/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") + } +} diff --git a/cmd/atelet/main.go b/cmd/atelet/main.go index 33a0d5331..069a3de30 100644 --- a/cmd/atelet/main.go +++ b/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 } diff --git a/cmd/atelet/main_test.go b/cmd/atelet/main_test.go index ad0785f7b..6f4761fa9 100644 --- a/cmd/atelet/main_test.go +++ b/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