Convert memoryFileSaved to MemoryFileMetadataProto.

pkg/sentry/pgalloc/BUILD

@@ -1,5 +1,5 @@
 load("//pkg/sync/locking:locking.bzl", "declare_mutex")
-load("//tools:defs.bzl", "go_library", "go_test")
+load("//tools:defs.bzl", "go_library", "go_test", "proto_library")
 load("//tools/go_generics:defs.bzl", "go_template_instance")
 
 package(
@@ -173,6 +173,7 @@ go_library(
     ],
     visibility = ["//pkg/sentry:internal"],
     deps = [
+        ":pgalloc_go_proto",
         "//pkg/abi/linux",
         "//pkg/atomicbitops",
         "//pkg/bitmap",
@@ -197,6 +198,7 @@ go_library(
         "//pkg/syncevent",
         "//pkg/timing",
         "//pkg/usermem",
+        "@org_golang_google_protobuf//proto:go_default_library",
         "@org_golang_x_sys//unix:go_default_library",
     ],
 )
@@ -214,3 +216,8 @@ go_test(
         "//pkg/sentry/memmap",
     ],
 )
+
+proto_library(
+    name = "pgalloc",
+    srcs = ["pgalloc.proto"],
+)

pkg/sentry/pgalloc/pgalloc.proto

@@ -0,0 +1,51 @@
+// Copyright 2026 The gVisor Authors.
+//
+// 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.
+
+syntax = "proto3";
+
+package gvisor.pgalloc;
+
+message FileRangeProto {
+  uint64 start = 1;
+  uint64 end = 2;
+}
+
+message UnfreeRangeProto {
+  uint64 start = 1;
+  uint64 end = 2;
+  uint64 refs = 3;
+}
+
+message MemAcctRangeProto {
+  uint64 start = 1;
+  uint64 end = 2;
+  uint32 kind = 3;
+  uint32 mem_cg_id = 4;
+  bool known_committed = 5;
+}
+
+message ChunkInfoProto {
+  bool huge = 1;
+}
+
+message MemoryFileMetadataProto {
+  uint64 version = 1;
+  repeated ChunkInfoProto chunks = 2;
+  repeated MemAcctRangeProto mem_acct = 3;
+  repeated UnfreeRangeProto unfree_small = 4;
+  repeated UnfreeRangeProto unfree_huge = 5;
+  repeated FileRangeProto unwaste_small = 6;
+  repeated FileRangeProto unwaste_huge = 7;
+  map<uint64, uint64> subreleased = 8;
+}

pkg/sentry/pgalloc/save_restore.go

