diff --git a/pkg/sentry/pgalloc/BUILD b/pkg/sentry/pgalloc/BUILD index 32f5a769ceb..14fc46e42cb 100644 --- a/pkg/sentry/pgalloc/BUILD +++ b/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"], +) diff --git a/pkg/sentry/pgalloc/pgalloc.proto b/pkg/sentry/pgalloc/pgalloc.proto new file mode 100644 index 00000000000..43ae4726eab --- /dev/null +++ b/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 subreleased = 8; +} diff --git a/pkg/sentry/pgalloc/save_restore.go b/pkg/sentry/pgalloc/save_restore.go index 163ee5ff0d6..a2b01dcb5df 100644 --- a/pkg/sentry/pgalloc/save_restore.go +++ b/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))