ateom-microvm: measure the guest agent for GetWorkloadStats

cmd/ateom-microvm/checkpoint.go

@@ -263,6 +263,13 @@ func listFiles(dir string) ([]string, error) {
 // snapshot is already on disk, so this only needs to release resources. ra may be
 // nil (e.g. ateom restarted and lost in-memory state).
 func (s *AteomService) teardownActor(ctx context.Context, id string, ra *runningActor, client *ch.Client) {
+	// Stop offering the guest to GetWorkloadStats first, before anything below
+	// makes it stop answering. Clearing it here rather than alongside the
+	// attribution is what keeps a poll that lands mid-teardown on the
+	// FAILED_PRECONDITION path ("no numbers right now") instead of surfacing a
+	// closed connection as a failed read.
+	s.guestStats.Store(nil)
+
 	if client != nil {
 		tShutdown := time.Now()
 		shutCtx, cancel := context.WithTimeout(ctx, 10*time.Second)
@@ -278,9 +285,9 @@ func (s *AteomService) teardownActor(ctx context.Context, id string, ra *running
 		// fails the forwarding goroutines' in-flight ReadStdout/ReadStderr calls, so
 		// they return io.EOF and exit (no goroutine leak). Guarded so a second
 		// teardown / a never-forwarded actor is a no-op.
-		if ra.logAgent != nil {
-			_ = ra.logAgent.Close()
-			ra.logAgent = nil
+		if ra.guestAgent != nil {
+			_ = ra.guestAgent.Close()
+			ra.guestAgent = nil
 		}
 
 		// Kill the CH process ateom launched.

cmd/ateom-microvm/internal/agentstats/agentstats.go

@@ -0,0 +1,164 @@
+// 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 agentstats turns the kata guest agent's per-container cgroup
+// accounting into the resource-usage sample ateom reports.
+//
+// The micro-VM ateom uses it to answer ateompb.Ateom/GetWorkloadStats. The host
+// cgroup is the wrong place to look on this runtime: the guest's RAM is a fixed
+// allocation cloud-hypervisor takes at boot, so the host cgroup reads roughly
+// the same whether the actor is idle or saturated. The numbers that move with
+// the workload are the ones the guest kernel keeps, and the agent is what can
+// read them.
+//
+// This package is deliberately pure — it converts an already-fetched
+// agentpb.CgroupStats and never talks to a guest — which keeps it testable
+// without a live micro-VM and, unlike the rest of the micro-VM ateom, without
+// the linux build tag.
+package agentstats
+
+import (
+	"github.com/agent-substrate/substrate/cmd/ateom-microvm/internal/third_party/kata/agentpb"
+)
+
+// Sample is a point-in-time reading for one container, or the sum of several.
+// It carries the same four numbers as the gVisor ateom's cgroupstats.Sample,
+// because both feed the same four fields of GetWorkloadStatsResponse.
+//
+// Anything the guest did not report reads as zero rather than failing the whole
+// sample, for the same reason as there: a partial reading is more useful than
+// none. FromCgroupStats says which fields can do that and why.
+type Sample struct {
+	// MemoryCurrentBytes is what is currently charged to the container's guest
+	// cgroup, page cache included.
+	MemoryCurrentBytes uint64
+
+	// MemoryPeakBytes is the high-water mark of MemoryCurrentBytes. Zero when
+	// the guest kernel does not expose one: the agent fills it from the cgroup's
+	// max-usage file, which cgroup v2 only grew in Linux 5.19.
+	//
+	// For a single container this is exact; summed across containers it is an
+	// upper bound, not an observed maximum — see Plus for why.
+	MemoryPeakBytes uint64
+
+	// MemoryWorkingSetBytes is MemoryCurrentBytes less the reclaimable page
+	// cache, floored at zero — the figure to compare against a memory limit,
+	// since MemoryCurrentBytes drifts upward with cache the kernel would drop
+	// for free under pressure.
+	MemoryWorkingSetBytes uint64
+
+	// CPUUsageUsec is cumulative CPU time consumed by the container, as seen by
+	// the guest kernel.
+	CPUUsageUsec uint64
+}
+
+// Keys of the memory.stat entry holding reclaimable file-backed pages. The
+// agent passes the guest's memory.stat through verbatim, so which one is
+// present depends on the cgroup version the guest kernel gave the container:
+// v2 names it inactive_file, v1 has a per-cgroup inactive_file and the
+// hierarchical total_inactive_file, and the total is the one that matches what
+// v2's figure means.
+const (
+	inactiveFileV2 = "inactive_file"
+	inactiveFileV1 = "total_inactive_file"
+)
+
+// FromCgroupStats converts one container's guest cgroup accounting.
+//
+// It never fails. Every field the agent left out reads as zero, and cs itself
+// may be nil — the agent answers without cgroup stats for a container it has no
+// accounting for, which is a normal state for one that has exited rather than
+// an error. The caller decides what an all-zero container means; see the
+// summing in the micro-VM ateom's GetWorkloadStats.
+func FromCgroupStats(cs *agentpb.CgroupStats) Sample {
+	mem := cs.GetMemoryStats().GetUsage()
+	current := mem.GetUsage()
+
+	// Saturating rather than wrapping. The guest reads usage and memory.stat a
+	// moment apart, so the reclaimable figure can legitimately exceed the usage
+	// read beside it; on uint64 the naive subtraction gives an absurd number
+	// instead of the near-zero the reading means.
+	workingSet := current
+	if inactiveFile, ok := inactiveFileBytes(cs); ok {
+		workingSet = 0
+		if inactiveFile < current {
+			workingSet = current - inactiveFile
+		}
+	}
+
+	return Sample{
+		MemoryCurrentBytes: current,
+		MemoryPeakBytes:    mem.GetMaxUsage(),
+
+		MemoryWorkingSetBytes: workingSet,
+
+		// The agent reports CPU time in nanoseconds, matching the runc stats
+		// struct its own is modeled on; the proto wants microseconds. Truncating
+		// division loses at most a microsecond per sample, and the field is
+		// cumulative, so the error does not accumulate across samples.
+		CPUUsageUsec: cs.GetCpuStats().GetCpuUsage().GetTotalUsage() / 1000,
+	}
+}
+
+// inactiveFileBytes returns the reclaimable page cache the guest reported, and
+// whether it reported one at all. Absent, the working set collapses to usage,
+// which over-reports by however much reclaimable cache the container holds —
+// the safe direction, since it never claims the workload is using less than it
+// is.
+func inactiveFileBytes(cs *agentpb.CgroupStats) (uint64, bool) {
+	stats := cs.GetMemoryStats().GetStats()
+	for _, key := range []string{inactiveFileV2, inactiveFileV1} {
+		if v, ok := stats[key]; ok {
+			return v, true
+		}
+	}
+	return 0, false
+}
+
+// Plus returns the sum of two samples, for accumulating an actor's containers
+// into the one figure the proto reports.
+//
+// Summing the peaks is an upper bound on the peak of the sum, not the peak of
+// the sum itself: two containers that peaked at different moments add up to a
+// total the actor never actually reached. The bound is the honest
+// approximation, and for the single-container actors this runtime mostly serves
+// it is exact.
+//
+// The guest kernel does track the true figure: every container cgroup sits
+// under the shared /ateomchv parent (see StartOverlayWorkload and the
+// CgroupsPath defaults in internal/kata), so with hierarchical accounting the
+// parent's memory.peak is the actor-level maximum this sum approximates. The
+// kata-agent just has no RPC that reads a cgroup by path — StatsContainer is
+// per-container only. If the agent ever grows a sandbox-level stats read, that
+// replaces this summing outright.
+//
+// Saturating on overflow, so a nonsensical reading from one container cannot
+// wrap the total to a small number and read as healthy.
+func (s Sample) Plus(o Sample) Sample {
+	return Sample{
+		MemoryCurrentBytes:    addSaturating(s.MemoryCurrentBytes, o.MemoryCurrentBytes),
+		MemoryPeakBytes:       addSaturating(s.MemoryPeakBytes, o.MemoryPeakBytes),
+		MemoryWorkingSetBytes: addSaturating(s.MemoryWorkingSetBytes, o.MemoryWorkingSetBytes),
+		CPUUsageUsec:          addSaturating(s.CPUUsageUsec, o.CPUUsageUsec),
+	}
+}
+
+// addSaturating returns a+b, or the maximum uint64 if that would wrap.
+func addSaturating(a, b uint64) uint64 {
+	if sum := a + b; sum >= a {
+		return sum
+	}
+	const maxUint64 = ^uint64(0)
+	return maxUint64
+}

cmd/ateom-microvm/internal/agentstats/agentstats_test.go

@@ -0,0 +1,190 @@
+// 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 agentstats
+
+import (
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+
+	"github.com/agent-substrate/substrate/cmd/ateom-microvm/internal/third_party/kata/agentpb"
+)
+
+// cgroupStats builds an agent reading: usage/max bytes, the memory.stat entries
+// the guest reported, and cumulative CPU nanoseconds.
+func cgroupStats(usage, maxUsage uint64, stats map[string]uint64, cpuNanos uint64) *agentpb.CgroupStats {
+	return &agentpb.CgroupStats{
+		MemoryStats: &agentpb.MemoryStats{
+			Usage: &agentpb.MemoryData{Usage: usage, MaxUsage: maxUsage},
+			Stats: stats,
+		},
+		CpuStats: &agentpb.CpuStats{
+			CpuUsage: &agentpb.CpuUsage{TotalUsage: cpuNanos},
+		},
+	}
+}
+
+func TestFromCgroupStats(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		cs   *agentpb.CgroupStats
+		want Sample
+	}{
+		{
+			name: "cgroup v2 guest",
+			cs:   cgroupStats(157286400, 209715200, map[string]uint64{"inactive_file": 20971520}, 1234567000),
+			want: Sample{
+				MemoryCurrentBytes:    157286400,
+				MemoryPeakBytes:       209715200,
+				MemoryWorkingSetBytes: 136314880,
+				CPUUsageUsec:          1234567,
+			},
+		},
+		{
+			// A v1 guest names the hierarchical figure differently. Dropping to
+			// the working-set-equals-usage fallback here would over-report every
+			// sample from such a guest rather than fail visibly, so it is worth
+			// pinning that the alternate key is understood.
+			name: "cgroup v1 guest reports total_inactive_file",
+			cs:   cgroupStats(1000, 2000, map[string]uint64{"total_inactive_file": 400}, 0),
+			want: Sample{MemoryCurrentBytes: 1000, MemoryPeakBytes: 2000, MemoryWorkingSetBytes: 600},
+		},
+		{
+			// v1 reports both: a per-cgroup figure and the hierarchical total.
+			// The total is the one that means what v2's inactive_file means.
+			name: "both keys present prefers the v2 name",
+			cs:   cgroupStats(1000, 0, map[string]uint64{"inactive_file": 100, "total_inactive_file": 400}, 0),
+			want: Sample{MemoryCurrentBytes: 1000, MemoryWorkingSetBytes: 900},
+		},
+		{
+			// The two figures are not a consistent snapshot, so this is
+			// reachable; on uint64 the naive subtraction would report ~1.8e19.
+			name: "reclaimable cache above usage floors at zero",
+			cs:   cgroupStats(1000, 0, map[string]uint64{"inactive_file": 4000}, 0),
+			want: Sample{MemoryCurrentBytes: 1000, MemoryWorkingSetBytes: 0},
+		},
+		{
+			name: "equal usage and reclaimable cache floors at zero",
+			cs:   cgroupStats(1000, 0, map[string]uint64{"inactive_file": 1000}, 0),
+			want: Sample{MemoryCurrentBytes: 1000, MemoryWorkingSetBytes: 0},
+		},
+		{
+			// Nothing to subtract, so the working set collapses to usage: an
+			// over-report, which is the safe direction.
+			name: "no memory.stat entries",
+			cs:   cgroupStats(1000, 2000, nil, 0),
+			want: Sample{MemoryCurrentBytes: 1000, MemoryPeakBytes: 2000, MemoryWorkingSetBytes: 1000},
+		},
+		{
+			name: "memory.stat without a reclaimable-cache entry",
+			cs:   cgroupStats(1000, 0, map[string]uint64{"anon": 900}, 0),
+			want: Sample{MemoryCurrentBytes: 1000, MemoryWorkingSetBytes: 1000},
+		},
+		{
+			// Guests below Linux 5.19 have no cgroup v2 high-water mark. The rest
+			// of the sample must survive that.
+			name: "no peak reported",
+			cs:   cgroupStats(1000, 0, map[string]uint64{"inactive_file": 100}, 5000),
+			want: Sample{MemoryCurrentBytes: 1000, MemoryWorkingSetBytes: 900, CPUUsageUsec: 5},
+		},
+		{
+			// The agent reports nanoseconds and the proto wants microseconds.
+			// A sub-microsecond total truncates to zero rather than rounding up.
+			name: "sub-microsecond cpu time truncates",
+			cs:   cgroupStats(0, 0, nil, 999),
+			want: Sample{},
+		},
+		{
+			name: "cpu time rounds down to whole microseconds",
+			cs:   cgroupStats(0, 0, nil, 1999),
+			want: Sample{CPUUsageUsec: 1},
+		},
+		{
+			name: "memory reported without cpu",
+			cs:   &agentpb.CgroupStats{MemoryStats: &agentpb.MemoryStats{Usage: &agentpb.MemoryData{Usage: 4096}}},
+			want: Sample{MemoryCurrentBytes: 4096, MemoryWorkingSetBytes: 4096},
+		},
+		{
+			name: "cpu reported without memory",
+			cs:   &agentpb.CgroupStats{CpuStats: &agentpb.CpuStats{CpuUsage: &agentpb.CpuUsage{TotalUsage: 2000}}},
+			want: Sample{CPUUsageUsec: 2},
+		},
+		{
+			// What the agent answers for a container it has no accounting for —
+			// one that has exited, most often. Must be a zero sample, not a
+			// panic: this is parsed on a timer for the life of every workload.
+			name: "nil cgroup stats",
+			cs:   nil,
+			want: Sample{},
+		},
+		{
+			name: "empty cgroup stats",
+			cs:   &agentpb.CgroupStats{},
+			want: Sample{},
+		},
+		{
+			// The reclaimable figure is there but the usage it subtracts from is
+			// not, so the working set floors rather than wrapping.
+			name: "reclaimable cache without a usage message",
+			cs:   &agentpb.CgroupStats{MemoryStats: &agentpb.MemoryStats{Stats: map[string]uint64{"inactive_file": 100}}},
+			want: Sample{},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			if diff := cmp.Diff(tc.want, FromCgroupStats(tc.cs)); diff != "" {
+				t.Errorf("FromCgroupStats() mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}
+
+func TestSamplePlus(t *testing.T) {
+	maxUint64 := ^uint64(0)
+
+	for _, tc := range []struct {
+		name string
+		a, b Sample
+		want Sample
+	}{
+		{
+			name: "adds every field",
+			a:    Sample{MemoryCurrentBytes: 100, MemoryPeakBytes: 200, MemoryWorkingSetBytes: 90, CPUUsageUsec: 10},
+			b:    Sample{MemoryCurrentBytes: 1, MemoryPeakBytes: 2, MemoryWorkingSetBytes: 3, CPUUsageUsec: 4},
+			want: Sample{MemoryCurrentBytes: 101, MemoryPeakBytes: 202, MemoryWorkingSetBytes: 93, CPUUsageUsec: 14},
+		},
+		{
+			// The accumulator starts here, so a zero left operand must be the
+			// identity or every actor's first container would be dropped.
+			name: "zero is the identity",
+			a:    Sample{},
+			b:    Sample{MemoryCurrentBytes: 7, MemoryPeakBytes: 8, MemoryWorkingSetBytes: 9, CPUUsageUsec: 10},
+			want: Sample{MemoryCurrentBytes: 7, MemoryPeakBytes: 8, MemoryWorkingSetBytes: 9, CPUUsageUsec: 10},
+		},
+		{
+			// A wrapped total would read as a nearly idle actor, which is the one
+			// wrong answer that looks plausible.
+			name: "saturates instead of wrapping",
+			a:    Sample{MemoryCurrentBytes: maxUint64, MemoryPeakBytes: maxUint64, MemoryWorkingSetBytes: maxUint64, CPUUsageUsec: maxUint64},
+			b:    Sample{MemoryCurrentBytes: 1, MemoryPeakBytes: 2, MemoryWorkingSetBytes: 3, CPUUsageUsec: 4},
+			want: Sample{MemoryCurrentBytes: maxUint64, MemoryPeakBytes: maxUint64, MemoryWorkingSetBytes: maxUint64, CPUUsageUsec: maxUint64},
+		},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			if diff := cmp.Diff(tc.want, tc.a.Plus(tc.b)); diff != "" {
+				t.Errorf("Sample.Plus() mismatch (-want +got):\n%s", diff)
+			}
+		})
+	}
+}