@@ -17,6 +17,7 @@ package pgalloc
 import (
 	"bytes"
 	"context"
+	"encoding/binary"
 	"fmt"
 	"io"
 	"math"
@@ -25,6 +26,7 @@ import (
 	"time"
 
 	"golang.org/x/sys/unix"
+	"google.golang.org/protobuf/proto"
 	"gvisor.dev/gvisor/pkg/atomicbitops"
 	"gvisor.dev/gvisor/pkg/bitmap"
 	"gvisor.dev/gvisor/pkg/errors/linuxerr"
@@ -35,6 +37,7 @@ import (
 	"gvisor.dev/gvisor/pkg/ringdeque"
 	"gvisor.dev/gvisor/pkg/sentry/checkpoint"
 	"gvisor.dev/gvisor/pkg/sentry/memmap"
+	pgallocpb "gvisor.dev/gvisor/pkg/sentry/pgalloc/pgalloc_go_proto"
 	"gvisor.dev/gvisor/pkg/sentry/state/stateio"
 	"gvisor.dev/gvisor/pkg/sentry/usage"
 	"gvisor.dev/gvisor/pkg/state"
@@ -63,17 +66,103 @@ func (f *MemoryFile) ResourceID() checkpoint.ResourceID {
 	return f.opts.ResourceID
 }
 
-// memoryFileSaved is the subset of MemoryFile that is stored in checkpoints.
-//
-// +stateify savable
-type memoryFileSaved struct {
-	unwasteSmall *unwasteSet
-	unwasteHuge  *unwasteSet
-	unfreeSmall  *unfreeSet
-	unfreeHuge   *unfreeSet
-	subreleased  map[uint64]uint64
-	memAcct      *memAcctSet
-	chunks       []chunkInfo
+func (f *MemoryFile) exportMetadataProto() *pgallocpb.MemoryFileMetadataProto {
+	pb := &pgallocpb.MemoryFileMetadataProto{
+		Version:     1,
+		Subreleased: make(map[uint64]uint64, len(f.subreleased)),
+	}
+	for k, v := range f.subreleased {
+		pb.Subreleased[k] = v
+	}
+	chunks := f.chunksLoad()
+	pb.Chunks = make([]*pgallocpb.ChunkInfoProto, len(chunks))
+	for i, c := range chunks {
+		pb.Chunks[i] = &pgallocpb.ChunkInfoProto{
+			Huge: c.huge,
+		}
+	}
+	for s := f.memAcct.FirstSegment(); s.Ok(); s = s.NextSegment() {
+		val := s.Value()
+		pb.MemAcct = append(pb.MemAcct, &pgallocpb.MemAcctRangeProto{
+			Start:          s.Start(),
+			End:            s.End(),
+			Kind:           uint32(val.kind),
+			MemCgId:        val.memCgID,
+			KnownCommitted: val.knownCommitted,
+		})
+	}
+	for s := f.unfreeSmall.FirstSegment(); s.Ok(); s = s.NextSegment() {
+		val := s.Value()
+		pb.UnfreeSmall = append(pb.UnfreeSmall, &pgallocpb.UnfreeRangeProto{
+			Start: s.Start(),
+			End:   s.End(),
+			Refs:  val.refs,
+		})
+	}
+	for s := f.unfreeHuge.FirstSegment(); s.Ok(); s = s.NextSegment() {
+		val := s.Value()
+		pb.UnfreeHuge = append(pb.UnfreeHuge, &pgallocpb.UnfreeRangeProto{
+			Start: s.Start(),
+			End:   s.End(),
+			Refs:  val.refs,
+		})
+	}
+	for s := f.unwasteSmall.FirstSegment(); s.Ok(); s = s.NextSegment() {
+		pb.UnwasteSmall = append(pb.UnwasteSmall, &pgallocpb.FileRangeProto{
+			Start: s.Start(),
+			End:   s.End(),
+		})
+	}
+	for s := f.unwasteHuge.FirstSegment(); s.Ok(); s = s.NextSegment() {
+		pb.UnwasteHuge = append(pb.UnwasteHuge, &pgallocpb.FileRangeProto{
+			Start: s.Start(),
+			End:   s.End(),
+		})
+	}
+	return pb
+}
+
+func (f *MemoryFile) importMetadataProto(pb *pgallocpb.MemoryFileMetadataProto) error {
+	if pb.Version != 1 {
+		return fmt.Errorf("unsupported MemoryFileMetadataProto version %d", pb.Version)
+	}
+	f.subreleased = make(map[uint64]uint64, len(pb.Subreleased))
+	for k, v := range pb.Subreleased {
+		f.subreleased[k] = v
+	}
+	chunks := make([]chunkInfo, len(pb.Chunks))
+	for i, c := range pb.Chunks {
+		chunks[i] = chunkInfo{
+			huge: c.Huge,
+		}
+	}
+	f.chunks.Store(&chunks)
+
+	f.memAcct.RemoveAll()
+	for _, s := range pb.MemAcct {
+		f.memAcct.InsertRange(memmap.FileRange{Start: s.Start, End: s.End}, memAcctInfo{
+			kind:           usage.MemoryKind(s.Kind),
+			memCgID:        s.MemCgId,
+			knownCommitted: s.KnownCommitted,
+		})
+	}
+	f.unfreeSmall.RemoveAll()
+	for _, s := range pb.UnfreeSmall {
+		f.unfreeSmall.InsertRange(memmap.FileRange{Start: s.Start, End: s.End}, unfreeInfo{refs: s.Refs})
+	}
+	f.unfreeHuge.RemoveAll()
+	for _, s := range pb.UnfreeHuge {
+		f.unfreeHuge.InsertRange(memmap.FileRange{Start: s.Start, End: s.End}, unfreeInfo{refs: s.Refs})
+	}
+	f.unwasteSmall.RemoveAll()
+	for _, s := range pb.UnwasteSmall {
+		f.unwasteSmall.InsertRange(memmap.FileRange{Start: s.Start, End: s.End}, unwasteInfo{})
+	}
+	f.unwasteHuge.RemoveAll()
+	for _, s := range pb.UnwasteHuge {
+		f.unwasteHuge.InsertRange(memmap.FileRange{Start: s.Start, End: s.End}, unwasteInfo{})
+	}
+	return nil
 }
 
 // SaveOpts provides options to MemoryFile.SaveTo().
@@ -386,16 +475,18 @@ func (f *MemoryFile) SaveTo(ctx context.Context, w io.Writer, opts *SaveOpts) er
 
 	// Save metadata.
 	timeMetadataStart := gohacks.Nanotime()
-	if _, err := state.Save(ctx, w, &memoryFileSaved{
-		unwasteSmall: &f.unwasteSmall,
-		unwasteHuge:  &f.unwasteHuge,
-		unfreeSmall:  &f.unfreeSmall,
-		unfreeHuge:   &f.unfreeHuge,
-		subreleased:  f.subreleased,
-		memAcct:      &f.memAcct,
-		chunks:       f.chunksLoad(),
-	}); err != nil {
-		return fmt.Errorf("failed to save metadata: %w", err)
+	pb := f.exportMetadataProto()
+	data, err := proto.Marshal(pb)
+	if err != nil {
+		return fmt.Errorf("failed to marshal metadata: %w", err)
+	}
+	var lengthBuf [8]byte
+	binary.LittleEndian.PutUint64(lengthBuf[:], uint64(len(data)))
+	if _, err := w.Write(lengthBuf[:]); err != nil {
+		return fmt.Errorf("failed to write metadata length: %w", err)
+	}
+	if _, err := w.Write(data); err != nil {
+		return fmt.Errorf("failed to write metadata: %w", err)
 	}
 	log.Infof("MemoryFile(%p): saved metadata in %s", f, time.Duration(gohacks.Nanotime()-timeMetadataStart))
 
@@ -935,18 +1026,23 @@ func (f *MemoryFile) LoadFrom(ctx context.Context, r io.Reader, opts *LoadOpts)
 
 	// Load metadata.
 	timeMetadataStart := gohacks.Nanotime()
-	var mfs memoryFileSaved
-	if _, err := state.Load(ctx, r, &mfs); err != nil {
-		return fmt.Errorf("failed to load metadata: %w", err)
-	}
-	f.unwasteSmall.MoveFrom(mfs.unwasteSmall)
-	f.unwasteHuge.MoveFrom(mfs.unwasteHuge)
-	f.unfreeSmall.MoveFrom(mfs.unfreeSmall)
-	f.unfreeHuge.MoveFrom(mfs.unfreeHuge)
-	f.subreleased = mfs.subreleased
-	f.memAcct.MoveFrom(mfs.memAcct)
-	chunks := mfs.chunks
-	f.chunks.Store(&chunks)
+	var lengthBuf [8]byte
+	if _, err := io.ReadFull(r, lengthBuf[:]); err != nil {
+		return fmt.Errorf("failed to read metadata length: %w", err)
+	}
+	length := binary.LittleEndian.Uint64(lengthBuf[:])
+	data := make([]byte, length)
+	if _, err := io.ReadFull(r, data); err != nil {
+		return fmt.Errorf("failed to read metadata: %w", err)
+	}
+	var pb pgallocpb.MemoryFileMetadataProto
+	if err := proto.Unmarshal(data, &pb); err != nil {
+		return fmt.Errorf("failed to unmarshal metadata: %w", err)
+	}
+	if err := f.importMetadataProto(&pb); err != nil {
+		return fmt.Errorf("failed to import metadata: %w", err)
+	}
+	chunks := f.chunksLoad()
 	mfTimeline.Reached("metadata loaded")
 	log.Infof("MemoryFile(%p): loaded metadata in %s", f, time.Duration(gohacks.Nanotime()-timeMetadataStart))