multi-actor worker API

benchmarking/locust/common/ateapi_pb2_grpc.py

@@ -164,6 +164,11 @@ def __init__(self, channel):
                 request_serializer=ateapi__pb2.DrainWorkerRequest.SerializeToString,
                 response_deserializer=ateapi__pb2.Worker.FromString,
                 _registered_method=True)
+        self.ListWorkerActorAssignments = channel.unary_unary(
+                '/ateapi.Control/ListWorkerActorAssignments',
+                request_serializer=ateapi__pb2.ListWorkerActorAssignmentsRequest.SerializeToString,
+                response_deserializer=ateapi__pb2.ListWorkerActorAssignmentsResponse.FromString,
+                _registered_method=True)
         self.ListActors = channel.unary_unary(
                 '/ateapi.Control/ListActors',
                 request_serializer=ateapi__pb2.ListActorsRequest.SerializeToString,
@@ -382,6 +387,14 @@ def DrainWorker(self, request, context):
         context.set_details('Method not implemented!')
         raise NotImplementedError('Method not implemented!')
 
+    def ListWorkerActorAssignments(self, request, context):
+        """List the Actors a Worker hosts. A subresource of Worker rather than a field
+        on it, so GetWorker and ListWorkers cost the same whatever the occupancy.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
     def ListActors(self, request, context):
         """List Actors.
         """
@@ -562,6 +575,11 @@ def add_ControlServicer_to_server(servicer, server):
                     request_deserializer=ateapi__pb2.DrainWorkerRequest.FromString,
                     response_serializer=ateapi__pb2.Worker.SerializeToString,
             ),
+            'ListWorkerActorAssignments': grpc.unary_unary_rpc_method_handler(
+                    servicer.ListWorkerActorAssignments,
+                    request_deserializer=ateapi__pb2.ListWorkerActorAssignmentsRequest.FromString,
+                    response_serializer=ateapi__pb2.ListWorkerActorAssignmentsResponse.SerializeToString,
+            ),
             'ListActors': grpc.unary_unary_rpc_method_handler(
                     servicer.ListActors,
                     request_deserializer=ateapi__pb2.ListActorsRequest.FromString,
@@ -1240,6 +1258,33 @@ def DrainWorker(request,
             metadata,
             _registered_method=True)
 
+    @staticmethod
+    def ListWorkerActorAssignments(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/ateapi.Control/ListWorkerActorAssignments',
+            ateapi__pb2.ListWorkerActorAssignmentsRequest.SerializeToString,
+            ateapi__pb2.ListWorkerActorAssignmentsResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)
+
     @staticmethod
     def ListActors(request,
             target,
@@ -1630,3 +1675,94 @@ def MintCert(request,
             timeout,
             metadata,
             _registered_method=True)
+
+
+class WorkerServiceStub:
+    """WorkerService is how a Worker tells the control plane about itself. It is
+    separate from Control because the two have different callers and different
+    authorization: Control is the client-facing API, while these RPCs are served
+    only to an atelet, and only for the Workers on its own node.
+    """
+
+    def __init__(self, channel):
+        """Constructor.
+
+        Args:
+            channel: A grpc.Channel.
+        """
+        self.SetWorkerCapacity = channel.unary_unary(
+                '/ateapi.WorkerService/SetWorkerCapacity',
+                request_serializer=ateapi__pb2.SetWorkerCapacityRequest.SerializeToString,
+                response_deserializer=ateapi__pb2.SetWorkerCapacityResponse.FromString,
+                _registered_method=True)
+
+
+class WorkerServiceServicer:
+    """WorkerService is how a Worker tells the control plane about itself. It is
+    separate from Control because the two have different callers and different
+    authorization: Control is the client-facing API, while these RPCs are served
+    only to an atelet, and only for the Workers on its own node.
+    """
+
+    def SetWorkerCapacity(self, request, context):
+        """SetWorkerCapacity records what a Worker can hold. Capacity is the Worker's
+        to report rather than the control plane's to infer: it is what the ateom
+        can actually supply, only its node can observe it, and a fleet may run
+        mixed ateom versions.
+
+        atelet calls this with its own client certificate, as it does for
+        MintCert. Idempotent: re-sending the same capacity is not a write.
+        """
+        context.set_code(grpc.StatusCode.UNIMPLEMENTED)
+        context.set_details('Method not implemented!')
+        raise NotImplementedError('Method not implemented!')
+
+
+def add_WorkerServiceServicer_to_server(servicer, server):
+    rpc_method_handlers = {
+            'SetWorkerCapacity': grpc.unary_unary_rpc_method_handler(
+                    servicer.SetWorkerCapacity,
+                    request_deserializer=ateapi__pb2.SetWorkerCapacityRequest.FromString,
+                    response_serializer=ateapi__pb2.SetWorkerCapacityResponse.SerializeToString,
+            ),
+    }
+    generic_handler = grpc.method_handlers_generic_handler(
+            'ateapi.WorkerService', rpc_method_handlers)
+    server.add_generic_rpc_handlers((generic_handler,))
+    server.add_registered_method_handlers('ateapi.WorkerService', rpc_method_handlers)
+
+
+ # This class is part of an EXPERIMENTAL API.
+class WorkerService:
+    """WorkerService is how a Worker tells the control plane about itself. It is
+    separate from Control because the two have different callers and different
+    authorization: Control is the client-facing API, while these RPCs are served
+    only to an atelet, and only for the Workers on its own node.
+    """
+
+    @staticmethod
+    def SetWorkerCapacity(request,
+            target,
+            options=(),
+            channel_credentials=None,
+            call_credentials=None,
+            insecure=False,
+            compression=None,
+            wait_for_ready=None,
+            timeout=None,
+            metadata=None):
+        return grpc.experimental.unary_unary(
+            request,
+            target,
+            '/ateapi.WorkerService/SetWorkerCapacity',
+            ateapi__pb2.SetWorkerCapacityRequest.SerializeToString,
+            ateapi__pb2.SetWorkerCapacityResponse.FromString,
+            options,
+            channel_credentials,
+            insecure,
+            call_credentials,
+            compression,
+            wait_for_ready,
+            timeout,
+            metadata,
+            _registered_method=True)

cmd/ateapi/internal/actoridentity/actoridentity.go

@@ -26,6 +26,7 @@ import (
 	"path"
 	"time"
 
+	"github.com/agent-substrate/substrate/cmd/ateapi/internal/ateletauth"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/controlapi"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/store"
 	"github.com/agent-substrate/substrate/cmd/ateapi/internal/workercache"
@@ -37,8 +38,6 @@ import (
 	"github.com/agent-substrate/substrate/internal/substratex509"
 	"github.com/agent-substrate/substrate/pkg/proto/ateapipb"
 	"google.golang.org/grpc/codes"
-	"google.golang.org/grpc/credentials"
-	"google.golang.org/grpc/peer"
 	"google.golang.org/grpc/status"
 	"k8s.io/apimachinery/pkg/api/operation"
 	"k8s.io/apimachinery/pkg/util/validation/field"
@@ -73,19 +72,7 @@ func New(actorIdentityJWTIssuer string, actorIDJWTPool localjwtauthority.Pool, a
 	}
 }
 
-// The SPIFFE identity that atelet client certs carry, as minted by the
-// podidentity signer (cmd/podcertcontroller/internal/podidentitysigner).
-//
-// These mirror the constants the atelet dialer verifies against in
-// cmd/ateapi/internal/controlapi/dialer.go. They are duplicated rather than
-// imported so that this package does not depend on controlapi for three
-// strings; if a third pkg that need these constants appears, they should move to a shared package.
-const (
-	ateletTrustDomain        = "cluster.local"
-	ateletNamespace          = "ate-system"
-	ateletSA                 = "atelet"
-	actorCertificateLifetime = time.Hour
-)
+const actorCertificateLifetime = time.Hour
 
 func (s *Server) MintJWT(ctx context.Context, req *ateapipb.MintJWTRequest) (*ateapipb.MintJWTResponse, error) {
 	caller, ok := principal.FromContext(ctx)
@@ -136,7 +123,7 @@ func (s *Server) MintJWT(ctx context.Context, req *ateapipb.MintJWTRequest) (*at
 }
 
 func (s *Server) MintCert(ctx context.Context, req *ateapipb.MintCertRequest) (*ateapipb.MintCertResponse, error) {
-	caller, err := authenticateAtelet(ctx)
+	caller, err := ateletauth.Authenticate(ctx)
 	if err != nil {
 		return nil, err
 	}
@@ -154,8 +141,9 @@ func (s *Server) MintCert(ctx context.Context, req *ateapipb.MintCertRequest) (*
 	}
 	atespace, actorName := actorRef.Atespace, actorRef.Name
 
-	// Actor identity comes only from ateapi state. expected_actor_uid is a
-	// fail-closed guard against a request crossing an assignment change.
+	// expected_actor_uid picked which actor to mint for (see authorizeActor);
+	// re-checking it here fails closed if the request crossed an assignment
+	// change.
 	actorUID := actor.GetMetadata().GetUid()
 	if actorUID == "" {
 		slog.ErrorContext(ctx, "MintCert: actor has no UID", slog.Any("actor", actorRef))
@@ -218,62 +206,6 @@ func (s *Server) MintCert(ctx context.Context, req *ateapipb.MintCertRequest) (*
 	}, nil
 }
 
-// ateletCaller is the verified identity of an atelet requesting an actor credential.
-type ateletCaller struct {
-	podName  string
-	nodeName string
-}
-
-// authenticateAtelet verifies that the RPC arrived over mTLS from an atelet,
-// and returns the identity that atelet's certificate asserts.
-//
-// The certificate chain is already verified by the TLS layer against the
-// pod-identity CA (see buildServerCreds in cmd/ateapi/main.go), so the
-// extensions read here are trustworthy: only the pod-identity signer can mint
-// a certificate carrying a given pod's node name.
-func authenticateAtelet(ctx context.Context) (*ateletCaller, error) {
-	p, ok := peer.FromContext(ctx)
-	if !ok {
-		return nil, status.Errorf(codes.Unauthenticated, "no peer transport information found")
-	}
-
-	tlsInfo, ok := p.AuthInfo.(credentials.TLSInfo)
-	if !ok {
-		return nil, status.Errorf(codes.Unauthenticated, "unexpected peer transport credentials")
-	}
-
-	if len(tlsInfo.State.PeerCertificates) == 0 {
-		return nil, status.Errorf(codes.Unauthenticated, "could not verify peer certificate")
-	}
-	leaf := tlsInfo.State.PeerCertificates[0]
-
-	// Only atelet may mint actor credentials. Everything else with a valid
-	// pod-identity certificate — including the actor workloads themselves — is
-	// rejected here.
-	expected := (&url.URL{
-		Scheme: "spiffe",
-		Host:   ateletTrustDomain,
-		Path:   path.Join("ns", ateletNamespace, "sa", ateletSA),
-	}).String()
-	if len(leaf.URIs) == 0 || leaf.URIs[0].String() != expected {
-		slog.WarnContext(ctx, "ActorIdentity denied: caller is not atelet",
-			slog.Any("uris", leaf.URIs), slog.String("expected", expected))
-		return nil, status.Errorf(codes.PermissionDenied, "caller is not permitted to mint actor credentials")
-	}
-
-	identity, err := substratex509.PodIdentityFromCertificate(leaf)
-	if err != nil {
-		slog.WarnContext(ctx, "ActorIdentity denied: malformed PodIdentity extension", slog.Any("err", err))
-		return nil, status.Errorf(codes.PermissionDenied, "caller is not permitted to mint actor credentials")
-	}
-	if identity == nil {
-		slog.WarnContext(ctx, "ActorIdentity denied: certificate has no PodIdentity extension")
-		return nil, status.Errorf(codes.PermissionDenied, "caller is not permitted to mint actor credentials")
-	}
-
-	return &ateletCaller{podName: identity.PodName, nodeName: identity.NodeName}, nil
-}
-
 func validateMintJWTRequest(ctx context.Context, req *ateapipb.MintJWTRequest) field.ErrorList {
 	// Call the generated validation.
 	op := operation.Operation{Type: operation.Create}
@@ -286,13 +218,14 @@ func validateMintCertRequest(ctx context.Context, req *ateapipb.MintCertRequest)
 	return controlapi.Validate_MintCertRequest(ctx, op, nil, req, nil)
 }
 
-// authorizeActor resolves the actor from the authenticated worker and verifies
-// that the worker and actor still point at one another. Actor identity supplied
-// by the requester never participates in this authorization decision.
+// authorizeActor resolves the actor the request names among those the worker is
+// hosting and verifies the two still point at one another. Requester-supplied
+// identity never participates in the decision.
+//
 // The worker is resolved from cache first (hot path), but cache misses and
 // denials fall back to the authoritative store to handle watch-delivery lag
 // right after ResumeActor.
-func (s *Server) authorizeActor(ctx context.Context, caller *ateletCaller, req *ateapipb.MintCertRequest) (*ateapipb.Actor, resources.ActorRef, error) {
+func (s *Server) authorizeActor(ctx context.Context, caller *ateletauth.Caller, req *ateapipb.MintCertRequest) (*ateapipb.Actor, resources.ActorRef, error) {
 	reason := "worker not found"
 	worker, err := s.workers.Worker(req.GetWorker().GetName())
 	if err != nil && !errors.Is(err, store.ErrNotFound) {
@@ -337,24 +270,32 @@ func (s *Server) authorizeActor(ctx context.Context, caller *ateletCaller, req *
 // denyMint logs the internal reason and returns a uniform PermissionDenied.
 // Denials are deliberately indistinguishable from each other: a caller that
 // is not entitled to a worker should not learn its assignment.
-func (s *Server) denyMint(ctx context.Context, caller *ateletCaller, req *ateapipb.MintCertRequest, reason string, args ...any) error {
+func (s *Server) denyMint(ctx context.Context, caller *ateletauth.Caller, req *ateapipb.MintCertRequest, reason string, args ...any) error {
 	slog.WarnContext(ctx, "ActorIdentity denied: "+reason,
-		append([]any{slog.String("worker", req.GetWorker().GetName()), slog.String("callerPod", caller.podName), slog.String("callerNode", caller.nodeName)}, args...)...)
+		append([]any{slog.String("worker", req.GetWorker().GetName()), slog.String("callerPod", caller.PodName), slog.String("callerNode", caller.NodeName)}, args...)...)
 	return status.Error(codes.PermissionDenied, "caller is not permitted to mint credentials for this actor")
 }
 
 var errAssignmentMismatch = errors.New("assignment mismatch")
 
 // authorizeWithWorker returns errAssignmentMismatch and a reason string if the authorization failed
 // due to an assignment mismatch, indicating the caller may want to refetch the worker and retry.
-func (s *Server) authorizeWithWorker(ctx context.Context, worker *ateapipb.Worker, caller *ateletCaller, req *ateapipb.MintCertRequest) (*ateapipb.Actor, resources.ActorRef, string, error) {
-	if worker.GetNodeName() != caller.nodeName {
+func (s *Server) authorizeWithWorker(ctx context.Context, worker *ateapipb.Worker, caller *ateletauth.Caller, req *ateapipb.MintCertRequest) (*ateapipb.Actor, resources.ActorRef, string, error) {
+	if worker.GetNodeName() != caller.NodeName {
 		return nil, resources.ActorRef{}, "worker is hosted on a different node", errAssignmentMismatch
 	}
 
-	actorRef := resources.ActorRefFromObjectRef(worker.GetStatus().GetAssignment().GetActor())
+	assigned, err := s.assignmentToMintFor(ctx, worker.GetMetadata().GetName(), req.GetExpectedActorUid())
+	if errors.Is(err, store.ErrNotFound) {
+		return nil, resources.ActorRef{}, "worker is not hosting the requested actor", errAssignmentMismatch
+	}
+	if err != nil {
+		slog.ErrorContext(ctx, "ActorIdentity: failed to read worker assignment", slog.Any("err", err))
+		return nil, resources.ActorRef{}, "", status.Error(codes.Internal, "failed to look up worker assignment")
+	}
+	actorRef := resources.ActorRefFromObjectRef(assigned.GetActor())
 	if actorRef == (resources.ActorRef{}) {
-		return nil, resources.ActorRef{}, "worker has no actor assignment", errAssignmentMismatch
+		return nil, resources.ActorRef{}, "worker assignment names no actor", errAssignmentMismatch
 	}
 
 	actor, err := s.store.GetActor(ctx, actorRef)
@@ -378,11 +319,38 @@ func (s *Server) authorizeWithWorker(ctx context.Context, worker *ateapipb.Worke
 		slog.ErrorContext(ctx, "ActorIdentity: running actor has no worker assignment", slog.Any("actor", actorRef))
 		return nil, resources.ActorRef{}, "", status.Error(codes.FailedPrecondition, "actor has no worker assigned")
 	}
-	if worker.GetStatus().GetAssignment().GetActorUid() != actor.GetMetadata().GetUid() {
+	if assigned.GetActorUid() != actor.GetMetadata().GetUid() {
 		return nil, resources.ActorRef{}, "worker is no longer assigned to this actor incarnation", errAssignmentMismatch
 	}
 	if assignment.GetWorker().GetName() != worker.GetMetadata().GetName() {
 		return nil, resources.ActorRef{}, "actor no longer points to the requesting worker", errAssignmentMismatch
 	}
 	return actor, actorRef, "", nil
 }
+
+// assignmentToMintFor picks which of the worker's actors to mint for, or
+// ErrNotFound when it hosts none. expected_actor_uid selects from what ateapi
+// records the worker as hosting; it does not assert.
+//
+// Falling back to another of the worker's assignments keeps a bad binding
+// (PermissionDenied) apart from a stale expectation (retryable). Reads go to
+// the store: the binding was committed moments ago and the watch has not
+// delivered it. Only one other assignment is needed to tell the two apart, so
+// the fallback reads a single row rather than the Worker's whole occupancy.
+func (s *Server) assignmentToMintFor(ctx context.Context, workerName, actorUID string) (*ateapipb.ActorAssignment, error) {
+	assigned, err := s.store.GetWorkerAssignment(ctx, workerName, actorUID)
+	if err == nil {
+		return assigned, nil
+	}
+	if !errors.Is(err, store.ErrNotFound) {
+		return nil, err
+	}
+	page, err := s.store.ListWorkerAssignments(ctx, workerName, store.ListOptions{PageSize: 1})
+	if err != nil {
+		return nil, err
+	}
+	if len(page.Items) == 0 {
+		return nil, store.ErrNotFound
+	}
+	return page.Items[0], nil
+}