Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -88,6 +88,11 @@ ate-env manifest \
kubectl get pods -n ate-env
```

`ate-env-api` waits for a suspended environment to wake before forwarding a guest call, up to
`--wake-budget` (default 2m; cold image pulls and cross-region snapshot reads exceed the router's own
parking budget) and, once the control plane reports the environment RUNNING, up to
`--wake-running-budget` (default 20s). Both are flags on the API deployment.

### 2. Register the ActorTemplate in Substrate

Substrate manages ActorTemplates directly in its control plane rather than Kubernetes CRDs. Use `ate-env manifest template` to generate the Substrate ActorTemplate manifest:
Expand Down
16 changes: 10 additions & 6 deletions cmd/ate-env-api/main.go
Original file line number Diff line number Diff line change
Expand Up @@ -35,12 +35,14 @@ import (

func main() {
var (
listen = flag.String("listen", "0.0.0.0:7777", "address to serve the API on")
ateapi = flag.String("ateapi", ate.DefaultControlAddr, "address of the ateapi gRPC control plane")
atenet = flag.String("atenet", ate.DefaultRouterAddr, "address of the atenet HTTP router")
hostSuffix = flag.String("host-suffix", ate.DefaultHostSuffix, "atenet router host suffix for actor routing")
skipVerify = flag.Bool("skip-verify", true, "skip TLS certificate verification on the control plane connection")
tokenFile = flag.String("ateapi-token-file", "", "file with a bearer token for the control plane, e.g. a projected ServiceAccount token with the ateapi audience")
listen = flag.String("listen", "0.0.0.0:7777", "address to serve the API on")
ateapi = flag.String("ateapi", ate.DefaultControlAddr, "address of the ateapi gRPC control plane")
atenet = flag.String("atenet", ate.DefaultRouterAddr, "address of the atenet HTTP router")
hostSuffix = flag.String("host-suffix", ate.DefaultHostSuffix, "atenet router host suffix for actor routing")
skipVerify = flag.Bool("skip-verify", true, "skip TLS certificate verification on the control plane connection")
tokenFile = flag.String("ateapi-token-file", "", "file with a bearer token for the control plane, e.g. a projected ServiceAccount token with the ateapi audience")
wakeBudget = flag.Duration("wake-budget", apiservice.DefaultWakeBudget, "how long a guest call waits for a suspended environment to wake before failing; cold image pulls and cross-region snapshot reads exceed the router's own parking budget")
wakeRunningBudget = flag.Duration("wake-running-budget", apiservice.DefaultWakeRunningBudget, "how long a guest call waits when the environment is already RUNNING but the router still answers Unavailable (guest not yet listening vs. path down)")
)
flag.Parse()

Expand Down Expand Up @@ -68,6 +70,8 @@ func main() {
log.Printf("default atespace %q ready", ate.DefaultAtespace)

apisvc := apiservice.New(client, *atenet, *hostSuffix)
apisvc.WakeBudget = *wakeBudget
apisvc.WakeRunningBudget = *wakeRunningBudget
defer apisvc.Close()

grpcServer := grpc.NewServer()
Expand Down
32 changes: 24 additions & 8 deletions internal/apiservice/server.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@ import (
"errors"
"io"
"strings"
"time"

"github.com/agent-substrate/env/internal/ate"
ateenvv1alpha "github.com/agent-substrate/env/proto/ateenv/v1alpha"
Expand Down Expand Up @@ -55,6 +56,13 @@ type Server struct {
client *ate.Client
routerAddr string
hostSuffix string

// WakeBudget bounds how long a guest call waits for a suspended
// environment to wake; see ensureAwake. Zero means DefaultWakeBudget.
WakeBudget time.Duration
// WakeRunningBudget bounds the wait once the control plane reports the
// environment RUNNING; see ensureAwake. Zero means DefaultWakeRunningBudget.
WakeRunningBudget time.Duration
}

// New creates a new Server.
Expand All @@ -65,9 +73,11 @@ func New(client *ate.Client, routerAddr, hostSuffix string) *Server {
routerAddr = strings.TrimPrefix(routerAddr, "http://")
routerAddr = strings.TrimPrefix(routerAddr, "https://")
return &Server{
client: client,
routerAddr: routerAddr,
hostSuffix: hostSuffix,
client: client,
routerAddr: routerAddr,
hostSuffix: hostSuffix,
WakeBudget: DefaultWakeBudget,
WakeRunningBudget: DefaultWakeRunningBudget,
}
}

Expand Down Expand Up @@ -261,7 +271,7 @@ func (s *Server) processClient(ctx context.Context) (ateenvv1alpha.ProcessServic
if err != nil {
return nil, nil, nil, err
}
conn, err := s.guestConn(atespace, envID)
conn, err := s.guestConn(ctx, atespace, envID)
if err != nil {
return nil, nil, nil, err
}
Expand Down Expand Up @@ -371,7 +381,7 @@ func (s *Server) ReadFile(req *ateenvv1alpha.ReadFileRequest, stream grpc.Server
if err != nil {
return err
}
conn, err := s.guestConn(atespace, envID)
conn, err := s.guestConn(ctx, atespace, envID)
if err != nil {
return err
}
Expand Down Expand Up @@ -403,7 +413,7 @@ func (s *Server) WriteFile(stream grpc.ClientStreamingServer[ateenvv1alpha.Write
if err != nil {
return err
}
conn, err := s.guestConn(atespace, envID)
conn, err := s.guestConn(ctx, atespace, envID)
if err != nil {
return err
}
Expand Down Expand Up @@ -455,15 +465,17 @@ func envFromContext(ctx context.Context) (string, string, error) {
return ids[0], atespace, nil
}

// guestConn dials the environment's guest through the router. Calls on the
// guestConn dials the environment's guest through the router and waits,
// within WakeBudget, for the guest to answer, so a suspended environment is
// awake before the caller's request goes through (see ensureAwake). Calls on the
// returned connection take the caller's context as is, for its deadline and
// cancellation only: the interceptors below are the only writer of
// ate-target-actor, and no caller metadata is forwarded. Forwarding it once
// sent that header twice, which the Substrate router rejects
// (agent-substrate/substrate#2308 answers duplicate ate-target-actor headers
// with "invalid actor reference"), and it would let a client pass routing
// headers through the proxy. Nothing in the guest reads caller metadata.
func (s *Server) guestConn(atespace, id string) (*grpc.ClientConn, error) {
func (s *Server) guestConn(ctx context.Context, atespace, id string) (*grpc.ClientConn, error) {
if s.routerAddr == "" {
return nil, status.Error(codes.FailedPrecondition, "router address not configured for guest operations")
}
Expand All @@ -486,6 +498,10 @@ func (s *Server) guestConn(atespace, id string) (*grpc.ClientConn, error) {
if err != nil {
return nil, status.Errorf(codes.Internal, "dialing guest router for %s/%s: %v", atespace, id, err)
}
if err := s.ensureAwake(ctx, conn, atespace, id); err != nil {
_ = conn.Close()
return nil, err
}
return conn, nil
}

Expand Down
2 changes: 2 additions & 0 deletions internal/apiservice/server_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -42,6 +42,7 @@ type testEnv struct {
conn *grpc.ClientConn
router *fakerouter.Router
control *fakecontrol.Server
srv *apiservice.Server
}

func newTestEnv(t *testing.T) (ateenvv1alpha.EnvironmentServiceClient, *fakecontrol.Server) {
Expand Down Expand Up @@ -105,6 +106,7 @@ func newFullTestEnv(t *testing.T) *testEnv {
conn: conn,
router: router,
control: control,
srv: srv,
}
}

Expand Down
135 changes: 135 additions & 0 deletions internal/apiservice/wakeretry.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,135 @@
// 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 apiservice

import (
"context"
"log"
"strings"
"time"

ateenvv1alpha "github.com/agent-substrate/env/proto/ateenv/v1alpha"
"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
"google.golang.org/grpc"
"google.golang.org/grpc/codes"
"google.golang.org/grpc/status"
)

// DefaultWakeBudget bounds how long a guest call waits for an environment
// that the control plane reports as suspended or in transition (resuming,
// pausing, suspending) to wake before the error reaches the caller. The first
// request that reaches the router is what resumes a suspended actor, and the
// router holds that request only for its own parking budget, a few seconds. A
// wake that has to pull the environment's image onto the node or read its
// snapshot across a region takes longer than that: 18 s for a cold image pull
// and 7 s for a 1 GiB snapshot read across regions, measured on GKE.
const DefaultWakeBudget = 2 * time.Minute

// DefaultWakeRunningBudget bounds the wait when the control plane already
// reports the environment RUNNING but the router still answers Unavailable.
// That has two causes: the guest is not yet listening after a restore, which
// resolves in seconds, or the path to the worker is down, which does not.
// A short budget covers the first without turning the second into a
// two-minute hang per call.
const DefaultWakeRunningBudget = 20 * time.Second

const (
wakeRetryInitialBackoff = 200 * time.Millisecond
wakeRetryMaxBackoff = 2 * time.Second
)

// ensureAwake sends a trivial request to the environment's guest over conn
// and repeats it while the router answers Unavailable, which is what an
// expired parking budget and a guest that is not yet listening both look
// like. It returns when the guest answers anything else, when the actor has
// crashed, when ctx ends, or when a budget is spent: WakeBudget bounds the
// whole wait, and WakeRunningBudget bounds the part of it that starts with
// the first attempt that sees the actor RUNNING, so a long suspended phase
// does not eat the time the guest needs to start listening after the state
// flips. Because the probe is the first request through the router, it is
// also what wakes the actor; the calls that follow on conn find it up. The
// state is read once per attempt, and the backoff caps attempts at roughly
// one per two seconds, so a full two-minute wait costs about sixty
// control-plane reads.
func (s *Server) ensureAwake(ctx context.Context, conn *grpc.ClientConn, atespace, id string) error {
budget := s.WakeBudget
if budget <= 0 {
budget = DefaultWakeBudget
}
runningBudget := s.WakeRunningBudget
if runningBudget <= 0 {
runningBudget = DefaultWakeRunningBudget
}
start := time.Now()
var runningSince time.Time // first attempt that saw the actor RUNNING
probe := ateenvv1alpha.NewProcessServiceClient(conn)
backoff := wakeRetryInitialBackoff
for attempt := 1; ; attempt++ {
_, err := probe.GetProcess(ctx, &ateenvv1alpha.GetProcessRequest{ProcessId: "ate-env-wake-probe"})
if status.Code(err) != codes.Unavailable {
// Anything else, NotFound for the probe id included, came from
// the guest itself.
return nil
}
state, serr := s.actorState(ctx, atespace, id)
if serr != nil {
return serr
}
if state == ateapipb.ActorState_ACTOR_STATE_CRASHED {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

[nits] internal/apiservice/wakeretry.go:90: Return immediately when state is ACTOR_STATE_DELETING so guest calls do not retry against a deleting environment.

return status.Errorf(codes.FailedPrecondition, "environment %s/%s crashed while waking: %v", atespace, id, err)
}
allowed, since, which := budget, start, "wake budget"
if state == ateapipb.ActorState_ACTOR_STATE_RUNNING {
if runningSince.IsZero() {
runningSince = time.Now()
}
allowed, since, which = runningBudget, runningSince, "running budget"
} else {
runningSince = time.Time{}
}
remaining := allowed - time.Since(since)
if overall := budget - time.Since(start); overall < remaining {
remaining, allowed, which = overall, budget, "wake budget"
}
if remaining <= 0 {
return status.Errorf(codes.Unavailable, "environment %s/%s did not answer within the %s of %s while %s (%d attempts): %v",
atespace, id, which, allowed, shortState(state), attempt, err)
}
if attempt == 1 {
log.Printf("waiting for environment %s/%s to wake (state %s, %s %s)", atespace, id, shortState(state), which, allowed)
}
select {
case <-ctx.Done():
return status.FromContextError(ctx.Err()).Err()
case <-time.After(min(backoff, remaining)):
}
backoff = min(backoff*2, wakeRetryMaxBackoff)
}
}

// actorState reads the actor's state from the control plane. A deleted
// environment surfaces here as NotFound (through toGRPCError), which is the
// right answer to a caller still holding its handle.
func (s *Server) actorState(ctx context.Context, atespace, id string) (ateapipb.ActorState, error) {
actor, err := s.client.Get(ctx, atespace, id)
if err != nil {
return ateapipb.ActorState_ACTOR_STATE_UNSPECIFIED, toGRPCError(err)
}
return actor.GetStatus().GetState(), nil
}

func shortState(st ateapipb.ActorState) string {
return strings.TrimPrefix(st.String(), "ACTOR_STATE_")
}
Loading
Loading