ateom: add GetActiveWorkloadStats RPC

cmd/ateom-gvisor/main.go

@@ -34,6 +34,7 @@ import (
 	"time"
 
 	"cloud.google.com/go/compute/metadata"
+	"github.com/agent-substrate/substrate/cmd/ateom-gvisor/internal/cgroupstats"
 	"github.com/agent-substrate/substrate/internal/actorlog"
 	"github.com/agent-substrate/substrate/internal/ateinterceptors"
 	"github.com/agent-substrate/substrate/internal/ateomnet"
@@ -365,6 +366,12 @@ type AteomService struct {
 	// own cgroup scope, which setupCgroupDelegation prepares. A field rather
 	// than a constant so tests can point GetWorkloadStats at a fixture tree.
 	cgroupRoot string
+
+	// readSandboxCgroup overrides cgroupstats.Read when set. Only tests set it:
+	// it is the seam that lets them interleave a lifecycle transition with the
+	// stats handlers' lock-free read, the way containerStatsReader does for the
+	// micro-VM runtime. nil means the real read.
+	readSandboxCgroup func(dir string) (cgroupstats.Sample, error)
 }
 
 var _ ateompb.AteomServer = (*AteomService)(nil)

cmd/ateom-gvisor/stats.go

@@ -27,6 +27,7 @@ import (
 	"google.golang.org/grpc/status"
 
 	"github.com/agent-substrate/substrate/cmd/ateom-gvisor/internal/cgroupstats"
+	"github.com/agent-substrate/substrate/internal/ateomstats"
 	"github.com/agent-substrate/substrate/internal/proto/ateompb"
 )
 
@@ -112,17 +113,16 @@ func (s *AteomService) GetWorkloadStats(ctx context.Context, req *ateompb.GetWor
 		return nil, status.Errorf(codes.NotFound, "ateom is executing actor %q, not the requested %q", active.UID, req.GetActorUid())
 	}
 
-	observedAt := time.Now()
-	sample, err := cgroupstats.Read(filepath.Join(s.cgroupRoot, sandboxCgroupContainer))
+	sample, err := s.sampleSandbox(active)
 	if err != nil {
-		// The requested actor is the active one but its cgroup is not there. Most
-		// often that is a poll landing in the boot: the ateom retains the
+		// The requested actor is the active one but its cgroup is not there.
+		// Most often that is a poll landing in the boot: the ateom retains the
 		// attribution from the moment it accepts the actor, before runsc has
-		// created the leaf. The other way in is a sandbox that went away between
-		// the check above and the read, which the next CheckpointWorkload turns
-		// into the NOT_FOUND above. Either way it is "no numbers right now" and the
-		// caller should take the next sample, so FAILED_PRECONDITION. Anything else
-		// is a real read failure.
+		// created the leaf. The other way in is a sandbox that went away
+		// underneath the read, which the next CheckpointWorkload turns into the
+		// NOT_FOUND above. Either way it is "no numbers right now" and the
+		// caller should take the next sample, so FAILED_PRECONDITION. Anything
+		// else is a real read failure.
 		if errors.Is(err, fs.ErrNotExist) {
 			return nil, status.Error(codes.FailedPrecondition, "no sandbox cgroup to measure yet")
 		}
@@ -144,7 +144,75 @@ func (s *AteomService) GetWorkloadStats(ctx context.Context, req *ateompb.GetWor
 		return nil, status.Errorf(codes.NotFound, "ateom stopped executing actor %q while the sample was being taken", req.GetActorUid())
 	}
 
-	return &ateompb.GetWorkloadStatsResponse{
+	return &ateompb.GetWorkloadStatsResponse{Sample: sample}, nil
+}
+
+// GetActiveWorkloadStats implements
+// ateompb.Ateom/GetActiveWorkloadStats: the discovery read, sampling
+// whatever is executing with no identity asserted. Same lock discipline as
+// GetWorkloadStats above, for the same reasons.
+func (s *AteomService) GetActiveWorkloadStats(ctx context.Context, req *ateompb.GetActiveWorkloadStatsRequest) (*ateompb.GetActiveWorkloadStatsResponse, error) {
+	active := s.activeActor.Load()
+	if active == nil {
+		return noSample(ateompb.NoSampleReason_NO_SAMPLE_REASON_NO_WORKLOAD), nil
+	}
+
+	sample, err := s.sampleSandbox(active)
+	if err != nil {
+		// A missing cgroup is a workload with no numbers yet -- a poll landing
+		// in the boot -- which for a caller with no prior knowledge is as
+		// normal a finding as an available ateom, so it is a reason, not an
+		// error. Anything else is a real read failure.
+		if errors.Is(err, fs.ErrNotExist) {
+			return noSample(ateompb.NoSampleReason_NO_SAMPLE_REASON_NOT_MEASURABLE_YET), nil
+		}
+		return nil, status.Errorf(codes.Internal, "reading sandbox cgroup: %v", err)
+	}
+
+	// Same re-check as GetWorkloadStats, different answer: with no uid asserted
+	// there is no "requested actor" for NOT_FOUND to disown, and a transition
+	// underneath the read just means these numbers cannot be attributed to any
+	// single actor. Report the reason as of now -- the next tick resolves it
+	// either way.
+	if latest := s.activeActor.Load(); latest != active {
+		reason := ateompb.NoSampleReason_NO_SAMPLE_REASON_NOT_MEASURABLE_YET
+		if latest == nil {
+			reason = ateompb.NoSampleReason_NO_SAMPLE_REASON_NO_WORKLOAD
+		}
+		return noSample(reason), nil
+	}
+
+	return &ateompb.GetActiveWorkloadStatsResponse{
+		Result: &ateompb.GetActiveWorkloadStatsResponse_Sample{Sample: sample},
+	}, nil
+}
+
+// noSample is the discovery read's "nothing to give, and that is normal"
+// answer.
+func noSample(reason ateompb.NoSampleReason) *ateompb.GetActiveWorkloadStatsResponse {
+	return &ateompb.GetActiveWorkloadStatsResponse{
+		Result: &ateompb.GetActiveWorkloadStatsResponse_NoSampleReason{NoSampleReason: reason},
+	}
+}
+
+// sampleSandbox reads the sandbox cgroup and builds the sample attributed to
+// active. Errors come back raw -- notably fs.ErrNotExist for a cgroup that is
+// not there yet -- because the two RPCs disagree on what that means: an error
+// code for the keyed read, a normal EXECUTING answer for the discovery read.
+// Callers re-check s.activeActor against the pointer they loaded after this
+// returns; the read holds no lock.
+func (s *AteomService) sampleSandbox(active *ateomstats.ActorAttribution) (*ateompb.WorkloadStatsSample, error) {
+	read := s.readSandboxCgroup
+	if read == nil {
+		read = cgroupstats.Read
+	}
+	observedAt := time.Now()
+	sample, err := read(filepath.Join(s.cgroupRoot, sandboxCgroupContainer))
+	if err != nil {
+		return nil, err
+	}
+
+	return &ateompb.WorkloadStatsSample{
 		Atespace:               active.Ref.Atespace,
 		ActorName:              active.Ref.Name,
 		ActorUid:               active.UID,

cmd/ateom-gvisor/stats_test.go

@@ -28,6 +28,7 @@ import (
 	"google.golang.org/grpc/status"
 	"google.golang.org/protobuf/testing/protocmp"
 
+	"github.com/agent-substrate/substrate/cmd/ateom-gvisor/internal/cgroupstats"
 	"github.com/agent-substrate/substrate/internal/ateomstats"
 	"github.com/agent-substrate/substrate/internal/proto/ateompb"
 	"github.com/agent-substrate/substrate/internal/resources"
@@ -89,13 +90,13 @@ func TestGetWorkloadStats(t *testing.T) {
 		t.Fatalf("GetWorkloadStats() error = %v, want nil", err)
 	}
 
-	if got.GetObservedAtUnixNano() < before || got.GetObservedAtUnixNano() > after {
-		t.Errorf("GetWorkloadStats() observed_at_unix_nano = %d, want within [%d, %d]", got.GetObservedAtUnixNano(), before, after)
+	if got.GetSample().GetObservedAtUnixNano() < before || got.GetSample().GetObservedAtUnixNano() > after {
+		t.Errorf("GetWorkloadStats() observed_at_unix_nano = %d, want within [%d, %d]", got.GetSample().GetObservedAtUnixNano(), before, after)
 	}
 	// Checked above; zeroed so the rest can be compared as a whole.
-	got.ObservedAtUnixNano = 0
+	got.GetSample().ObservedAtUnixNano = 0
 
-	want := &ateompb.GetWorkloadStatsResponse{
+	want := &ateompb.GetWorkloadStatsResponse{Sample: &ateompb.WorkloadStatsSample{
 		Atespace:               "space-a",
 		ActorName:              "actor-a",
 		ActorUid:               "uid-a",
@@ -107,7 +108,7 @@ func TestGetWorkloadStats(t *testing.T) {
 		MemoryPeakBytes:        209715200,
 		MemoryWorkingSetBytes:  136314880,
 		CpuUsageUsec:           1234567,
-	}
+	}}
 	if diff := cmp.Diff(want, got, protocmp.Transform()); diff != "" {
 		t.Errorf("GetWorkloadStats() mismatch (-want +got):\n%s", diff)
 	}
@@ -218,3 +219,122 @@ func TestAteomServiceStartsAvailable(t *testing.T) {
 		t.Errorf("new AteomService.activeActor = %v, want nil", got)
 	}
 }
+
+func TestGetActiveWorkloadStats(t *testing.T) {
+	s := newStatsService(t, healthyCgroup)
+	s.activeActor.Store(&testActor)
+
+	got, err := s.GetActiveWorkloadStats(context.Background(), &ateompb.GetActiveWorkloadStatsRequest{})
+	if err != nil {
+		t.Fatalf("GetActiveWorkloadStats() error = %v, want nil", err)
+	}
+	if got.GetSample() == nil {
+		t.Fatalf("GetActiveWorkloadStats() = %v, want a sample", got)
+	}
+
+	// The keyed read against the same fixture is the reference: the discovery
+	// read must produce the identical sample, since both are the same
+	// measurement with a different addressing mode.
+	want, err := s.GetWorkloadStats(context.Background(), &ateompb.GetWorkloadStatsRequest{ActorUid: "uid-a"})
+	if err != nil {
+		t.Fatalf("GetWorkloadStats() error = %v, want nil", err)
+	}
+	sample := got.GetSample()
+	sample.ObservedAtUnixNano = 0
+	want.GetSample().ObservedAtUnixNano = 0
+	if diff := cmp.Diff(want.GetSample(), sample, protocmp.Transform()); diff != "" {
+		t.Errorf("discovery sample differs from keyed sample (-keyed +discovery):\n%s", diff)
+	}
+}
+
+// TestGetActiveWorkloadStatsAvailable pins the contract that makes the
+// discovery read scrapeable: an idle ateom is a reason, never an error.
+func TestGetActiveWorkloadStatsAvailable(t *testing.T) {
+	s := newStatsService(t, healthyCgroup)
+
+	got, err := s.GetActiveWorkloadStats(context.Background(), &ateompb.GetActiveWorkloadStatsRequest{})
+	if err != nil {
+		t.Fatalf("GetActiveWorkloadStats() on an available ateom: error = %v, want nil", err)
+	}
+	if got.GetNoSampleReason() != ateompb.NoSampleReason_NO_SAMPLE_REASON_NO_WORKLOAD {
+		t.Errorf("GetActiveWorkloadStats() = %v, want NO_WORKLOAD reason", got)
+	}
+}
+
+// TestGetActiveWorkloadStatsBooting: executing but nothing to measure yet is
+// a NOT_MEASURABLE_YET reason, not an error, unlike the keyed read's
+// FAILED_PRECONDITION. A blind caller finds boots as routinely as idle
+// workers.
+func TestGetActiveWorkloadStatsBooting(t *testing.T) {
+	s := newStatsService(t, nil) // no cgroup directory: a poll landing mid-boot
+	s.activeActor.Store(&testActor)
+
+	got, err := s.GetActiveWorkloadStats(context.Background(), &ateompb.GetActiveWorkloadStatsRequest{})
+	if err != nil {
+		t.Fatalf("GetActiveWorkloadStats() mid-boot: error = %v, want nil", err)
+	}
+	if got.GetNoSampleReason() != ateompb.NoSampleReason_NO_SAMPLE_REASON_NOT_MEASURABLE_YET {
+		t.Errorf("GetActiveWorkloadStats() mid-boot = %v, want NOT_MEASURABLE_YET reason", got)
+	}
+}
+
+// The transition tests cover the re-check that runs after the lock-free
+// measurement, via the readSandboxCgroup seam: flipping activeActor inside the
+// read lands in exactly the window a checkpoint plus a fresh run (or a
+// checkpoint alone) can land in.
+
+func TestGetActiveWorkloadStatsTransition(t *testing.T) {
+	otherActor := testActor
+	otherActor.UID = "uid-b"
+
+	tests := []struct {
+		name string
+		to   *ateomstats.ActorAttribution
+		want ateompb.NoSampleReason
+	}{
+		// A new actor took the slot: there is a workload, its numbers are just
+		// not attributable this tick.
+		{name: "to another actor", to: &otherActor, want: ateompb.NoSampleReason_NO_SAMPLE_REASON_NOT_MEASURABLE_YET},
+		// A checkpoint emptied the slot: report what is true now.
+		{name: "to available", to: nil, want: ateompb.NoSampleReason_NO_SAMPLE_REASON_NO_WORKLOAD},
+	}
+
+	for _, tc := range tests {
+		t.Run(tc.name, func(t *testing.T) {
+			s := newStatsService(t, healthyCgroup)
+			s.activeActor.Store(&testActor)
+			s.readSandboxCgroup = func(dir string) (cgroupstats.Sample, error) {
+				s.activeActor.Store(tc.to)
+				return cgroupstats.Read(dir)
+			}
+
+			got, err := s.GetActiveWorkloadStats(context.Background(), &ateompb.GetActiveWorkloadStatsRequest{})
+			if err != nil {
+				t.Fatalf("GetActiveWorkloadStats() during transition: error = %v, want nil", err)
+			}
+			if got.GetSample() != nil {
+				t.Errorf("GetActiveWorkloadStats() during transition returned sample %v, want none", got.GetSample())
+			}
+			if got.GetNoSampleReason() != tc.want {
+				t.Errorf("GetActiveWorkloadStats() during transition = %v, want %v reason", got, tc.want)
+			}
+		})
+	}
+}
+
+// TestGetWorkloadStatsTransition pins the keyed read's side of the same
+// window: the caller asserted an actor that is gone by the time the sample
+// exists, so the answer is NOT_FOUND -- its mapping wants re-resolving.
+func TestGetWorkloadStatsTransition(t *testing.T) {
+	s := newStatsService(t, healthyCgroup)
+	s.activeActor.Store(&testActor)
+	s.readSandboxCgroup = func(dir string) (cgroupstats.Sample, error) {
+		s.activeActor.Store(nil)
+		return cgroupstats.Read(dir)
+	}
+
+	_, err := s.GetWorkloadStats(context.Background(), &ateompb.GetWorkloadStatsRequest{ActorUid: "uid-a"})
+	if got := status.Code(err); got != codes.NotFound {
+		t.Errorf("GetWorkloadStats() during transition: code = %v, want %v (err: %v)", got, codes.NotFound, err)
+	}
+}