mirror of
https://github.com/agent-substrate/substrate.git
synced 2026-11-11 10:23:17 +08:00
# Implement Pagination and Pipelining for ListActors ## Overview **Fixes:** #153 This PR introduces robust, memory-safe pagination to the `ListActors` operation and optimizes the Valkey retrieval layer. Previously, `ListActors` performed an unbounded parallel scatter-gather across all Valkey cluster shards, pulling all actors into memory at once. At scale (millions/billions of actors), this would cause OOM errors and severe latency. This update replaces the concurrent scatter-gather with a deterministic, sequential shard scanning approach, adhering to [AIP-158](https://google.aip.dev/158) pagination standards. ## Changes Included - **gRPC API (`ateapi.proto`)**: Added `page_size`, `page_token`, and `next_page_token` fields. - **Client (`kubectl-ate`)**: Updated to seamlessly loop through `next_page_token` values until all actors are fetched, presenting a continuous stream to the user without blowing up memory on a single request. - **API Handler (`controlapi`)**: Bounded requests with a max `page_size` limit of `1000`. - **Persistence (`ateredis`)**: - Implemented `listActorsPageToken` to safely serialize and encode cross-shard cursors. - Replaced `ForEachMaster` concurrent execution with deterministic sequential iteration. - Implemented Valkey Pipeline batching (`Pipelined`) for `GET` requests to eliminate N+1 network latency. - **Testing**: Added `TestListActors_Pagination` at both the integration level (`functional_test.go`) and data layer (`ateredis_test.go`). Refactored sorting to use `cmpopts.SortSlices`. ## Design Decisions During the design phase, several explicit trade-offs were made to balance performance with API simplicity: 1. **Opaque Tokens (AIP-158):** The `page_token` is strictly an opaque string at the `ateapi` layer. The client (`kubectl-ate`) is completely unaware that it contains base64-encoded JSON mapping to `{"shard", "cursor"}`. This keeps our API database-agnostic. 2. **Empty Pages & Client-Side Looping:** The persistence layer is kept simple and "dumb". If a shard scan returns 0 keys but the cursor is not `0`, we return an empty page. We rely on the client `kubectl-ate` to `for` loop until `next_page_token` is completely exhausted. Update: I have server side looping across shards until pagesize is met to limit round trips due to high shard count or low actor count on distributed shards. 3. **"Soft" Consistency (Missing/Duplicate Data):** `ListActors` is an infrequent operations command, not on the critical path. Taking a global cluster lock for iteration would severely degrade performance. We accept eventual consistency here: if actors migrate shards during an active scan, they might be duplicated or missed. 4. **Topology Flux Errors:** If a client passes a valid token pointing to a shard address that no longer exists (due to a cluster resharding mid-scan), the API simply returns an `Aborted` error. We do not attempt complex recovery; the client must just restart the `ListActors` command. 5. **Bounded Page Size:** Requests are currently capped at a maximum of `1000` to prevent single massive payloads. We can revisit/change the max later. NOTE: Actor filtering will be implemented later. ## Architecture Flow (Valkey) Note: If there are future plugin-db models they will need to handle the cusor/pages at their own implementation layer per the store.go interface. The system now seamlessly iterates across shards, handing off cursors safely between nodes: ```mermaid sequenceDiagram participant CLI as kubectl-ate participant API as ateapi participant Valkey as Valkey Cluster (Multiple Shards) CLI->>API: ListActors(page_size=1000, page_token="") Note over API: Request 1: Empty Token API->>API: Decode page_token (empty -> start fresh) API->>Valkey: Fetch all Master Nodes Valkey-->>API: [Master C, Master A, Master B] API->>API: Sort Masters alphabetically by IP/Port Note over API: Sorted: [Master A, Master B, Master C] API->>Valkey: SCAN (Master A, cursor=0, count=1000) Valkey-->>API: keys=[k1..k300], next_cursor=45 Note over API: OPTIMIZATION: Pipelining avoids N+1 problem API->>Valkey: PIPELINE GET (k1..k300) directly on Master A Valkey-->>API: values=[v1..v300] API->>API: Decode Protobufs API->>API: Encode page_token={"shard": "Master A", "cursor": 45} API-->>CLI: 300 Actors, next_page_token="base64(eyJzaGFyZ..." %% SECOND REQUEST CLI->>API: ListActors(page_size=1000, page_token="base64(...)") Note over API: Request 2: Resuming Mid-Shard API->>API: Decode page_token -> Shard="Master A", Cursor=45 API->>API: Fast-forward directly to Master A API->>Valkey: SCAN (Master A, cursor=45, count=1000) Valkey-->>API: keys=[k301..k400], next_cursor=0 API->>Valkey: PIPELINE GET (k301..k400) directly on Master A Valkey-->>API: values=[v301..v400] Note over API: Shard A is empty. Token points to Shard B! API->>API: Encode page_token={"shard": "Master B", "cursor": 0} API-->>CLI: 100 Actors, next_page_token="base64(eyJzaGFyZ..." %% THIRD REQUEST CLI->>API: ListActors(page_size=1000, page_token="base64(...)") Note over API: Request 3: Starting New Shard API->>API: Decode page_token -> Shard="Master B", Cursor=0 Note over API: Continues scanning Master B... ``` ## Testing: ``` > ./bin/kubectl-ate get actors NAMESPACE TEMPLATE ID STATUS ATEOM POD ATEOM IP VERSION ate-demo-counter counter 61b478cc-2e44-4ec7-ac6c-1af3f385de86 STATUS_SUSPENDED <none> 5 ``` ``` ~/go/bin/grpcurl -insecure -d '{"page_size": 2}' localhost:9090 ateapi.Control/ListActors { "actors": [ { "actorId": "test-counter-4", "version": "1", "actorTemplateNamespace": "ate-demo-counter", "actorTemplateName": "counter", "status": "STATUS_SUSPENDED" } ], "nextPageToken": "eyJzaGFyZCI6IjEwLjI0NC4wLjE5OjYzNzkiLCJjdXJzb3IiOjI1MjU1fQ==" } ``` <details> ``` bash -c ' export NEXT_TOKEN="" while :; do echo "--- Fetching Page ---" # Fetch the data and store the response RESPONSE=$(~/go/bin/grpcurl -insecure -d "{\"page_size\": 2, \"page_token\": \"$NEXT_TOKEN\"}" localhost:9090 ateapi.Control/ListActors) # Print the response to the terminal echo "$RESPONSE" # Extract the token without needing jq NEXT_TOKEN=$(echo "$RESPONSE" | grep -o "\"nextPageToken\": \"[^\"]*\"" | cut -d"\"" -f4) # If no token is found, we reached the end if [ -z "$NEXT_TOKEN" ]; then echo "--- Reached the end! ---" break fi echo "" read -p "Press Enter to fetch the next page using the new token..." done ' ``` --- Fetching Page --- { "actors": [ { "actorId": "test-counter-4", "version": "1", "actorTemplateNamespace": "ate-demo-counter", "actorTemplateName": "counter", "status": "STATUS_SUSPENDED" } ], "nextPageToken": "eyJzaGFyZCI6IjEwLjI0NC4wLjE5OjYzNzkiLCJjdXJzb3IiOjI1MjU1fQ==" } Press Enter to fetch the next page using the new token... --- Fetching Page --- { "nextPageToken": "eyJzaGFyZCI6IjEwLjI0NC4wLjIxOjYzNzkiLCJjdXJzb3IiOjB9" } Press Enter to fetch the next page using the new token... --- Fetching Page --- { "actors": [ { "actorId": "test-counter-1", "version": "1", "actorTemplateNamespace": "ate-demo-counter", "actorTemplateName": "counter", "status": "STATUS_SUSPENDED" } ], "nextPageToken": "eyJzaGFyZCI6IjEwLjI0NC4wLjIxOjYzNzkiLCJjdXJzb3IiOjQ1NjAwfQ==" } Press Enter to fetch the next page using the new token... --- Fetching Page --- { "actors": [ { "actorId": "test-counter-5", "version": "1", "actorTemplateNamespace": "ate-demo-counter", "actorTemplateName": "counter", "status": "STATUS_SUSPENDED" } ], "nextPageToken": "eyJzaGFyZCI6IjEwLjI0NC4wLjIxOjYzNzkiLCJjdXJzb3IiOjYyMDg2fQ==" } Press Enter to fetch the next page using the new token... --- Fetching Page --- { "nextPageToken": "eyJzaGFyZCI6IjEwLjI0NC4wLjI1OjYzNzkiLCJjdXJzb3IiOjB9" } Press Enter to fetch the next page using the new token... --- Fetching Page --- { "actors": [ { "actorId": "test-counter-2", "version": "1", "actorTemplateNamespace": "ate-demo-counter", "actorTemplateName": "counter", "status": "STATUS_SUSPENDED" } ], "nextPageToken": "eyJzaGFyZCI6IjEwLjI0NC4wLjI1OjYzNzkiLCJjdXJzb3IiOjE2OTkzfQ==" } Press Enter to fetch the next page using the new token... --- Fetching Page --- { "actors": [ { "actorId": "test-counter-3", "version": "1", "actorTemplateNamespace": "ate-demo-counter", "actorTemplateName": "counter", "status": "STATUS_SUSPENDED" }, { "actorId": "61b478cc-2e44-4ec7-ac6c-1af3f385de86", "version": "5", "actorTemplateNamespace": "ate-demo-counter", "actorTemplateName": "counter", "status": "STATUS_SUSPENDED", "lastSnapshot": "gs://ate-snapshots/ate-demo-counter/61b478cc-2e44-4ec7-ac6c-1af3f385de86/2026-06-03T03:17:34Z-UMYVUZOWW6LDPS5X5FSG3KC4VV" } ] } --- Reached the end! --- </details>
1398 lines
41 KiB
Go
1398 lines
41 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"
|
|
"net"
|
|
"os"
|
|
"os/exec"
|
|
"strings"
|
|
"sync"
|
|
"testing"
|
|
"time"
|
|
|
|
"github.com/agent-substrate/substrate/cmd/ateapi/internal/store/ateredis"
|
|
"github.com/agent-substrate/substrate/internal/ateinterceptors"
|
|
"github.com/agent-substrate/substrate/internal/proto/ateletpb"
|
|
atev1alpha1 "github.com/agent-substrate/substrate/pkg/api/v1alpha1"
|
|
"github.com/agent-substrate/substrate/pkg/client/clientset/versioned"
|
|
"github.com/agent-substrate/substrate/pkg/client/informers/externalversions"
|
|
listersv1alpha1 "github.com/agent-substrate/substrate/pkg/client/listers/api/v1alpha1"
|
|
"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
|
|
"github.com/alicebob/miniredis/v2"
|
|
"github.com/google/go-cmp/cmp"
|
|
"github.com/google/go-cmp/cmp/cmpopts"
|
|
"github.com/redis/go-redis/v9"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/codes"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
"google.golang.org/grpc/status"
|
|
"google.golang.org/protobuf/testing/protocmp"
|
|
corev1 "k8s.io/api/core/v1"
|
|
metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
|
|
"k8s.io/apimachinery/pkg/util/wait"
|
|
"k8s.io/client-go/kubernetes"
|
|
"k8s.io/client-go/rest"
|
|
"sigs.k8s.io/controller-runtime/pkg/envtest"
|
|
)
|
|
|
|
var (
|
|
testEnv *envtest.Environment
|
|
cfg *rest.Config
|
|
fakeAtelet = &FakeAteletServer{}
|
|
)
|
|
|
|
func TestMain(m *testing.M) {
|
|
cmd := exec.Command("bash", "../../../../hack/run-tool.sh", "setup-envtest", "use", "--print", "path")
|
|
out, err := cmd.Output()
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
binaryAssetsDirectory := strings.TrimSpace(string(out))
|
|
|
|
testEnv = &envtest.Environment{
|
|
CRDDirectoryPaths: []string{"../../../../manifests/ate-install/generated"},
|
|
BinaryAssetsDirectory: binaryAssetsDirectory,
|
|
}
|
|
|
|
cfg, err = testEnv.Start()
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
|
|
// Create ate-system namespace
|
|
k8sClient, err := kubernetes.NewForConfig(cfg)
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
_, err = k8sClient.CoreV1().Namespaces().Create(context.Background(), &corev1.Namespace{
|
|
ObjectMeta: metav1.ObjectMeta{Name: "ate-system"},
|
|
}, metav1.CreateOptions{})
|
|
if err != nil && !strings.Contains(err.Error(), "already exists") {
|
|
os.Exit(1)
|
|
}
|
|
|
|
// Create shared Atelet Pod
|
|
ateletPod := &corev1.Pod{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "atelet-shared",
|
|
Namespace: "ate-system",
|
|
Labels: map[string]string{
|
|
"app": "atelet",
|
|
},
|
|
},
|
|
Spec: corev1.PodSpec{
|
|
NodeName: "node1",
|
|
Containers: []corev1.Container{
|
|
{Name: "main", Image: "nginx"},
|
|
},
|
|
},
|
|
}
|
|
createdAtelet, err := k8sClient.CoreV1().Pods("ate-system").Create(context.Background(), ateletPod, metav1.CreateOptions{})
|
|
if err != nil && !strings.Contains(err.Error(), "already exists") {
|
|
os.Exit(1)
|
|
}
|
|
if err == nil {
|
|
createdAtelet.Status.PodIPs = []corev1.PodIP{{IP: "127.0.0.1"}}
|
|
createdAtelet.Status.Phase = corev1.PodRunning
|
|
_, err = k8sClient.CoreV1().Pods("ate-system").UpdateStatus(context.Background(), createdAtelet, metav1.UpdateOptions{})
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
}
|
|
|
|
// Start Fake Atelet Server on port 8085
|
|
ateletGrpcServer := grpc.NewServer()
|
|
ateletpb.RegisterAteomHerderServer(ateletGrpcServer, fakeAtelet)
|
|
ateletLis, err := net.Listen("tcp", "127.0.0.1:8085")
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
go func() {
|
|
if err := ateletGrpcServer.Serve(ateletLis); err != nil {
|
|
fmt.Printf("atelet grpc server exited: %v\n", err)
|
|
}
|
|
}()
|
|
|
|
code := m.Run()
|
|
|
|
ateletGrpcServer.Stop()
|
|
|
|
err = testEnv.Stop()
|
|
if err != nil {
|
|
os.Exit(1)
|
|
}
|
|
|
|
os.Exit(code)
|
|
}
|
|
|
|
// FakeAteletServer implements ateletpb.WorkersServer
|
|
type FakeAteletServer struct {
|
|
ateletpb.UnimplementedAteomHerderServer
|
|
|
|
Lock sync.Mutex
|
|
|
|
RunCalled bool
|
|
|
|
CheckpointCalled bool
|
|
|
|
RestoreCalled bool
|
|
FailRestore error
|
|
RestoreDelay time.Duration
|
|
}
|
|
|
|
func (f *FakeAteletServer) Reset() {
|
|
f.Lock.Lock()
|
|
defer f.Lock.Unlock()
|
|
|
|
f.RunCalled = false
|
|
|
|
f.CheckpointCalled = false
|
|
|
|
f.RestoreCalled = false
|
|
f.FailRestore = nil
|
|
f.RestoreDelay = 0
|
|
}
|
|
|
|
func (f *FakeAteletServer) Run(ctx context.Context, req *ateletpb.RunRequest) (*ateletpb.RunResponse, error) {
|
|
f.Lock.Lock()
|
|
defer f.Lock.Unlock()
|
|
|
|
f.RunCalled = true
|
|
|
|
return &ateletpb.RunResponse{}, nil
|
|
}
|
|
|
|
func (f *FakeAteletServer) Checkpoint(ctx context.Context, req *ateletpb.CheckpointRequest) (*ateletpb.CheckpointResponse, error) {
|
|
f.Lock.Lock()
|
|
defer f.Lock.Unlock()
|
|
|
|
f.CheckpointCalled = true
|
|
|
|
return &ateletpb.CheckpointResponse{}, nil
|
|
}
|
|
|
|
func (f *FakeAteletServer) Restore(ctx context.Context, req *ateletpb.RestoreRequest) (*ateletpb.RestoreResponse, error) {
|
|
f.Lock.Lock()
|
|
defer f.Lock.Unlock()
|
|
|
|
f.RestoreCalled = true
|
|
if f.RestoreDelay > 0 {
|
|
time.Sleep(f.RestoreDelay)
|
|
}
|
|
if f.FailRestore != nil {
|
|
return nil, f.FailRestore
|
|
}
|
|
return &ateletpb.RestoreResponse{}, nil
|
|
}
|
|
|
|
type testContext struct {
|
|
mr *miniredis.Miniredis
|
|
service *Service
|
|
client ateapipb.ControlClient
|
|
k8sClient kubernetes.Interface
|
|
substrateClient versioned.Interface
|
|
persistence *ateredis.Persistence
|
|
fakeAtelet *FakeAteletServer
|
|
cleanup func()
|
|
actorTemplateLister listersv1alpha1.ActorTemplateLister
|
|
}
|
|
|
|
// setupTest sets up a fully isolated test environment.
|
|
func setupTest(t *testing.T, ns string) *testContext {
|
|
t.Helper()
|
|
// 1. Start Miniredis
|
|
mr, err := miniredis.Run()
|
|
if err != nil {
|
|
t.Fatalf("failed to start miniredis: %v", err)
|
|
}
|
|
|
|
rdb := redis.NewClusterClient(&redis.ClusterOptions{
|
|
Addrs: []string{mr.Addr()},
|
|
})
|
|
persistence := ateredis.NewPersistence(rdb)
|
|
|
|
// 2. Initialize Clientsets using global cfg
|
|
k8sClient, err := kubernetes.NewForConfig(cfg)
|
|
if err != nil {
|
|
mr.Close()
|
|
t.Fatalf("failed to create k8s clientset: %v", err)
|
|
}
|
|
|
|
substrateClient, err := versioned.NewForConfig(cfg)
|
|
if err != nil {
|
|
mr.Close()
|
|
t.Fatalf("failed to create substrate clientset: %v", err)
|
|
}
|
|
|
|
// 3. Initialize Informers
|
|
workerFactory, workerInformer := WorkerPodInformer(k8sClient)
|
|
ateletFactory, ateletInformer := AteletInformer(k8sClient)
|
|
|
|
substrateInformerFactory := externalversions.NewSharedInformerFactory(substrateClient, 0)
|
|
actorTemplateLister := substrateInformerFactory.Api().V1alpha1().ActorTemplates().Lister()
|
|
|
|
ctx, cancel := context.WithCancel(context.Background())
|
|
|
|
syncer := NewWorkerPoolSyncer(persistence, workerInformer)
|
|
syncer.Start(ctx)
|
|
|
|
workerFactory.Start(ctx.Done())
|
|
ateletFactory.Start(ctx.Done())
|
|
substrateInformerFactory.Start(ctx.Done())
|
|
|
|
workerFactory.WaitForCacheSync(ctx.Done())
|
|
ateletFactory.WaitForCacheSync(ctx.Done())
|
|
substrateInformerFactory.WaitForCacheSync(ctx.Done())
|
|
|
|
// 4. Initialize Service
|
|
dialer := NewAteletDialer(workerInformer.GetIndexer(), ateletInformer.GetIndexer())
|
|
service := NewService(persistence, actorTemplateLister, dialer)
|
|
|
|
// 5. Start REAL gRPC Server for ATE API
|
|
grpcServer := grpc.NewServer(grpc.UnaryInterceptor(ateinterceptors.ServerUnaryInterceptor))
|
|
ateapipb.RegisterControlServer(grpcServer, service)
|
|
|
|
lis, err := net.Listen("tcp", "localhost:0")
|
|
if err != nil {
|
|
cancel()
|
|
mr.Close()
|
|
t.Fatalf("failed to listen: %v", err)
|
|
}
|
|
|
|
go func() {
|
|
if err := grpcServer.Serve(lis); err != nil {
|
|
t.Logf("grpc server exited: %v", err)
|
|
}
|
|
}()
|
|
|
|
conn, err := grpc.NewClient(lis.Addr().String(), grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
if err != nil {
|
|
grpcServer.Stop()
|
|
cancel()
|
|
mr.Close()
|
|
t.Fatalf("failed to connect: %v", err)
|
|
}
|
|
|
|
client := ateapipb.NewControlClient(conn)
|
|
|
|
// Call Reset on global mock
|
|
fakeAtelet.Reset()
|
|
|
|
// Create namespace
|
|
_, err = k8sClient.CoreV1().Namespaces().Create(context.Background(), &corev1.Namespace{
|
|
ObjectMeta: metav1.ObjectMeta{Name: ns},
|
|
}, metav1.CreateOptions{})
|
|
if err != nil {
|
|
conn.Close()
|
|
grpcServer.Stop()
|
|
cancel()
|
|
mr.Close()
|
|
t.Fatalf("failed to create namespace %s: %v", ns, err)
|
|
}
|
|
|
|
cleanup := func() {
|
|
conn.Close()
|
|
grpcServer.Stop()
|
|
cancel()
|
|
mr.Close()
|
|
}
|
|
|
|
return &testContext{
|
|
mr: mr,
|
|
service: service,
|
|
client: client,
|
|
k8sClient: k8sClient,
|
|
substrateClient: substrateClient,
|
|
persistence: persistence,
|
|
fakeAtelet: fakeAtelet,
|
|
cleanup: cleanup,
|
|
actorTemplateLister: actorTemplateLister,
|
|
}
|
|
}
|
|
|
|
func namespaceForTest(baseName string) string {
|
|
return fmt.Sprintf("%s-%d", baseName, time.Now().UnixNano())
|
|
}
|
|
|
|
func createTemplate(t *testing.T, tc *testContext, ns string) {
|
|
t.Helper()
|
|
actorTemplate := &atev1alpha1.ActorTemplate{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: "tmpl1",
|
|
Namespace: ns,
|
|
},
|
|
Spec: atev1alpha1.ActorTemplateSpec{
|
|
Runsc: atev1alpha1.RunscConfig{
|
|
AMD64: &atev1alpha1.RunscPlatformConfig{
|
|
URL: "gs://gvisor/releases/nightly/2026-05-19/x86_64/runsc",
|
|
SHA256Hash: "a397be1abc2420d26bce6c70e6e2ff96c73aaaab929756c56f5e2089ea842b63",
|
|
},
|
|
},
|
|
PauseImage: "pause@sha256:abc",
|
|
Containers: []atev1alpha1.Container{
|
|
{
|
|
Name: "main",
|
|
Image: "main@sha256:abc",
|
|
Command: []string{"/main"},
|
|
},
|
|
},
|
|
WorkerPoolRef: corev1.ObjectReference{
|
|
Namespace: ns,
|
|
Name: "pool1",
|
|
},
|
|
},
|
|
}
|
|
createdTemplate, err := tc.substrateClient.ApiV1alpha1().ActorTemplates(ns).Create(context.Background(), actorTemplate, metav1.CreateOptions{})
|
|
if err != nil {
|
|
t.Fatalf("failed to create actor template: %v", err)
|
|
}
|
|
|
|
createdTemplate.Status = atev1alpha1.ActorTemplateStatus{
|
|
GoldenSnapshot: "gs://my-bucket/my-folder",
|
|
}
|
|
|
|
_, err = tc.substrateClient.ApiV1alpha1().ActorTemplates(ns).UpdateStatus(context.Background(), createdTemplate, metav1.UpdateOptions{})
|
|
if err != nil {
|
|
t.Fatalf("failed to update status: %v", err)
|
|
}
|
|
|
|
// Wait for Informer cache to sync
|
|
err = wait.PollUntilContextTimeout(context.Background(), 100*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
|
|
tmpl, err := tc.actorTemplateLister.ActorTemplates(ns).Get("tmpl1")
|
|
if err != nil {
|
|
return false, nil // Retry if not found in cache yet
|
|
}
|
|
return tmpl.Status.GoldenSnapshot != "", nil
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("failed to wait for template status update in informer: %v", err)
|
|
}
|
|
}
|
|
|
|
func createWorkerPod(t *testing.T, tc *testContext, ns string, name string, nodeName string) {
|
|
t.Helper()
|
|
pod := &corev1.Pod{
|
|
ObjectMeta: metav1.ObjectMeta{
|
|
Name: name,
|
|
Namespace: ns,
|
|
Labels: map[string]string{
|
|
"ate.dev/worker-pool": "pool1",
|
|
},
|
|
},
|
|
Spec: corev1.PodSpec{
|
|
NodeName: nodeName,
|
|
Containers: []corev1.Container{
|
|
{Name: "main", Image: "nginx"},
|
|
},
|
|
},
|
|
}
|
|
createdPod, err := tc.k8sClient.CoreV1().Pods(ns).Create(context.Background(), pod, metav1.CreateOptions{})
|
|
if err != nil {
|
|
t.Fatalf("failed to create worker pod: %v", err)
|
|
}
|
|
createdPod.Status.PodIPs = []corev1.PodIP{{IP: "127.0.0.1"}}
|
|
createdPod.Status.Phase = corev1.PodRunning
|
|
_, err = tc.k8sClient.CoreV1().Pods(ns).UpdateStatus(context.Background(), createdPod, metav1.UpdateOptions{})
|
|
if err != nil {
|
|
t.Fatalf("failed to update worker pod status: %v", err)
|
|
}
|
|
|
|
// Wait for worker to be registered via API
|
|
err = wait.PollUntilContextTimeout(context.Background(), 100*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
|
|
resp, err := tc.client.ListWorkers(ctx, &ateapipb.ListWorkersRequest{})
|
|
if err != nil {
|
|
return false, nil // Retry on API error
|
|
}
|
|
for _, w := range resp.GetWorkers() {
|
|
if w.GetWorkerNamespace() == ns && w.GetWorkerPod() == name {
|
|
return true, nil
|
|
}
|
|
}
|
|
return false, nil
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("failed to wait for worker to be registered: %v", err)
|
|
}
|
|
}
|
|
|
|
func deleteWorkerPod(t *testing.T, tc *testContext, ns string, name string) {
|
|
t.Helper()
|
|
err := tc.k8sClient.CoreV1().Pods(ns).Delete(context.Background(), name, metav1.DeleteOptions{})
|
|
if err != nil {
|
|
t.Fatalf("failed to delete worker pod %s: %v", name, err)
|
|
}
|
|
|
|
// Wait for worker to be removed from API
|
|
err = wait.PollUntilContextTimeout(context.Background(), 100*time.Millisecond, 5*time.Second, true, func(ctx context.Context) (bool, error) {
|
|
resp, err := tc.client.ListWorkers(ctx, &ateapipb.ListWorkersRequest{})
|
|
if err != nil {
|
|
return false, nil // Retry on API error
|
|
}
|
|
for _, w := range resp.GetWorkers() {
|
|
if w.GetWorkerNamespace() == ns && w.GetWorkerPod() == name {
|
|
return false, nil // Still there
|
|
}
|
|
}
|
|
return true, nil // Gone!
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("failed to wait for worker to be removed: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestCreateActor_Success tests the happy path for creating an actor.
|
|
// Workflow:
|
|
// 1. Creates a mock ActorTemplate in the test namespace.
|
|
// 2. Calls CreateActor RPC.
|
|
// 3. Verifies that the actor is successfully created and returned in the response with a generated ID.
|
|
func TestCreateActor_Success(t *testing.T) {
|
|
ns := namespaceForTest("ns-create-success")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
createResp, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
|
|
want := &ateapipb.CreateActorResponse{
|
|
Actor: &ateapipb.Actor{
|
|
ActorId: "id1",
|
|
Version: 1,
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
Status: ateapipb.Actor_STATUS_SUSPENDED,
|
|
},
|
|
}
|
|
|
|
if diff := cmp.Diff(want, createResp, protocmp.Transform()); diff != "" {
|
|
t.Errorf("CreateActor response mismatch (-want +got):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
// TestCreateActor_TemplateNotFound tests that creating an actor with a non-existent template fails with FailedPrecondition.
|
|
func TestCreateActor_TemplateNotFound(t *testing.T) {
|
|
ns := namespaceForTest("ns-create-notfound")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "non-existent",
|
|
ActorId: "id1",
|
|
})
|
|
assertGrpcError(t, err, codes.FailedPrecondition, fmt.Sprintf("ActorTemplate %s/non-existent not found", ns))
|
|
}
|
|
|
|
// TestCreateActor_Duplicate tests that creating an actor with an existing ID fails.
|
|
func TestCreateActor_Duplicate(t *testing.T) {
|
|
ns := namespaceForTest("ns-create-dup")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("first CreateActor failed: %v", err)
|
|
}
|
|
|
|
_, err = tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
assertGrpcError(t, err, codes.AlreadyExists, "Actor id1 already exists")
|
|
}
|
|
|
|
// TestGetActor_Found tests that an existing actor can be retrieved.
|
|
func TestGetActor_Found(t *testing.T) {
|
|
ns := namespaceForTest("ns-get-found")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
createResp, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
|
|
id := createResp.GetActor().GetActorId()
|
|
|
|
getResp, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("GetActor failed: %v", err)
|
|
}
|
|
|
|
want := &ateapipb.GetActorResponse{
|
|
Actor: createResp.GetActor(),
|
|
}
|
|
|
|
if diff := cmp.Diff(want, getResp, protocmp.Transform()); diff != "" {
|
|
t.Errorf("GetActor response mismatch (-want +got):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
// TestGetActor_NotFound tests that retrieving a non-existent actor fails.
|
|
// Workflow:
|
|
// 1. Calls GetActor RPC with a non-existent ID.
|
|
// 2. Verifies that it returns an error (NotFound).
|
|
func TestGetActor_NotFound(t *testing.T) {
|
|
ns := namespaceForTest("ns-get-notfound")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
_, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{
|
|
ActorId: "non-existent",
|
|
})
|
|
assertGrpcError(t, err, codes.NotFound, "Actor non-existent not found")
|
|
}
|
|
|
|
// TestListActors tests that all created actors can be listed.
|
|
// Workflow:
|
|
// 1. Creates a mock ActorTemplate.
|
|
// 2. Calls CreateActor twice to create two actors.
|
|
// 3. Calls ListActors RPC.
|
|
// 4. Verifies that both actors are returned in the list.
|
|
func TestListActors(t *testing.T) {
|
|
ns := namespaceForTest("ns-list-actors")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
resp1, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor 1 failed: %v", err)
|
|
}
|
|
resp2, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id2",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor 2 failed: %v", err)
|
|
}
|
|
|
|
listResp, err := tc.client.ListActors(context.Background(), &ateapipb.ListActorsRequest{})
|
|
if err != nil {
|
|
t.Fatalf("ListActors failed: %v", err)
|
|
}
|
|
|
|
if len(listResp.Actors) != 2 {
|
|
t.Fatalf("expected 2 actors, got %d", len(listResp.Actors))
|
|
}
|
|
|
|
want := []*ateapipb.Actor{
|
|
resp1.GetActor(),
|
|
resp2.GetActor(),
|
|
}
|
|
|
|
opts := []cmp.Option{
|
|
protocmp.Transform(),
|
|
cmpopts.SortSlices(func(a, b *ateapipb.Actor) bool {
|
|
return a.ActorId < b.ActorId
|
|
}),
|
|
}
|
|
|
|
if diff := cmp.Diff(want, listResp.Actors, opts...); diff != "" {
|
|
t.Errorf("ListActors response mismatch (-want +got):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
// TestListActors_Pagination tests that ListActors correctly paginates results.
|
|
func TestListActors_Pagination(t *testing.T) {
|
|
ns := namespaceForTest("ns-list-actors-pagination")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
var want []*ateapipb.Actor
|
|
for i := 0; i < 5; i++ {
|
|
resp, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: fmt.Sprintf("id%d", i),
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor %d failed: %v", i, err)
|
|
}
|
|
want = append(want, resp.GetActor())
|
|
}
|
|
|
|
var allActors []*ateapipb.Actor
|
|
pageToken := ""
|
|
|
|
for {
|
|
listResp, err := tc.client.ListActors(context.Background(), &ateapipb.ListActorsRequest{
|
|
PageSize: 2,
|
|
PageToken: pageToken,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("ListActors failed: %v", err)
|
|
}
|
|
|
|
allActors = append(allActors, listResp.Actors...)
|
|
pageToken = listResp.GetNextPageToken()
|
|
if pageToken == "" {
|
|
break
|
|
}
|
|
}
|
|
|
|
if len(allActors) != 5 {
|
|
t.Fatalf("expected 5 actors total, got %d", len(allActors))
|
|
}
|
|
|
|
opts := []cmp.Option{
|
|
protocmp.Transform(),
|
|
cmpopts.SortSlices(func(a, b *ateapipb.Actor) bool {
|
|
return a.ActorId < b.ActorId
|
|
}),
|
|
}
|
|
|
|
if diff := cmp.Diff(want, allActors, opts...); diff != "" {
|
|
t.Errorf("ListActors pagination response mismatch (-want +got):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
func TestListActors_PageSizeValidation(t *testing.T) {
|
|
ns := namespaceForTest("ns-list-actors-validation")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
// 1. Negative page size
|
|
_, err := tc.client.ListActors(context.Background(), &ateapipb.ListActorsRequest{
|
|
PageSize: -1,
|
|
})
|
|
if status.Code(err) != codes.InvalidArgument {
|
|
t.Errorf("expected InvalidArgument error for negative page_size, got: %v", err)
|
|
}
|
|
|
|
// 2. Page size exceeding maxPageSize (1000)
|
|
_, err = tc.client.ListActors(context.Background(), &ateapipb.ListActorsRequest{
|
|
PageSize: 1001,
|
|
})
|
|
if status.Code(err) != codes.InvalidArgument {
|
|
t.Errorf("expected InvalidArgument error for page_size > 1000, got: %v", err)
|
|
}
|
|
}
|
|
|
|
// TestListWorkers tests that workers mirrored to Redis are listed.
|
|
// Workflow:
|
|
// 1. Creates a mock worker Pod in Kubernetes.
|
|
// 2. Waits for the background WorkerPoolSyncer to mirror it to Redis.
|
|
// 3. Calls ListWorkers RPC.
|
|
// 4. Verifies that the worker appears in the response.
|
|
func TestListWorkers(t *testing.T) {
|
|
ns := namespaceForTest("ns-list-workers")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createWorkerPod(t, tc, ns, "worker-1", "")
|
|
|
|
listResp, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{})
|
|
if err != nil {
|
|
t.Fatalf("ListWorkers failed: %v", err)
|
|
}
|
|
|
|
var filteredWorkers []*ateapipb.Worker
|
|
for _, w := range listResp.GetWorkers() {
|
|
if w.GetWorkerNamespace() == ns {
|
|
filteredWorkers = append(filteredWorkers, w)
|
|
}
|
|
}
|
|
|
|
want := []*ateapipb.Worker{
|
|
{
|
|
WorkerNamespace: ns,
|
|
WorkerPool: "pool1",
|
|
WorkerPod: "worker-1",
|
|
Ip: "127.0.0.1",
|
|
Version: 1,
|
|
},
|
|
}
|
|
|
|
if diff := cmp.Diff(want, filteredWorkers, protocmp.Transform(), protocmp.IgnoreFields(&ateapipb.Worker{}, "worker_pod_uid")); diff != "" {
|
|
t.Errorf("ListWorkers response mismatch (-want +got):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
// TestResumeActor tests the full workflow of resuming a suspended actor.
|
|
// Workflow:
|
|
// 1. Creates a mock ActorTemplate.
|
|
// 2. Creates a mock Atelet Pod in 'ate-system' namespace on 'node1'.
|
|
// 3. Creates a mock worker Pod in the test namespace on 'node1'.
|
|
// 4. Waits for the WorkerPoolSyncer to mirror the worker to Redis.
|
|
// 5. Creates an actor (starts as SUSPENDED).
|
|
// 6. Calls ResumeActor RPC.
|
|
// 7. Verifies that the fake Atelet received the Restore call.
|
|
// 8. Verifies that the actor status is updated to RUNNING.
|
|
func TestResumeActor(t *testing.T) {
|
|
ns := namespaceForTest("ns-resume")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
createWorkerPod(t, tc, ns, "worker-1", "node1")
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
id := "id1"
|
|
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("ResumeActor failed: %v", err)
|
|
}
|
|
|
|
if !tc.fakeAtelet.RestoreCalled {
|
|
t.Errorf("expected Restore to be called")
|
|
}
|
|
|
|
getResp, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("GetActor failed: %v", err)
|
|
}
|
|
want := &ateapipb.GetActorResponse{
|
|
Actor: &ateapipb.Actor{
|
|
ActorId: id,
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
Status: ateapipb.Actor_STATUS_RUNNING,
|
|
AteomPodNamespace: ns,
|
|
AteomPodName: "worker-1",
|
|
AteomPodIp: "127.0.0.1",
|
|
},
|
|
}
|
|
if diff := cmp.Diff(want, getResp, protocmp.Transform(), protocmp.IgnoreFields(&ateapipb.Actor{}, "version"), protocmp.IgnoreFields(&ateapipb.Actor{}, "ateom_pod_uid")); diff != "" {
|
|
t.Errorf("GetActor response mismatch (-want +got):\n%s", diff)
|
|
}
|
|
|
|
// Verify that the worker record also has the assigned actor details
|
|
listWorkersResp, err := tc.client.ListWorkers(context.Background(), &ateapipb.ListWorkersRequest{})
|
|
if err != nil {
|
|
t.Fatalf("ListWorkers failed: %v", err)
|
|
}
|
|
var actorWorker *ateapipb.Worker
|
|
for _, w := range listWorkersResp.GetWorkers() {
|
|
if w.GetWorkerNamespace() == ns && w.GetWorkerPod() == "worker-1" {
|
|
actorWorker = w
|
|
break
|
|
}
|
|
}
|
|
if actorWorker == nil {
|
|
t.Fatalf("expected worker-1 in namespace %s not found in ListWorkers", ns)
|
|
}
|
|
|
|
wantWorker := &ateapipb.Worker{
|
|
WorkerNamespace: ns,
|
|
WorkerPool: "pool1",
|
|
WorkerPod: "worker-1",
|
|
ActorNamespace: ns,
|
|
ActorTemplate: "tmpl1",
|
|
ActorId: id,
|
|
Ip: "127.0.0.1",
|
|
}
|
|
|
|
if diff := cmp.Diff(wantWorker, actorWorker, protocmp.Transform(), protocmp.IgnoreFields(&ateapipb.Worker{}, "version"), protocmp.IgnoreFields(&ateapipb.Worker{}, "worker_pod_uid")); diff != "" {
|
|
t.Errorf("Worker state mismatch (-want +got):\n%s", diff)
|
|
}
|
|
}
|
|
|
|
// TestResumeActor_NoWorkers tests that resuming an actor fails when no free workers are available.
|
|
// Workflow:
|
|
// 1. Creates a mock ActorTemplate.
|
|
// 2. Creates an actor.
|
|
// 3. Calls ResumeActor RPC without creating any workers.
|
|
// 4. Verifies that ResumeActor fails with FailedPrecondition status.
|
|
// TestResumeActor_NoWorkers tests that resuming an actor fails when no free workers are available.
|
|
// Workflow:
|
|
// 1. Creates a mock ActorTemplate.
|
|
// 2. Creates an actor.
|
|
// 3. Calls ResumeActor RPC without creating any workers.
|
|
// 4. Verifies that ResumeActor fails with FailedPrecondition status.
|
|
func TestResumeActor_NoWorkers(t *testing.T) {
|
|
ns := namespaceForTest("ns-resume-no-workers")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
createResp, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
|
|
id := createResp.GetActor().GetActorId()
|
|
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
assertGrpcError(t, err, codes.FailedPrecondition, "no free workers available")
|
|
}
|
|
|
|
// TestResumeActor_Reentrancy tests the failure recovery and re-entrancy of ResumeActor.
|
|
// Workflow:
|
|
// 1. Creates a mock ActorTemplate.
|
|
// 2. Creates a mock Atelet Pod and a mock Worker Pod.
|
|
// 3. Waits for the WorkerPoolSyncer to mirror the worker to store.
|
|
// 4. Creates an actor in SUSPENDED state.
|
|
// 5. Configures fake Atelet to FAIL on Restore.
|
|
// 6. Calls ResumeActor and verifies it fails, but actor status becomes RESUMING.
|
|
// 7. Configures fake Atelet to SUCCEED on Restore.
|
|
// 8. Calls ResumeActor again and verifies it succeeds and actor status becomes RUNNING.
|
|
func TestResumeActor_Reentrancy(t *testing.T) {
|
|
ns := namespaceForTest("ns-resume-reentrancy")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
// Create Worker Pod
|
|
createWorkerPod(t, tc, ns, "worker-1", "node1")
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
id := "id1"
|
|
|
|
// STEP 1: Make Atelet FAIL on Restore!
|
|
tc.fakeAtelet.FailRestore = fmt.Errorf("mock atelet failure")
|
|
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err == nil {
|
|
t.Fatalf("expected ResumeActor to fail due to atelet error")
|
|
}
|
|
|
|
// Verify actor state is RESUMING in Redis!
|
|
actor, err := tc.persistence.GetActor(context.Background(), id)
|
|
if err != nil {
|
|
t.Fatalf("failed to get actor from store: %v", err)
|
|
}
|
|
if actor.GetStatus() != ateapipb.Actor_STATUS_RESUMING {
|
|
t.Errorf("expected status RESUMING, got %v", actor.GetStatus())
|
|
}
|
|
|
|
// STEP 2: Make Atelet SUCCEED!
|
|
tc.fakeAtelet.FailRestore = nil
|
|
tc.fakeAtelet.RestoreCalled = false // reset for verification
|
|
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("ResumeActor failed on retry: %v", err)
|
|
}
|
|
|
|
if !tc.fakeAtelet.RestoreCalled {
|
|
t.Errorf("expected Restore to be called on retry")
|
|
}
|
|
|
|
// Verify actor state is RUNNING!
|
|
actor, err = tc.persistence.GetActor(context.Background(), id)
|
|
if err != nil {
|
|
t.Fatalf("failed to get actor from store: %v", err)
|
|
}
|
|
if actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING {
|
|
t.Errorf("expected status RUNNING, got %v", actor.GetStatus())
|
|
}
|
|
}
|
|
|
|
// TestSuspendActor tests the full workflow of suspending a running actor.
|
|
// Workflow:
|
|
// 1. Creates a mock ActorTemplate.
|
|
// 2. Creates a mock Atelet Pod on 'node1'.
|
|
// 3. Creates a mock worker Pod on 'node1'.
|
|
// 4. Waits for the WorkerPoolSyncer to mirror the worker to Redis.
|
|
// 5. Creates an actor.
|
|
// 6. Calls ResumeActor to transition it to RUNNING.
|
|
// 7. Calls SuspendActor RPC.
|
|
// 8. Verifies that the fake Atelet received the Suspend call.
|
|
func TestSuspendActor(t *testing.T) {
|
|
ns := namespaceForTest("ns-suspend")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
createWorkerPod(t, tc, ns, "worker-1", "node1")
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
id := "id1"
|
|
|
|
// Resume first to make it running
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("ResumeActor failed: %v", err)
|
|
}
|
|
|
|
// Suspend
|
|
_, err = tc.client.SuspendActor(context.Background(), &ateapipb.SuspendActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("SuspendActor failed: %v", err)
|
|
}
|
|
|
|
if !tc.fakeAtelet.CheckpointCalled {
|
|
t.Errorf("expected atelet Checkpoint to be called")
|
|
}
|
|
|
|
getResp, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("GetActor failed: %v", err)
|
|
}
|
|
want := &ateapipb.GetActorResponse{
|
|
Actor: &ateapipb.Actor{
|
|
ActorId: id,
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
Status: ateapipb.Actor_STATUS_SUSPENDED,
|
|
LastSnapshot: fmt.Sprintf("gs://my-bucket/%s/tmpl1/%s/", ns, id),
|
|
},
|
|
}
|
|
|
|
if diff := cmp.Diff(want, getResp, protocmp.Transform(), protocmp.IgnoreFields(&ateapipb.Actor{}, "version"), protocmp.IgnoreFields(&ateapipb.Actor{}, "last_snapshot"), protocmp.IgnoreFields(&ateapipb.Actor{}, "ateom_pod_uid")); diff != "" {
|
|
t.Errorf("GetActor response mismatch (-want +got):\n%s", diff)
|
|
}
|
|
|
|
}
|
|
|
|
// TestValidation tests the negative validation cases for all gRPC methods.
|
|
// Workflow:
|
|
// 1. Uses table-driven tests for each RPC method (CreateActor, GetActor, ResumeActor, SuspendActor).
|
|
// 2. Passes invalid requests (missing required fields).
|
|
// 3. Verifies that all requests fail with an error.
|
|
func TestValidation(t *testing.T) {
|
|
ns := namespaceForTest("ns-validation")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
t.Run("CreateActor", func(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
req *ateapipb.CreateActorRequest
|
|
wantMsg string
|
|
}{
|
|
{"missing namespace", &ateapipb.CreateActorRequest{ActorTemplateName: "tmpl1", ActorId: "id1"}, "actor_template_namespace is required"},
|
|
{"missing template name", &ateapipb.CreateActorRequest{ActorTemplateNamespace: "ns1", ActorId: "id1"}, "actor_template_name is required"},
|
|
{"missing actor id", &ateapipb.CreateActorRequest{ActorTemplateNamespace: "ns1", ActorTemplateName: "tmpl1"}, "actor_id is required"},
|
|
{"invalid actor id (capitals)", &ateapipb.CreateActorRequest{ActorTemplateNamespace: "ns1", ActorTemplateName: "tmpl1", ActorId: "ID1"}, "invalid actor_id: must start and end with a lower case alphanumeric character, and consist only of lower case alphanumeric characters or '-'"},
|
|
{"invalid actor id (special chars)", &ateapipb.CreateActorRequest{ActorTemplateNamespace: "ns1", ActorTemplateName: "tmpl1", ActorId: "id_1"}, "invalid actor_id: must start and end with a lower case alphanumeric character, and consist only of lower case alphanumeric characters or '-'"},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
_, err := tc.client.CreateActor(context.Background(), tt.req)
|
|
assertGrpcError(t, err, codes.InvalidArgument, tt.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("GetActor", func(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
req *ateapipb.GetActorRequest
|
|
wantMsg string
|
|
}{
|
|
{"missing id", &ateapipb.GetActorRequest{}, "id is required"},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
_, err := tc.client.GetActor(context.Background(), tt.req)
|
|
assertGrpcError(t, err, codes.InvalidArgument, tt.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("ResumeActor", func(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
req *ateapipb.ResumeActorRequest
|
|
wantMsg string
|
|
}{
|
|
{"missing id", &ateapipb.ResumeActorRequest{}, "id is required"},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
_, err := tc.client.ResumeActor(context.Background(), tt.req)
|
|
assertGrpcError(t, err, codes.InvalidArgument, tt.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("SuspendActor", func(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
req *ateapipb.SuspendActorRequest
|
|
wantMsg string
|
|
}{
|
|
{"missing id", &ateapipb.SuspendActorRequest{}, "id is required"},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
_, err := tc.client.SuspendActor(context.Background(), tt.req)
|
|
assertGrpcError(t, err, codes.InvalidArgument, tt.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
|
|
t.Run("DeleteActor", func(t *testing.T) {
|
|
tests := []struct {
|
|
name string
|
|
req *ateapipb.DeleteActorRequest
|
|
wantMsg string
|
|
}{
|
|
{"missing id", &ateapipb.DeleteActorRequest{}, "actor_id is required"},
|
|
{"invalid actor id (capitals)", &ateapipb.DeleteActorRequest{ActorId: "ID1"}, "invalid actor_id: must start and end with a lower case alphanumeric character, and consist only of lower case alphanumeric characters or '-'"},
|
|
{"invalid actor id (special chars)", &ateapipb.DeleteActorRequest{ActorId: "id_1"}, "invalid actor_id: must start and end with a lower case alphanumeric character, and consist only of lower case alphanumeric characters or '-'"},
|
|
}
|
|
for _, tt := range tests {
|
|
t.Run(tt.name, func(t *testing.T) {
|
|
_, err := tc.client.DeleteActor(context.Background(), tt.req)
|
|
assertGrpcError(t, err, codes.InvalidArgument, tt.wantMsg)
|
|
})
|
|
}
|
|
})
|
|
}
|
|
|
|
func TestResumeActor_LockConflict(t *testing.T) {
|
|
ns := namespaceForTest("ns-resume-conflict")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
createWorkerPod(t, tc, ns, "worker-1", "node1")
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
id := "id1"
|
|
|
|
// Set a delay on the fake Atelet to hold the lock
|
|
tc.fakeAtelet.RestoreDelay = 1 * time.Second
|
|
|
|
// Launch Request A in a goroutine
|
|
errChan := make(chan error, 1)
|
|
go func() {
|
|
_, err := tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
errChan <- err
|
|
}()
|
|
|
|
// Sleep a bit to ensure Request A acquired the lock
|
|
time.Sleep(200 * time.Millisecond)
|
|
|
|
// Launch Request B (should fail due to lock conflict)
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
assertGrpcError(t, err, codes.Aborted, "another operation is in progress for this actor")
|
|
|
|
// Wait for Request A to finish
|
|
if errA := <-errChan; errA != nil {
|
|
t.Fatalf("Request A failed: %v", errA)
|
|
}
|
|
}
|
|
|
|
func TestResumeActor_DanglingWorker(t *testing.T) {
|
|
ns := namespaceForTest("ns-resume-dangling")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
// 1. Create Worker Pod A
|
|
createWorkerPod(t, tc, ns, "worker-a", "node1")
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
id := "id1"
|
|
|
|
// 2. Configure fake Atelet to FAIL on Restore!
|
|
tc.fakeAtelet.FailRestore = fmt.Errorf("mock atelet failure")
|
|
|
|
// 3. Call ResumeActor -> Expect failure
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err == nil {
|
|
t.Fatalf("expected ResumeActor to fail due to atelet error")
|
|
}
|
|
|
|
// Verify actor state is RESUMING with worker A assigned
|
|
getResp, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("GetActor failed: %v", err)
|
|
}
|
|
actor := getResp.GetActor()
|
|
if actor.GetStatus() != ateapipb.Actor_STATUS_RESUMING {
|
|
t.Fatalf("expected status RESUMING, got %v", actor.GetStatus())
|
|
}
|
|
if actor.GetAteomPodName() != "worker-a" {
|
|
t.Fatalf("expected worker-a assigned, got %v", actor.GetAteomPodName())
|
|
}
|
|
|
|
deleteWorkerPod(t, tc, ns, "worker-a")
|
|
|
|
// 6. Create Worker Pod B
|
|
createWorkerPod(t, tc, ns, "worker-b", "node1")
|
|
|
|
// 7. Configure fake Atelet to SUCCEED on Restore
|
|
tc.fakeAtelet.FailRestore = nil
|
|
tc.fakeAtelet.RestoreCalled = false // reset
|
|
|
|
// 8. Call ResumeActor again -> Expect success and picking Worker B!
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("ResumeActor failed on retry: %v", err)
|
|
}
|
|
|
|
if !tc.fakeAtelet.RestoreCalled {
|
|
t.Errorf("expected Restore to be called on retry")
|
|
}
|
|
|
|
// Verify actor state is RUNNING with worker B assigned
|
|
actor, err = tc.persistence.GetActor(context.Background(), id)
|
|
if err != nil {
|
|
t.Fatalf("failed to get actor from store: %v", err)
|
|
}
|
|
if actor.GetStatus() != ateapipb.Actor_STATUS_RUNNING {
|
|
t.Errorf("expected status RUNNING, got %v", actor.GetStatus())
|
|
}
|
|
if actor.GetAteomPodName() != "worker-b" {
|
|
t.Errorf("expected worker-b assigned, got %v", actor.GetAteomPodName())
|
|
}
|
|
}
|
|
|
|
func TestSuspendActor_DanglingWorker(t *testing.T) {
|
|
ns := namespaceForTest("ns-sd")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
// 1. Create Worker Pod
|
|
createWorkerPod(t, tc, ns, "worker-1", "node1")
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
id := "id1"
|
|
|
|
// Resume first to make it running
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("ResumeActor failed: %v", err)
|
|
}
|
|
|
|
deleteWorkerPod(t, tc, ns, "worker-1")
|
|
|
|
// 3. Call SuspendActor -> Should succeed (our fix skips missing pod execution)
|
|
actors, _, _ := tc.persistence.ListActors(context.Background(), maxPageSize, "")
|
|
t.Logf("Actors in Redis before Suspend: %d", len(actors))
|
|
for _, a := range actors {
|
|
t.Logf(" Actor: %s/%s/%s", a.GetActorTemplateNamespace(), a.GetActorTemplateName(), a.GetActorId())
|
|
}
|
|
|
|
_, err = tc.client.SuspendActor(context.Background(), &ateapipb.SuspendActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("SuspendActor failed: %v", err)
|
|
}
|
|
|
|
// 4. Verify it becomes SUSPENDED in Redis
|
|
getResp, err := tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{
|
|
ActorId: id,
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("GetActor failed: %v", err)
|
|
}
|
|
if getResp.GetActor().GetStatus() != ateapipb.Actor_STATUS_SUSPENDED {
|
|
t.Errorf("expected status SUSPENDED, got %v", getResp.GetActor().GetStatus())
|
|
}
|
|
if getResp.GetActor().GetAteomPodNamespace() != "" {
|
|
t.Errorf("expected ateom_pod_namespace to be empty, got %v", getResp.GetActor().GetAteomPodNamespace())
|
|
}
|
|
}
|
|
|
|
func TestDeleteActor_Success(t *testing.T) {
|
|
ns := namespaceForTest("ns-delete-success")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
|
|
_, err = tc.client.DeleteActor(context.Background(), &ateapipb.DeleteActorRequest{
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("DeleteActor failed: %v", err)
|
|
}
|
|
|
|
_, err = tc.client.GetActor(context.Background(), &ateapipb.GetActorRequest{
|
|
ActorId: "id1",
|
|
})
|
|
assertGrpcError(t, err, codes.NotFound, "Actor id1 not found")
|
|
}
|
|
|
|
func TestDeleteActor_NotSuspended(t *testing.T) {
|
|
ns := namespaceForTest("ns-delete-notsuspended")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
createTemplate(t, tc, ns)
|
|
createWorkerPod(t, tc, ns, "worker-1", "node1")
|
|
|
|
_, err := tc.client.CreateActor(context.Background(), &ateapipb.CreateActorRequest{
|
|
ActorTemplateNamespace: ns,
|
|
ActorTemplateName: "tmpl1",
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("CreateActor failed: %v", err)
|
|
}
|
|
|
|
_, err = tc.client.ResumeActor(context.Background(), &ateapipb.ResumeActorRequest{
|
|
ActorId: "id1",
|
|
})
|
|
if err != nil {
|
|
t.Fatalf("ResumeActor failed: %v", err)
|
|
}
|
|
|
|
_, err = tc.client.DeleteActor(context.Background(), &ateapipb.DeleteActorRequest{
|
|
ActorId: "id1",
|
|
})
|
|
assertGrpcError(t, err, codes.FailedPrecondition, "Actor id1 is not suspended (status: STATUS_RUNNING)")
|
|
}
|
|
|
|
func TestDeleteActor_NotFound(t *testing.T) {
|
|
ns := namespaceForTest("ns-delete-notfound")
|
|
tc := setupTest(t, ns)
|
|
defer tc.cleanup()
|
|
|
|
_, err := tc.client.DeleteActor(context.Background(), &ateapipb.DeleteActorRequest{
|
|
ActorId: "non-existent",
|
|
})
|
|
assertGrpcError(t, err, codes.NotFound, "Actor non-existent not found")
|
|
}
|
|
|
|
func assertGrpcError(t *testing.T, err error, wantCode codes.Code, wantMsg string) {
|
|
t.Helper()
|
|
if err == nil {
|
|
t.Fatalf("expected error, got nil")
|
|
}
|
|
st, ok := status.FromError(err)
|
|
if !ok {
|
|
t.Fatalf("expected gRPC status error, got: %v", err)
|
|
}
|
|
if st.Code() != wantCode {
|
|
t.Errorf("expected status %v, got %v", wantCode, st.Code())
|
|
}
|
|
if st.Message() != wantMsg {
|
|
t.Errorf("expected message %q, got %q", wantMsg, st.Message())
|
|
}
|
|
}
|