ateom: add GetWorkloadStats RPC and retain actor attribution

cmd/ateom-gvisor/main.go

@@ -33,6 +33,7 @@ import (
 	"github.com/agent-substrate/substrate/internal/ateinterceptors"
 	"github.com/agent-substrate/substrate/internal/ateomnet"
 	"github.com/agent-substrate/substrate/internal/ateompath"
+	"github.com/agent-substrate/substrate/internal/ateomstats"
 	"github.com/agent-substrate/substrate/internal/atunnel"
 	"github.com/agent-substrate/substrate/internal/contextlogging"
 	"github.com/agent-substrate/substrate/internal/imagecache"
@@ -48,7 +49,9 @@ import (
 	sdktrace "go.opentelemetry.io/otel/sdk/trace"
 	"golang.org/x/sys/unix"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/reflection"
+	"google.golang.org/grpc/status"
 )
 
 var (
@@ -239,6 +242,17 @@ type AteomService struct {
 	atunnelEgressPort        uint16
 	atunnelCredentialBundle  string
 	atunnelEgressTrustBundle string
+
+	// activeActor is the actor whose workload this ateom is currently running,
+	// or nil when it is "available". An ateom serves one actor at a time, so a
+	// single slot is enough (the micro-VM ateom holds the same field on each
+	// entry of its s.running map). Guarded by lock, like every other RPC-visible
+	// field.
+	//
+	// Set by RunWorkload / RestoreWorkload and cleared by CheckpointWorkload, so
+	// it tracks exactly the available/executing state machine described on the
+	// Ateom service. GetWorkloadStats reads it to attribute its sample.
+	activeActor *ateomstats.ActorAttribution
 }
 
 var _ ateompb.AteomServer = (*AteomService)(nil)
@@ -266,6 +280,12 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload
 	actorRef := resources.ActorRef{Atespace: req.GetAtespace(), Name: req.GetActorName()}
 	s.actorLogger.EmitLifecycleLog("Actor starting", actorRef, req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName())
 
+	// Retain the attribution before the boot rather than after it, so a sample
+	// taken against a workload that dies mid-boot is still attributable. The
+	// cleanup below drops it again if the boot fails outright.
+	attribution := ateomstats.ActorAttributionFromRequest(req)
+	s.activeActor = &attribution
+
 	// Contract with atelet:
 	//
 	//   * Correct runsc version is downloaded and placed on disk.
@@ -276,10 +296,14 @@ func (s *AteomService) RunWorkload(ctx context.Context, req *ateompb.RunWorkload
 		DumpNetInfo:        true,
 		EgressRedirectPort: s.egressRedirectPort(req.GetEgressGatewayAddress() != ""),
 	}); err != nil {
+		// Cleared here as well as in the deferred cleanup below, because that
+		// defer is not registered until after this check.
+		s.activeActor = nil
 		return nil, fmt.Errorf("while setting up actor network: %w", err)
 	}
 	defer func() {
 		if retErr != nil {
+			s.activeActor = nil
 			// Detach any bundle rootfs overlays a partially-completed setup
 			// mounted, mirroring the post-checkpoint cleanup — otherwise they
 			// linger in this namespace until atelet wipes the bundle dirs.
@@ -396,6 +420,19 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec
 		return nil, fmt.Errorf("unsupported snapshot scope: %v", req.GetScope())
 	}
 
+	// The sandbox is gone as of the checkpoint above, so the ateom is back to
+	// "available" from here on: there is nothing left to measure, and holding
+	// the attribution would let a later GetWorkloadStats report a checkpointed
+	// actor as though it were still running.
+	//
+	// Cleared here rather than at the end of the function because everything
+	// below is bookkeeping over a dead sandbox and can still fail (listing the
+	// snapshot files returns an error), which would otherwise leave the
+	// attribution behind. Conversely nothing above this point clears it: a
+	// checkpoint that failed may well have left the workload running, and
+	// reporting its usage is then the honest answer.
+	s.activeActor = nil
+
 	// After checkpointing the sandbox root, runsc may no longer have a usable
 	// control server for state/delete calls. Keep this as best-effort cleanup:
 	// atelet resets the actor runsc, bundle, pidfile, and checkpoint
@@ -433,6 +470,16 @@ func (s *AteomService) CheckpointWorkload(ctx context.Context, req *ateompb.Chec
 	return &ateompb.CheckpointWorkloadResponse{SnapshotFiles: snapshotFiles}, nil
 }
 
+// GetWorkloadStats implements ateompb.Ateom/GetWorkloadStats.
+//
+// The attribution half is wired up here; the measurement half is not. Reading the
+// sandbox's cgroup (/sys/fs/cgroup/pause) lands in the follow-up to
+// https://github.com/agent-substrate/substrate/issues/594, at which point this
+// stops returning Unimplemented.
+func (s *AteomService) GetWorkloadStats(ctx context.Context, req *ateompb.GetWorkloadStatsRequest) (*ateompb.GetWorkloadStatsResponse, error) {
+	return nil, status.Error(codes.Unimplemented, "GetWorkloadStats is not implemented yet")
+}
+
 // listSnapshotFiles returns the (relative) names of regular files directly under
 // dir, which atelet ships to object storage as the snapshot.
 func listSnapshotFiles(dir string) ([]string, error) {
@@ -486,6 +533,10 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore
 	actorRef := resources.ActorRef{Atespace: req.GetAtespace(), Name: req.GetActorName()}
 	s.actorLogger.EmitLifecycleLog("Actor restoring", actorRef, req.GetActorUid(), req.GetActorTemplateNamespace(), req.GetActorTemplateName())
 
+	// Same as RunWorkload: retain before the boot, drop again if it fails.
+	attribution := ateomstats.ActorAttributionFromRequest(req)
+	s.activeActor = &attribution
+
 	// Contract with atelet:
 	//
 	//   * Correct runsc version is downloaded and placed on disk.
@@ -497,10 +548,13 @@ func (s *AteomService) RestoreWorkload(ctx context.Context, req *ateompb.Restore
 		DumpNetInfo:        true,
 		EgressRedirectPort: s.egressRedirectPort(req.GetEgressGatewayAddress() != ""),
 	}); err != nil {
+		// Same as the Run path: the defer below is not registered yet.
+		s.activeActor = nil
 		return nil, fmt.Errorf("while setting up actor network: %w", err)
 	}
 	defer func() {
 		if retErr != nil {
+			s.activeActor = nil
 			// Same overlay detach as the Run-failure path above.
 			if err := imagecache.UnmountAllUnder(ateompath.OCIBundleDir(req.GetActorUid())); err != nil {
 				slog.WarnContext(ctx, "Failed to unmount bundle rootfs overlays after Restore failure",

cmd/ateom-gvisor/stats_test.go

@@ -0,0 +1,59 @@
+//go:build linux
+
+// 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 main
+
+import (
+	"context"
+	"testing"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/agent-substrate/substrate/internal/proto/ateompb"
+)
+
+// TestGetWorkloadStatsUnimplemented pins the stub's advertised contract. The
+// cgroup read replaces this body; until then a caller that gets any other code back
+// would be reading numbers that are not there.
+//
+// The retention this stub will eventually read — s.activeActor, set by
+// RunWorkload / RestoreWorkload and cleared by CheckpointWorkload — has no unit
+// test, because those three RPCs each reach for netlink, runsc, and the worker
+// pod's netns within a few lines of entry and cannot be driven from `go test`.
+// Its mapping is covered in internal/ateomstats; the transitions are verified
+// end to end once GetWorkloadStats returns real data.
+func TestGetWorkloadStatsUnimplemented(t *testing.T) {
+	s := &AteomService{}
+
+	resp, err := s.GetWorkloadStats(context.Background(), &ateompb.GetWorkloadStatsRequest{ActorUid: "uid-c"})
+	if resp != nil {
+		t.Errorf("GetWorkloadStats() returned response %v, want nil", resp)
+	}
+	if got := status.Code(err); got != codes.Unimplemented {
+		t.Errorf("GetWorkloadStats() error code = %v, want %v (err: %v)", got, codes.Unimplemented, err)
+	}
+}
+
+// TestAteomServiceStartsAvailable checks that a freshly constructed service
+// retains no attribution. GetWorkloadStats's FAILED_PRECONDITION-when-available
+// behavior is built on this: a non-nil zero value here would make an idle ateom
+// report an empty actor's usage instead of refusing.
+func TestAteomServiceStartsAvailable(t *testing.T) {
+	if s := (&AteomService{}); s.activeActor != nil {
+		t.Errorf("new AteomService.activeActor = %v, want nil", s.activeActor)
+	}
+}

cmd/ateom-microvm/restore.go

@@ -270,6 +270,7 @@ func (s *AteomService) restoreFullScope(ctx context.Context, p actorBootParams,
 	ra := &runningActor{
 		chCmd: chCmd, vfsdCmd: vfsdCmd, durableVfsdCmd: durableVfsdCmd,
 		apiSocket: apiSocket, baseID: srcID, restoreSourceDir: restoreDir,
+		activeActor: p.actorAttribution(),
 	}
 
 	// Re-attach stdout/stderr forwarding for each container: the restored guest's

cmd/ateom-microvm/run.go

@@ -34,6 +34,7 @@ import (
 	"github.com/agent-substrate/substrate/cmd/ateom-microvm/internal/kata"
 	"github.com/agent-substrate/substrate/cmd/ateom-microvm/internal/third_party/kata/agentpb"
 	"github.com/agent-substrate/substrate/internal/ateompath"
+	"github.com/agent-substrate/substrate/internal/ateomstats"
 	"github.com/agent-substrate/substrate/internal/imagecache"
 	"github.com/agent-substrate/substrate/internal/proto/ateompb"
 	"github.com/agent-substrate/substrate/internal/readyz"
@@ -55,6 +56,14 @@ type runningActor struct {
 	// base-id file so the chain survives suspend->resume->suspend.
 	baseID string
 
+	// activeActor is the actor from the Run/Restore request that started this
+	// micro-VM, retained so GetWorkloadStats can attribute its samples. The rest
+	// of this struct is about owning processes; this field is the one piece of
+	// the original request the service has to remember. Named to match the gVisor
+	// ateom's AteomService.activeActor, which holds the same thing for the one
+	// actor that ateom serves.
+	activeActor ateomstats.ActorAttribution
+
 	// ateom owns this CH process (booted at Run or relaunched at Restore).
 	chCmd *exec.Cmd
 	// vfsdCmd is the virtiofsd serving the overlay RO lower (the CH fs device
@@ -243,6 +252,17 @@ type actorBootParams struct {
 	egressGatewayAddress string
 }
 
+// actorAttribution regroups the actor fields that arrived on the Run/Restore
+// request, for retention on the resulting runningActor.
+func (p actorBootParams) actorAttribution() ateomstats.ActorAttribution {
+	return ateomstats.ActorAttribution{
+		Ref:               p.actorRef,
+		UID:               p.actorUID,
+		TemplateNamespace: p.templateNS,
+		TemplateName:      p.templateName,
+	}
+}
+
 // coldBootAttempts is how many times a cold boot is tried when the micro-VM
 // stops before the kata-agent answers. Two: one retry covers a transient guest
 // death (a contended host makes the guest's boot pathologically slow, and a
@@ -458,7 +478,7 @@ func (s *AteomService) coldBootActor(ctx context.Context, p actorBootParams) (re
 		return fmt.Errorf("while waiting for container readyz: %w", err)
 	}
 
-	ra := &runningActor{chCmd: chCmd, vfsdCmd: vfsdCmd, durableVfsdCmd: durableVfsdCmd, apiSocket: apiSocket, baseID: actorUID, logAgent: ac}
+	ra := &runningActor{chCmd: chCmd, vfsdCmd: vfsdCmd, durableVfsdCmd: durableVfsdCmd, apiSocket: apiSocket, baseID: actorUID, logAgent: ac, activeActor: p.actorAttribution()}
 	if err := s.activateActorNetworking(p.actorRef.Atespace, p.actorRef.Name, p.actorVersion, p.egressGatewayAddress); err != nil {
 		return err
 	}

cmd/ateom-microvm/stats.go

@@ -0,0 +1,38 @@
+//go:build linux
+
+// 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 main
+
+import (
+	"context"
+
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/agent-substrate/substrate/internal/proto/ateompb"
+)
+
+// GetWorkloadStats implements ateompb.Ateom/GetWorkloadStats.
+//
+// The attribution half is wired up (see runningActor.activeActor); the measurement
+// half is not. The micro-VM read goes to the guest agent's StatsContainer
+// rather than the host cgroup — guest RAM is a fixed allocation, so the host
+// cgroup barely moves with the workload — and lands in the follow-up to
+// https://github.com/agent-substrate/substrate/issues/594, at which point this
+// stops returning Unimplemented.
+func (s *AteomService) GetWorkloadStats(ctx context.Context, req *ateompb.GetWorkloadStatsRequest) (*ateompb.GetWorkloadStatsResponse, error) {
+	return nil, status.Error(codes.Unimplemented, "GetWorkloadStats is not implemented yet")
+}

cmd/ateom-microvm/stats_test.go

@@ -0,0 +1,113 @@
+//go:build linux
+
+// 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 main
+
+import (
+	"context"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+
+	"github.com/agent-substrate/substrate/internal/ateomstats"
+	"github.com/agent-substrate/substrate/internal/proto/ateompb"
+	"github.com/agent-substrate/substrate/internal/resources"
+)
+
+// TestActorBootParamsAttribution pins the mapping from actorBootParams onto the
+// attribution a runningActor retains. The three loose fields are the ones most
+// likely to get crossed, since actorBootParams names them differently than
+// ActorAttribution does; the distinct placeholders below make a swap visible.
+func TestActorBootParamsAttribution(t *testing.T) {
+	tests := []struct {
+		name string
+		p    actorBootParams
+		want ateomstats.ActorAttribution
+	}{
+		{
+			name: "fully populated",
+			p: actorBootParams{
+				actorRef:     resources.ActorRef{Atespace: "atespace-a", Name: "actor-b"},
+				actorUID:     "uid-c",
+				templateNS:   "template-ns-d",
+				templateName: "template-name-e",
+			},
+			want: ateomstats.ActorAttribution{
+				Ref:               resources.ActorRef{Atespace: "atespace-a", Name: "actor-b"},
+				UID:               "uid-c",
+				TemplateNamespace: "template-ns-d",
+				TemplateName:      "template-name-e",
+			},
+		},
+		{
+			name: "zero params",
+			p:    actorBootParams{},
+			want: ateomstats.ActorAttribution{},
+		},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			if diff := cmp.Diff(tc.want, tc.p.actorAttribution()); diff != "" {
+				t.Errorf("actorBootParams.actorAttribution() mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+// TestActorBootParamsAttributionMatchesRequest checks the two hops the
+// attribution makes — request to actorBootParams (in RunWorkload /
+// RestoreWorkload) and actorBootParams to ActorAttribution — compose back into
+// what the caller sent. The two hops are written in different files, so this is
+// the assertion that catches them drifting apart.
+func TestActorBootParamsAttributionMatchesRequest(t *testing.T) {
+	req := &ateompb.RunWorkloadRequest{
+		Atespace:               "atespace-a",
+		ActorName:              "actor-b",
+		ActorUid:               "uid-c",
+		ActorTemplateNamespace: "template-ns-d",
+		ActorTemplateName:      "template-name-e",
+	}
+
+	// Mirrors the actorBootParams literal in RunWorkload and RestoreWorkload.
+	p := actorBootParams{
+		actorRef:     resources.ActorRef{Atespace: req.GetAtespace(), Name: req.GetActorName()},
+		actorUID:     req.GetActorUid(),
+		templateNS:   req.GetActorTemplateNamespace(),
+		templateName: req.GetActorTemplateName(),
+	}
+
+	if diff := cmp.Diff(ateomstats.ActorAttributionFromRequest(req), p.actorAttribution()); diff != "" {
+		t.Errorf("attribution via actorBootParams differs from attribution via request (-request +params):\n%s", diff)
+	}
+}
+
+// TestGetWorkloadStatsUnimplemented pins the stub's advertised contract. The
+// guest agent read replaces this body; until then a caller that gets any other code
+// back would be reading numbers that are not there.
+func TestGetWorkloadStatsUnimplemented(t *testing.T) {
+	s := &AteomService{}
+
+	resp, err := s.GetWorkloadStats(context.Background(), &ateompb.GetWorkloadStatsRequest{ActorUid: "uid-c"})
+	if resp != nil {
+		t.Errorf("GetWorkloadStats() returned response %v, want nil", resp)
+	}
+	if got := status.Code(err); got != codes.Unimplemented {
+		t.Errorf("GetWorkloadStats() error code = %v, want %v (err: %v)", got, codes.Unimplemented, err)
+	}
+}