substrate/cmd/ateapi/internal/controlapi/functionaltest/worker_contention_test.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 functionaltest

import (
	"context"
	"fmt"
	"sync"
	"testing"
	"time"

	"k8s.io/apimachinery/pkg/util/wait"

	"go.ate.dev/substrate/api/ateapi/v1"
	"go.ate.dev/substrate/cmd/ateapi/internal/store"
	"go.ate.dev/substrate/internal/hardware"
	"google.golang.org/protobuf/proto"
)

// TestResumeActor_ConcurrentOntoOneWorker is the regression test for a refusal
// that only concurrency produces, and that a worker hosting one actor could
// never have shown.
//
// Claiming a worker rewrites that worker's whole record, so actors activating
// onto the SAME worker are N writers compare-and-swapping one row. All but one
// lose each round. That is fine and expected -- the loser re-reads and tries
// again -- but only if the retry budget is sized for how many writers there
// are. It was five steps from 10ms, and on a real cluster that rejected 21% of
// activations at 12 in flight and 67% at 24, every one of them failing at
// ~235ms with "timed out waiting for the condition".
//
// Nothing sequential catches it: the same activations one at a time all
// succeed, which is exactly what the pre-existing resume tests do.
func TestResumeActor_ConcurrentOntoOneWorker(t *testing.T) {
	const actors = 128

	ns := namespaceForTest("ns-resume-contention")
	tc := setupTest(t, ns)
	defer tc.cleanup()

	// createTemplate makes pool1 for us; widen it so one worker has room for all
	// of them. The point is that they contend, not that they are turned away for
	// lack of capacity.
	createTemplate(t, tc, ns)
	createWorkerPod(t, tc, ns, "worker-1", "node1", "pool1")
	setWorkerActorCapacity(t, tc, "pool1", actors)

	for i := range actors {
		if _, err := tc.client.CreateActor(context.Background(), &ateapi.CreateActorRequest{Actor: &ateapi.Actor{
			Metadata:      &ateapi.ResourceMetadata{Atespace: testAtespace, Name: fmt.Sprintf("id%d", i)},
			ActorTemplate: &ateapi.ObjectRef{Atespace: testAtespace, Name: "tmpl1"},
		}}); err != nil {
			t.Fatalf("CreateActor %d: %v", i, err)
		}
	}

	// Resume them all at once. Starting the goroutines is not enough to make
	// them race -- release them together so they arrive at the claim inside the
	// same window.
	var start sync.WaitGroup
	start.Add(1)
	var wg sync.WaitGroup
	errs := make([]error, actors)
	for i := range actors {
		wg.Go(func() {
			start.Wait()
			_, errs[i] = tc.client.ResumeActor(context.Background(), &ateapi.ResumeActorRequest{
				Actor: &ateapi.ObjectRef{Atespace: testAtespace, Name: fmt.Sprintf("id%d", i)},
			})
		})
	}
	start.Done()
	wg.Wait()

	var failed int
	for i, err := range errs {
		if err != nil {
			failed++
			if failed <= 3 {
				t.Errorf("ResumeActor id%d: %v", i, err)
			}
		}
	}
	if failed > 0 {
		t.Fatalf("%d of %d concurrent activations onto one worker were refused; "+
			"they contend on the worker record and the retry budget has to absorb that", failed, actors)
	}

	// And every one of them is really on the worker: a claim that was lost but
	// reported as won would show up here rather than as an error above. Found by
	// pod rather than by taking the only worker in the list -- the list is
	// process-wide, so a -count>1 run sees the workers earlier iterations made.
	listed, err := tc.client.ListWorkers(context.Background(), &ateapi.ListWorkersRequest{})
	if err != nil {
		t.Fatalf("ListWorkers: %v", err)
	}
	var worker *ateapi.Worker
	for _, candidate := range listed.GetWorkers() {
		if candidate.GetWorkerNamespace() == ns && candidate.GetWorkerPod() == "worker-1" {
			worker = candidate
			break
		}
	}
	if worker == nil {
		t.Fatalf("worker %s/worker-1 not in the list of %d", ns, len(listed.GetWorkers()))
	}
	// A listing reports occupancy through the allocation total; the assignments
	// themselves are their own records.
	if got := int(worker.GetStatus().GetAllocated().GetActors()); got != actors {
		t.Errorf("worker allocation counts %d actors, want %d", got, actors)
	}
	page, err := tc.persistence.ListWorkerAssignments(context.Background(), worker.GetMetadata().GetName(), store.ListOptions{})
	if err != nil {
		t.Fatalf("ListWorkerAssignments: %v", err)
	}
	if len(page.Items) != actors {
		t.Errorf("worker holds %d assignments, want %d", len(page.Items), actors)
	}
}

