Files
substrate/cmd/ateapi/internal/controlapi/workload_spec.go
T
Dmitry Berkovich 125180ea65 feat: implement container readiness probes with custom HTTP endpoint … (#330)
Extends ActorTemplate.Container spec to allow defining a readyz probe. 
* ateoms (both gVisor & microVM) are updated and wait for readyz signal
if it was enabled on container when creating a new container or resume
from a snapshot.
* actorTemplates controller has been updated not to wait 20 sec during
golden snapshot creation, if readyz is enabled on actor. It assumes,
atelet call returns when all containers passed readyz.

Tested manually both counters: 
* gVisor 
* microvm 

Fixes #316 

- [x] e2e tests pass 
- [X] Appropriate changes to documentation are included in the PR
2026-06-26 18:59:56 -07:00

216 lines
6.1 KiB
Go

// 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 (
"context"
"fmt"
"sync"
"time"
"github.com/agent-substrate/substrate/internal/proto/ateletpb"
atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
corev1 "k8s.io/api/core/v1"
apierrors "k8s.io/apimachinery/pkg/api/errors"
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
"k8s.io/client-go/kubernetes"
)
const envSecretCacheTTL = 30 * time.Second
func workloadSpecFromActorTemplate(ctx context.Context, kubeClient kubernetes.Interface, secretCache *envSecretCache, actorTemplate *atev1alpha1.ActorTemplate) (*ateletpb.WorkloadSpec, error) {
workloadSpec := &ateletpb.WorkloadSpec{
PauseImage: actorTemplate.Spec.PauseImage,
}
resolver := envResolver{
kubeClient: kubeClient,
namespace: actorTemplate.Namespace,
cache: secretCache,
}
for _, ctr := range actorTemplate.Spec.Containers {
ateletCtr := &ateletpb.Container{
Name: ctr.Name,
Image: ctr.Image,
Command: ctr.Command,
Readyz: toAteletReadyz(ctr.Readyz),
}
for _, env := range ctr.Env {
ateletEnv, err := resolver.resolve(ctx, ctr.Name, env)
if err != nil {
return nil, err
}
if ateletEnv != nil {
ateletCtr.Env = append(ateletCtr.Env, ateletEnv)
}
}
workloadSpec.Containers = append(workloadSpec.Containers, ateletCtr)
}
return workloadSpec, nil
}
// toAteletReadyz projects the CRD readyz field onto the ateletpb wire type.
// Returns nil when the source is nil so containers without a probe stay
// unchanged on the wire.
func toAteletReadyz(in *atev1alpha1.ContainerReadyz) *ateletpb.Readyz {
if in == nil {
return nil
}
out := &ateletpb.Readyz{}
if in.HTTPGet != nil {
out.HttpGet = &ateletpb.HTTPGetAction{
Path: in.HTTPGet.Path,
Port: in.HTTPGet.Port,
}
}
return out
}
type envResolver struct {
kubeClient kubernetes.Interface
namespace string
cache *envSecretCache
}
func (r *envResolver) resolve(ctx context.Context, containerName string, env atev1alpha1.EnvVar) (*ateletpb.EnvEntry, error) {
envID := fmt.Sprintf("container %q env %q", containerName, env.Name)
switch {
case env.Value != nil:
return &ateletpb.EnvEntry{
Name: env.Name,
Value: *env.Value,
}, nil
case env.ValueFrom != nil:
value, include, err := r.resolveValueFrom(ctx, envID, env.ValueFrom)
if err != nil {
return nil, err
}
if !include {
return nil, nil
}
return &ateletpb.EnvEntry{
Name: env.Name,
Value: value,
}, nil
}
return nil, status.Errorf(codes.FailedPrecondition, "%s has unknown value source", envID)
}
func (r *envResolver) resolveValueFrom(ctx context.Context, envID string, valueFrom *atev1alpha1.EnvVarSource) (string, bool, error) {
if ref := valueFrom.SecretKeyRef; ref != nil {
return r.resolveSecretKeyRef(ctx, envID, ref)
}
return "", false, status.Errorf(codes.FailedPrecondition, "%s uses unsupported valueFrom source; only secretKeyRef is supported", envID)
}
func (r *envResolver) resolveSecretKeyRef(ctx context.Context, envID string, ref *atev1alpha1.SecretKeySelector) (string, bool, error) {
if r.kubeClient == nil {
return "", false, status.Errorf(codes.FailedPrecondition, "%s cannot resolve secretKeyRef because Kubernetes client is unavailable", envID)
}
secret, err := r.secret(ctx, ref.Name)
if err != nil {
if apierrors.IsNotFound(err) {
if isOptional(ref.Optional) {
return "", false, nil
}
return "", false, status.Errorf(codes.FailedPrecondition, "%s references missing secret %s/%s", envID, r.namespace, ref.Name)
}
return "", false, status.Errorf(codes.Internal, "while resolving %s secretKeyRef %s/%s: %v", envID, r.namespace, ref.Name, err)
}
value, ok := secret.Data[ref.Key]
if !ok {
if isOptional(ref.Optional) {
return "", false, nil
}
return "", false, status.Errorf(codes.FailedPrecondition, "%s references missing key %q in secret %s/%s", envID, ref.Key, r.namespace, ref.Name)
}
return string(value), true, nil
}
func (r *envResolver) secret(ctx context.Context, name string) (*corev1.Secret, error) {
if r.cache != nil {
return r.cache.get(ctx, r.kubeClient, r.namespace, name)
}
return r.kubeClient.CoreV1().Secrets(r.namespace).Get(ctx, name, metav1.GetOptions{})
}
type envSecretCache struct {
mu sync.RWMutex
ttl time.Duration
entries map[envSecretCacheKey]envSecretCacheEntry
}
type envSecretCacheKey struct {
namespace string
name string
}
type envSecretCacheEntry struct {
secret *corev1.Secret
expiresAt time.Time
}
func newEnvSecretCache(ttl time.Duration) *envSecretCache {
return &envSecretCache{
ttl: ttl,
entries: map[envSecretCacheKey]envSecretCacheEntry{},
}
}
func (c *envSecretCache) get(ctx context.Context, kubeClient kubernetes.Interface, namespace, name string) (*corev1.Secret, error) {
key := envSecretCacheKey{
namespace: namespace,
name: name,
}
now := time.Now()
c.mu.RLock()
entry, ok := c.entries[key]
if ok && now.Before(entry.expiresAt) {
secret := entry.secret.DeepCopy()
c.mu.RUnlock()
return secret, nil
}
c.mu.RUnlock()
// TODO: Make refresh smarter if this pattern sticks, for example by
// refreshing asynchronously or watching referenced Secrets.
secret, err := kubeClient.CoreV1().Secrets(namespace).Get(ctx, name, metav1.GetOptions{})
if err != nil {
return nil, err
}
secret = secret.DeepCopy()
c.mu.Lock()
c.entries[key] = envSecretCacheEntry{
secret: secret,
expiresAt: time.Now().Add(c.ttl),
}
c.mu.Unlock()
return secret.DeepCopy(), nil
}
func isOptional(optional *bool) bool {
return optional != nil && *optional
}