// setWorkerActorCapacity raises the actors ceiling on every Worker in the pool.
//
// The ceiling is the Worker's, reported by its worker, so a test that wants more
// than the unset default of one writes it where the reporter would. Waits for
// the scheduler's cache to see it, since placement reads that and not the store.
func setWorkerActorCapacity(t *testing.T, tc *testContext, pool string, actors int32) {
	t.Helper()
	page, err := tc.persistence.ListWorkers(context.Background(), store.ListOptions{PageSize: 1000})
	if err != nil {
		t.Fatalf("listing workers: %v", err)
	}
	var reported int
	for _, w := range page.Items {
		if w.GetWorkerPool() != pool {
			continue
		}
		reportWorkerCapacity(t, tc, w.GetMetadata().GetName(), actors)
		reported++
	}
	if reported == 0 {
		t.Fatalf("no workers in pool %q to give capacity", pool)
	}
}

// reportedEpoch is the epoch reportWorkerCapacity registers, and the one the
// fake atelet reports its workloads ran under.
const reportedEpoch = 1

// defaultWorkerRuntime returns the default ActorSandboxRuntime that
// reportWorkerCapacity records on a test worker of the given sandbox class.
func defaultWorkerRuntime(class string) *ateapi.ActorSandboxRuntime {
	if class == "" {
		class = "gvisor"
	}
	return &ateapi.ActorSandboxRuntime{
		SandboxClass: class,
		Version: &ateapi.VersionedSandboxCompat{
			SchemaVersion: hardware.SchemaVersionV1,
			Attributes: []*ateapi.AttributeEntry{
				{Key: hardware.AttrArchitecture, Value: "amd64"},
				{Key: hardware.AttrSandboxBundleHash, Value: "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855"},
			},
		},
	}
}

// workerSandboxClasses is the sandbox_classes a worker reports to run def as
// its default runtime and restore restorable beyond it, all of def's class.
// The arguments are the ActorSandboxRuntimes a test also expects on snapshots and
// assignments, so both sides of a match derive from one value; the compats
// are copied so the stored row never aliases them.
func workerSandboxClasses(def *ateapi.ActorSandboxRuntime, restorable ...*ateapi.ActorSandboxRuntime) []*ateapi.WorkerSandboxClass {
	class := &ateapi.WorkerSandboxClass{SandboxClass: def.GetSandboxClass(), DefaultRuntime: &ateapi.WorkerSandboxRuntime{Version: proto.CloneOf(def.GetVersion())}}
	for _, r := range restorable {
		class.RestorableRuntimes = append(class.RestorableRuntimes, &ateapi.WorkerSandboxRuntime{Version: proto.CloneOf(r.GetVersion())})
	}
	return []*ateapi.WorkerSandboxClass{class}
}

// reportWorkerCapacity stands in for the worker process's registration and
// runtimes report, which atelet forwards to WorkerService in a real cluster,
// with compute for actors of the size test templates declare. A Worker not yet
// registered gets reportedEpoch, an IP and a default runtime; one already
// registered keeps them. It waits for the worker cache, which placement reads,
// to catch up.
func reportWorkerCapacity(t *testing.T, tc *testContext, name string, actors int32) {
	t.Helper()
	ctx := context.Background()
	worker, err := tc.persistence.GetWorker(ctx, name)
	if err != nil {
		t.Fatalf("getting worker %s: %v", name, err)
	}
	if _, err := tc.persistence.UpdateWorker(ctx, name, store.PreconditionFrom(worker), func(toUpdate *ateapi.Worker) error {
		if toUpdate.Status.RegisteredEpoch == 0 {
			toUpdate.Status.RegisteredEpoch = reportedEpoch
			toUpdate.Status.Ips = []string{"127.0.0.1"}
		}
		toUpdate.Status.Capacity = &ateapi.WorkerResources{
			Actors:        actors,
			CpuMillicores: int64(actors) * 1000, MemoryBytes: int64(actors) * 512 << 20,
		}
		if len(toUpdate.Status.SandboxClasses) == 0 {
			toUpdate.Status.SandboxClasses = workerSandboxClasses(defaultWorkerRuntime("gvisor"))
		}
		return nil
	}); err != nil {
		t.Fatalf("setting capacity on worker %s: %v", name, err)
	}
	if err := wait.PollUntilContextTimeout(ctx, 10*time.Millisecond, 5*time.Second, true,
		func(context.Context) (bool, error) {
			got, err := tc.workerCache.Worker(name)
			return err == nil && got.GetStatus().GetCapacity().GetActors() == actors, nil
		}); err != nil {
		t.Fatalf("worker %s did not reach capacity.actors=%d in the cache: %v", name, actors, err)
	}
}