From c30491ffe1c76d925ba8cc3172e66c43a2c69f91 Mon Sep 17 00:00:00 2001 From: Liquidzk <492928615@qq.com> Date: Tue, 21 Jul 2026 14:18:34 +0800 Subject: [PATCH 1/6] Store UPF working sets with snapshots Signed-off-by: Liquidzk <492928615@qq.com> --- ctriface/iface.go | 17 +++++++++-------- memory/manager/snapshot_state.go | 15 +++++++-------- snapshotting/manager_test.go | 12 ++++++++++++ snapshotting/snapshot.go | 8 ++++++++ 4 files changed, 36 insertions(+), 16 deletions(-) diff --git a/ctriface/iface.go b/ctriface/iface.go index a3ff9db4e..e2f386088 100644 --- a/ctriface/iface.go +++ b/ctriface/iface.go @@ -595,14 +595,15 @@ func (o *Orchestrator) LoadSnapshot(ctx context.Context, vmID string, snap *snap configureSnapshotMemoryBackend(conf, "Uffd", uffdSock) if err := o.memoryManager.PrepareSnapshotLoad(manager.SnapshotStateCfg{ - VMID: vmID, - VMMStatePath: snap.GetSnapshotFilePath(), - GuestMemPath: snap.GetMemFilePath(), - InstanceSockAddr: uffdSock, - BaseDir: o.getVMBaseDir(vmID), - GuestMemSize: int(conf.MachineCfg.MemSizeMib) * 1024 * 1024, - IsLazyMode: o.isLazyMode, - WorkingSetPath: o.getWorkingSetFile(vmID), + VMID: vmID, + VMMStatePath: snap.GetSnapshotFilePath(), + GuestMemPath: snap.GetMemFilePath(), + InstanceSockAddr: uffdSock, + BaseDir: o.getVMBaseDir(vmID), + GuestMemSize: int(conf.MachineCfg.MemSizeMib) * 1024 * 1024, + IsLazyMode: o.isLazyMode, + WorkingSetPath: snap.GetWorkingSetFilePath(), + WorkingSetTracePath: snap.GetWorkingSetTraceFilePath(), }); err != nil { return nil, nil, err } diff --git a/memory/manager/snapshot_state.go b/memory/manager/snapshot_state.go index f7f54eea3..cd0df11af 100644 --- a/memory/manager/snapshot_state.go +++ b/memory/manager/snapshot_state.go @@ -39,7 +39,7 @@ import ( type SnapshotStateCfg struct { VMID string - VMMStatePath, GuestMemPath, WorkingSetPath string + VMMStatePath, GuestMemPath, WorkingSetPath, WorkingSetTracePath string InstanceSockAddr string BaseDir string // base directory for the instance @@ -90,7 +90,7 @@ func NewSnapshotState(cfg SnapshotStateCfg) *SnapshotState { cfg = normalizeSnapshotStateCfg(cfg) s.SnapshotStateCfg = cfg - s.trace = initTrace(s.getTraceFile()) + s.trace = initTrace(s.WorkingSetTracePath) if s.metricsModeOn { s.totalPFServed = make([]float64, 0) s.uniquePFServed = make([]float64, 0) @@ -105,6 +105,9 @@ func normalizeSnapshotStateCfg(cfg SnapshotStateCfg) SnapshotStateCfg { if cfg.WorkingSetPath == "" && cfg.BaseDir != "" { cfg.WorkingSetPath = filepath.Join(cfg.BaseDir, "working_set_pages") } + if cfg.WorkingSetTracePath == "" && cfg.BaseDir != "" { + cfg.WorkingSetTracePath = filepath.Join(cfg.BaseDir, "trace") + } return cfg } @@ -113,9 +116,9 @@ func (s *SnapshotState) refreshSnapshotLoad(cfg SnapshotStateCfg) { trace := s.trace if trace == nil { - trace = initTrace(filepath.Join(cfg.BaseDir, "trace")) + trace = initTrace(cfg.WorkingSetTracePath) } else { - trace.traceFileName = filepath.Join(cfg.BaseDir, "trace") + trace.traceFileName = cfg.WorkingSetTracePath } isRecordReady := s.isRecordReady isEverActivated := s.isEverActivated @@ -196,10 +199,6 @@ func (s *SnapshotState) processMetrics() { s.currentMetric = nil } -func (s *SnapshotState) getTraceFile() string { - return filepath.Join(s.BaseDir, "trace") -} - func (s *SnapshotState) mapGuestMemory() error { fd, err := os.OpenFile(s.GuestMemPath, os.O_RDONLY, 0444) if err != nil { diff --git a/snapshotting/manager_test.go b/snapshotting/manager_test.go index 81f43e221..26e163e2e 100644 --- a/snapshotting/manager_test.go +++ b/snapshotting/manager_test.go @@ -24,6 +24,8 @@ package snapshotting_test import ( "fmt" + "path/filepath" + ctrdlog "github.com/containerd/log" log "github.com/sirupsen/logrus" "github.com/stretchr/testify/require" @@ -79,6 +81,16 @@ func TestSnapshotManagerSingle(t *testing.T) { testSnapshotManager(t, mgr, revision, imageName) } +func TestSnapshotWorkingSetPaths(t *testing.T) { + baseDir := t.TempDir() + revision := "myrevision-working-set" + snap := snapshotting.NewSnapshot(revision, baseDir, "testImage") + snapshotDir := filepath.Join(baseDir, revision) + + require.Equal(t, filepath.Join(snapshotDir, "working_set_pages"), snap.GetWorkingSetFilePath()) + require.Equal(t, filepath.Join(snapshotDir, "working_set_trace"), snap.GetWorkingSetTraceFilePath()) +} + func TestSnapshotManagerConcurrent(t *testing.T) { // Create snapshot manager mgr := snapshotting.NewSnapshotManager(snapshotsDir) diff --git a/snapshotting/snapshot.go b/snapshotting/snapshot.go index e5382a71d..bc447d768 100644 --- a/snapshotting/snapshot.go +++ b/snapshotting/snapshot.go @@ -82,6 +82,14 @@ func (snp *Snapshot) GetMemFilePath() string { return filepath.Join(snp.snapDir, "mem_file") } +func (snp *Snapshot) GetWorkingSetFilePath() string { + return filepath.Join(snp.snapDir, "working_set_pages") +} + +func (snp *Snapshot) GetWorkingSetTraceFilePath() string { + return filepath.Join(snp.snapDir, "working_set_trace") +} + func (snp *Snapshot) GetPatchFilePath() string { return filepath.Join(snp.snapDir, "patch_file") } From bcda0faaa61b90c9cd17cdc7a9b1f9d664de7e7f Mon Sep 17 00:00:00 2001 From: Liquidzk <492928615@qq.com> Date: Tue, 21 Jul 2026 14:26:31 +0800 Subject: [PATCH 2/6] Persist UPF working sets across restores Signed-off-by: Liquidzk <492928615@qq.com> --- memory/manager/manager.go | 4 +- memory/manager/manager_test.go | 140 +++++++++++++++++-- memory/manager/snapshot_state.go | 45 +++--- memory/manager/snapshot_state_test.go | 17 ++- memory/manager/trace.go | 192 ++++++++++++++++---------- 5 files changed, 295 insertions(+), 103 deletions(-) diff --git a/memory/manager/manager.go b/memory/manager/manager.go index e25ab5c1a..7994a6f1f 100644 --- a/memory/manager/manager.go +++ b/memory/manager/manager.go @@ -231,12 +231,12 @@ func (m *MemoryManager) FetchState(vmID string) error { if state.metricsModeOn && state.currentMetric == nil { state.currentMetric = metrics.NewMetric() } - if state.metricsModeOn && state.isRecordReady && !state.IsLazyMode { + if state.metricsModeOn && !state.IsLazyMode { tStart = time.Now() } err := state.fetchState() - if err == nil && !tStart.IsZero() { + if err == nil && state.isRecordReady && !tStart.IsZero() { state.currentMetric.MetricMap[fetchStateMetric] = metrics.ToUS(time.Since(tStart)) } return err diff --git a/memory/manager/manager_test.go b/memory/manager/manager_test.go index 04b0e871d..b1f4efcbc 100644 --- a/memory/manager/manager_test.go +++ b/memory/manager/manager_test.go @@ -102,24 +102,26 @@ func TestMemoryManagerRegisterFetchPrepareDeregister(t *testing.T) { } } -func TestPrepareSnapshotLoadPreservesWorkingSetState(t *testing.T) { +func TestPrepareSnapshotLoadResetsWorkingSetState(t *testing.T) { baseDir := t.TempDir() vmID := "vm-prepare-ws" guestMemPath := filepath.Join(baseDir, "guest_mem") vmmStatePath := filepath.Join(baseDir, "state") workingSetPath := filepath.Join(baseDir, "working_set_pages") + workingSetTracePath := filepath.Join(baseDir, "working_set_trace") prepareGuestMemoryFile(t, guestMemPath, 2*os.Getpagesize()) writeTestFile(t, vmmStatePath, "state") manager := NewMemoryManager(MemoryManagerCfg{}) cfg := SnapshotStateCfg{ - VMID: vmID, - BaseDir: baseDir, - VMMStatePath: vmmStatePath, - GuestMemPath: guestMemPath, - WorkingSetPath: workingSetPath, - GuestMemSize: 2 * os.Getpagesize(), + VMID: vmID, + BaseDir: baseDir, + VMMStatePath: vmmStatePath, + GuestMemPath: guestMemPath, + WorkingSetPath: workingSetPath, + WorkingSetTracePath: workingSetTracePath, + GuestMemSize: 2 * os.Getpagesize(), } if err := manager.RegisterVM(cfg); err != nil { t.Fatalf("RegisterVM returned error: %v", err) @@ -132,11 +134,13 @@ func TestPrepareSnapshotLoadPreservesWorkingSetState(t *testing.T) { nextVMMStatePath := filepath.Join(baseDir, "next_state") nextSocketPath := filepath.Join(baseDir, "next_uffd.sock") + nextTracePath := filepath.Join(baseDir, "next_working_set_trace") writeTestFile(t, nextVMMStatePath, "next-state") nextCfg := cfg nextCfg.VMMStatePath = nextVMMStatePath nextCfg.InstanceSockAddr = nextSocketPath + nextCfg.WorkingSetTracePath = nextTracePath nextCfg.IsLazyMode = true if err := manager.PrepareSnapshotLoad(nextCfg); err != nil { @@ -144,11 +148,14 @@ func TestPrepareSnapshotLoadPreservesWorkingSetState(t *testing.T) { } got := manager.instances[vmID] - if got.trace != trace { - t.Fatal("PrepareSnapshotLoad replaced trace state") + if got.trace == trace { + t.Fatal("PrepareSnapshotLoad retained the previous trace state") } - if !got.isRecordReady { - t.Fatal("PrepareSnapshotLoad cleared isRecordReady") + if got.isRecordReady { + t.Fatal("PrepareSnapshotLoad retained working set readiness") + } + if got.trace.traceFileName != nextTracePath { + t.Fatalf("trace path = %q, want %q", got.trace.traceFileName, nextTracePath) } if got.VMMStatePath != nextVMMStatePath { t.Fatalf("VMMStatePath = %q, want %q", got.VMMStatePath, nextVMMStatePath) @@ -161,6 +168,117 @@ func TestPrepareSnapshotLoadPreservesWorkingSetState(t *testing.T) { } } +func TestFetchStateLoadsWorkingSetAcrossVMIDs(t *testing.T) { + baseDir := t.TempDir() + snapshotDir := filepath.Join(baseDir, "revision") + if err := os.Mkdir(snapshotDir, 0755); err != nil { + t.Fatalf("os.Mkdir returned error: %v", err) + } + + guestMemPath := filepath.Join(snapshotDir, "mem_file") + vmmStatePath := filepath.Join(snapshotDir, "snap_file") + workingSetPath := filepath.Join(snapshotDir, "working_set_pages") + workingSetTracePath := filepath.Join(snapshotDir, "working_set_trace") + pageSize := uint64(os.Getpagesize()) + prepareGuestMemoryFile(t, guestMemPath, 5*int(pageSize)) + writeTestFile(t, vmmStatePath, "state") + + manager := NewMemoryManager(MemoryManagerCfg{}) + recordCfg := SnapshotStateCfg{ + VMID: "vm-record", + BaseDir: filepath.Join(baseDir, "vm-record"), + VMMStatePath: vmmStatePath, + GuestMemPath: guestMemPath, + WorkingSetPath: workingSetPath, + WorkingSetTracePath: workingSetTracePath, + GuestMemSize: 5 * int(pageSize), + } + if err := manager.RegisterVM(recordCfg); err != nil { + t.Fatalf("RegisterVM returned error: %v", err) + } + recordState := manager.instances[recordCfg.VMID] + recordState.trace.AppendRecord(Record{offset: 3 * pageSize}) + recordState.trace.AppendRecord(Record{offset: pageSize}) + if err := recordState.trace.ProcessRecord(guestMemPath, workingSetPath, pageSize); err != nil { + t.Fatalf("ProcessRecord returned error: %v", err) + } + + replayCfg := recordCfg + replayCfg.VMID = "vm-replay" + replayCfg.BaseDir = filepath.Join(baseDir, replayCfg.VMID) + if err := manager.PrepareSnapshotLoad(replayCfg); err != nil { + t.Fatalf("PrepareSnapshotLoad returned error: %v", err) + } + if err := manager.FetchState(replayCfg.VMID); err != nil { + t.Fatalf("FetchState returned error: %v", err) + } + + replayState := manager.instances[replayCfg.VMID] + if !replayState.isRecordReady { + t.Fatal("FetchState did not mark the persisted working set ready") + } + if got, want := replayState.trace.pageSize, pageSize; got != want { + t.Fatalf("trace page size = %#x, want %#x", got, want) + } + if got, want := len(replayState.trace.trace), 2; got != want { + t.Fatalf("trace length = %d, want %d", got, want) + } + if got, want := replayState.trace.trace[0].offset, pageSize; got != want { + t.Fatalf("first trace offset = %#x, want %#x", got, want) + } + if got, want := replayState.trace.trace[1].offset, 3*pageSize; got != want { + t.Fatalf("second trace offset = %#x, want %#x", got, want) + } + + guestMem, err := os.ReadFile(guestMemPath) + if err != nil { + t.Fatalf("os.ReadFile guest memory returned error: %v", err) + } + wantWorkingSet := append([]byte{}, guestMem[pageSize:2*pageSize]...) + wantWorkingSet = append(wantWorkingSet, guestMem[3*pageSize:4*pageSize]...) + if !reflect.DeepEqual(replayState.workingSet, wantWorkingSet) { + t.Fatal("loaded working set does not match recorded pages") + } +} + +func TestFetchStateRejectsInvalidWorkingSetArtifacts(t *testing.T) { + tests := []struct { + name string + traceContent string + writePages bool + }{ + {name: "invalid trace", traceContent: "{"}, + {name: "missing pages", traceContent: `{"version":1,"page_size":4096,"offsets":[0]}`}, + {name: "wrong page data size", traceContent: `{"version":1,"page_size":4096,"offsets":[0]}`, writePages: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + baseDir := t.TempDir() + vmmStatePath := filepath.Join(baseDir, "snap_file") + tracePath := filepath.Join(baseDir, "working_set_trace") + writeTestFile(t, vmmStatePath, "state") + writeTestFile(t, tracePath, tt.traceContent) + if tt.writePages { + writeTestFile(t, filepath.Join(baseDir, "working_set_pages"), "pages") + } + + state := NewSnapshotState(SnapshotStateCfg{ + BaseDir: baseDir, + VMMStatePath: vmmStatePath, + WorkingSetPath: filepath.Join(baseDir, "working_set_pages"), + WorkingSetTracePath: tracePath, + }) + if err := state.fetchState(); err == nil { + t.Fatal("fetchState succeeded for invalid working set artifacts") + } + if state.isRecordReady { + t.Fatal("invalid working set artifacts were marked ready") + } + }) + } +} + func TestMemoryManagerActivateReceivesFirecrackerMappings(t *testing.T) { baseDir := t.TempDir() vmID := "vm-activate" diff --git a/memory/manager/snapshot_state.go b/memory/manager/snapshot_state.go index cd0df11af..c69982030 100644 --- a/memory/manager/snapshot_state.go +++ b/memory/manager/snapshot_state.go @@ -114,13 +114,6 @@ func normalizeSnapshotStateCfg(cfg SnapshotStateCfg) SnapshotStateCfg { func (s *SnapshotState) refreshSnapshotLoad(cfg SnapshotStateCfg) { cfg = normalizeSnapshotStateCfg(cfg) - trace := s.trace - if trace == nil { - trace = initTrace(cfg.WorkingSetTracePath) - } else { - trace.traceFileName = cfg.WorkingSetTracePath - } - isRecordReady := s.isRecordReady isEverActivated := s.isEverActivated totalPFServed := s.totalPFServed uniquePFServed := s.uniquePFServed @@ -131,14 +124,14 @@ func (s *SnapshotState) refreshSnapshotLoad(cfg SnapshotStateCfg) { s.firstPageFaultOnce = nil s.userFaultFD = nil s.guestRegionMappings = nil - s.trace = trace + s.trace = initTrace(cfg.WorkingSetTracePath) s.epfd = 0 s.wakeFD = -1 s.quitCh = nil s.pollDoneCh = nil s.isEverActivated = isEverActivated s.isActive = false - s.isRecordReady = isRecordReady + s.isRecordReady = false s.guestMem = nil s.workingSet = nil s.totalPFServed = totalPFServed @@ -225,17 +218,30 @@ func (s *SnapshotState) unmapGuestMemory() error { return nil } -// fetchState verifies snapshot state and loads the replay working set when ready. +// fetchState verifies snapshot state and loads a persisted working set when available. func (s *SnapshotState) fetchState() error { if _, err := os.ReadFile(s.VMMStatePath); err != nil { log.Errorf("Failed to fetch VMM state: %v\n", err) return err } - if s.isRecordReady && !s.IsLazyMode { - return s.fetchWorkingSet() + s.isRecordReady = false + s.workingSet = nil + s.trace.reset() + if err := s.trace.readTrace(); err != nil { + if os.IsNotExist(err) { + return nil + } + return fmt.Errorf("failed to read working set trace: %w", err) } + if !s.IsLazyMode { + if err := s.fetchWorkingSet(); err != nil { + return err + } + } + + s.isRecordReady = true return nil } @@ -249,10 +255,6 @@ func (s *SnapshotState) fetchWorkingSet() error { if size > uint64(int(^uint(0)>>1)) { return fmt.Errorf("working set too large: %#x", size) } - if size == 0 { - s.workingSet = nil - return nil - } f, err := os.Open(s.WorkingSetPath) if err != nil { @@ -260,6 +262,17 @@ func (s *SnapshotState) fetchWorkingSet() error { return err } defer func() { _ = f.Close() }() + info, err := f.Stat() + if err != nil { + return err + } + if info.Size() < 0 || uint64(info.Size()) != size { + return fmt.Errorf("working set size is %#x, want %#x", info.Size(), size) + } + if size == 0 { + s.workingSet = nil + return nil + } s.workingSet = make([]byte, int(size)) n, err := io.ReadFull(f, s.workingSet) diff --git a/memory/manager/snapshot_state_test.go b/memory/manager/snapshot_state_test.go index df814d225..e02bf172f 100644 --- a/memory/manager/snapshot_state_test.go +++ b/memory/manager/snapshot_state_test.go @@ -213,11 +213,12 @@ func TestTraceProcessRecordWritesWorkingSet(t *testing.T) { baseDir := t.TempDir() guestMemPath := filepath.Join(baseDir, "guest_mem") workingSetPath := filepath.Join(baseDir, "working_set_pages") + tracePath := filepath.Join(baseDir, "working_set_trace") pageSize := uint64(os.Getpagesize()) prepareGuestMemoryFile(t, guestMemPath, 5*int(pageSize)) - trace := initTrace(filepath.Join(baseDir, "trace")) + trace := initTrace(tracePath) trace.AppendRecord(Record{offset: 3 * pageSize}) trace.AppendRecord(Record{offset: pageSize}) trace.AppendRecord(Record{offset: 2 * pageSize}) @@ -246,6 +247,20 @@ func TestTraceProcessRecordWritesWorkingSet(t *testing.T) { if !reflect.DeepEqual(got, want) { t.Fatal("working set contents do not match recorded guest memory pages") } + + loadedTrace := initTrace(tracePath) + if err := loadedTrace.readTrace(); err != nil { + t.Fatalf("readTrace returned error: %v", err) + } + if got, want := loadedTrace.pageSize, pageSize; got != want { + t.Fatalf("loaded trace page size = %#x, want %#x", got, want) + } + if !reflect.DeepEqual(loadedTrace.trace, trace.trace) { + t.Fatal("loaded trace records do not match persisted records") + } + if !reflect.DeepEqual(loadedTrace.regions, trace.regions) { + t.Fatal("loaded trace regions do not match persisted regions") + } } func TestReceiveUffdMappingsAndFD(t *testing.T) { diff --git a/memory/manager/trace.go b/memory/manager/trace.go index 1bfe3f0a8..86b239c60 100644 --- a/memory/manager/trace.go +++ b/memory/manager/trace.go @@ -23,16 +23,17 @@ package manager import ( - "encoding/csv" - "errors" + "encoding/json" "fmt" "io" "os" + "path/filepath" "sort" - "strconv" "sync" ) +const workingSetTraceVersion = 1 + // Record identifies one guest memory page by its offset in the full memory file. type Record struct { offset uint64 @@ -49,6 +50,12 @@ type Trace struct { regions map[uint64]int } +type workingSetTraceMetadata struct { + Version int `json:"version"` + PageSize uint64 `json:"page_size"` + Offsets []uint64 `json:"offsets"` +} + func initTrace(traceFileName string) *Trace { return &Trace{ traceFileName: traceFileName, @@ -69,60 +76,52 @@ func (t *Trace) AppendRecord(r Record) { t.containedOffsets[r.offset] = struct{}{} } -func (t *Trace) WriteTrace() error { - t.Lock() - defer t.Unlock() - - file, err := os.Create(t.traceFileName) - if err != nil { - return err - } - defer func() { _ = file.Close() }() - - writer := csv.NewWriter(file) - for _, rec := range t.trace { - if err := writer.Write([]string{strconv.FormatUint(rec.offset, 16)}); err != nil { - return err - } - } - writer.Flush() - return writer.Error() -} - -//nolint:unused func (t *Trace) readTrace() error { - f, err := os.Open(t.traceFileName) + data, err := os.ReadFile(t.traceFileName) if err != nil { return err } - defer func() { _ = f.Close() }() - lines, err := csv.NewReader(f).ReadAll() - if err != nil { + var metadata workingSetTraceMetadata + if err := json.Unmarshal(data, &metadata); err != nil { return err } + if metadata.Version != workingSetTraceVersion { + return fmt.Errorf("unsupported working set trace version: %d", metadata.Version) + } + if metadata.PageSize == 0 { + return errInvalidGuestRegionPageSize + } - for _, line := range lines { - rec, err := readRecord(line) - if err != nil { - return err + containedOffsets := make(map[uint64]struct{}, len(metadata.Offsets)) + records := make([]Record, 0, len(metadata.Offsets)) + for _, offset := range metadata.Offsets { + if _, ok := containedOffsets[offset]; ok { + return fmt.Errorf("duplicate working set trace offset: %#x", offset) } - t.AppendRecord(rec) + records = append(records, Record{offset: offset}) + containedOffsets[offset] = struct{}{} } + + t.Lock() + defer t.Unlock() + + t.pageSize = metadata.PageSize + t.containedOffsets = containedOffsets + t.trace = records + t.buildRegionsLocked() + return nil } -//nolint:unused -func readRecord(line []string) (Record, error) { - if len(line) == 0 { - return Record{}, errors.New("empty trace record") - } - offset, err := strconv.ParseUint(line[0], 16, 64) - if err != nil { - return Record{}, err - } +func (t *Trace) reset() { + t.Lock() + defer t.Unlock() - return Record{offset: offset}, nil + t.pageSize = 0 + t.containedOffsets = make(map[uint64]struct{}) + t.trace = make([]Record, 0) + t.regions = make(map[uint64]int) } func (t *Trace) containsRecord(rec Record) bool { @@ -139,6 +138,15 @@ func (t *Trace) ProcessRecord(guestMemPath, workingSetPath string, pageSize uint defer t.Unlock() t.pageSize = pageSize + t.buildRegionsLocked() + + if err := t.writeWorkingSetPagesToFileLocked(guestMemPath, workingSetPath, pageSize); err != nil { + return err + } + return t.writeTraceLocked() +} + +func (t *Trace) buildRegionsLocked() { sort.Slice(t.trace, func(i, j int) bool { return t.trace[i].offset < t.trace[j].offset }) @@ -149,7 +157,7 @@ func (t *Trace) ProcessRecord(guestMemPath, workingSetPath string, pageSize uint regionStart uint64 ) for i, rec := range t.trace { - if i == 0 || rec.offset != last+pageSize { + if i == 0 || rec.offset != last+t.pageSize { regionStart = rec.offset t.regions[regionStart] = 1 } else { @@ -157,22 +165,35 @@ func (t *Trace) ProcessRecord(guestMemPath, workingSetPath string, pageSize uint } last = rec.offset } - - return t.writeWorkingSetPagesToFileLocked(guestMemPath, workingSetPath, pageSize) } -func (t *Trace) writeWorkingSetPagesToFileLocked(guestMemPath, workingSetPath string, pageSize uint64) error { - fSrc, err := os.Open(guestMemPath) +func (t *Trace) writeTraceLocked() error { + offsets := make([]uint64, len(t.trace)) + for i, rec := range t.trace { + offsets[i] = rec.offset + } + + data, err := json.Marshal(workingSetTraceMetadata{ + Version: workingSetTraceVersion, + PageSize: t.pageSize, + Offsets: offsets, + }) if err != nil { return err } - defer func() { _ = fSrc.Close() }() - fDst, err := os.Create(workingSetPath) + return writeFileAtomically(t.traceFileName, func(file *os.File) error { + _, err := file.Write(data) + return err + }) +} + +func (t *Trace) writeWorkingSetPagesToFileLocked(guestMemPath, workingSetPath string, pageSize uint64) error { + fSrc, err := os.Open(guestMemPath) if err != nil { return err } - defer func() { _ = fDst.Close() }() + defer func() { _ = fSrc.Close() }() keys := make([]uint64, 0, len(t.regions)) for k := range t.regions { @@ -180,31 +201,56 @@ func (t *Trace) writeWorkingSetPagesToFileLocked(guestMemPath, workingSetPath st } sort.Slice(keys, func(i, j int) bool { return keys[i] < keys[j] }) - var dstOffset int64 - for _, offset := range keys { - copyLen := uint64(t.regions[offset]) * pageSize - if copyLen > uint64(int(^uint(0)>>1)) { - return fmt.Errorf("working set region too large: %#x", copyLen) - } + return writeFileAtomically(workingSetPath, func(fDst *os.File) error { + var dstOffset int64 + for _, offset := range keys { + copyLen := uint64(t.regions[offset]) * pageSize + if copyLen > uint64(int(^uint(0)>>1)) { + return fmt.Errorf("working set region too large: %#x", copyLen) + } - buf := make([]byte, int(copyLen)) - n, err := fSrc.ReadAt(buf, int64(offset)) - if err != nil && err != io.EOF { - return err - } - if n != len(buf) { - return io.ErrUnexpectedEOF - } + buf := make([]byte, int(copyLen)) + n, err := fSrc.ReadAt(buf, int64(offset)) + if err != nil && err != io.EOF { + return err + } + if n != len(buf) { + return io.ErrUnexpectedEOF + } - n, err = fDst.WriteAt(buf, dstOffset) - if err != nil { - return err + n, err = fDst.WriteAt(buf, dstOffset) + if err != nil { + return err + } + if n != len(buf) { + return io.ErrShortWrite + } + dstOffset += int64(copyLen) } - if n != len(buf) { - return io.ErrShortWrite - } - dstOffset += int64(copyLen) + + return nil + }) +} + +func writeFileAtomically(path string, write func(*os.File) error) error { + tempFile, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+"-*") + if err != nil { + return err + } + tempPath := tempFile.Name() + defer func() { _ = os.Remove(tempPath) }() + + if err := write(tempFile); err != nil { + _ = tempFile.Close() + return err + } + if err := tempFile.Sync(); err != nil { + _ = tempFile.Close() + return err + } + if err := tempFile.Close(); err != nil { + return err } - return fDst.Sync() + return os.Rename(tempPath, path) } From 2abf0fef93604ec4c50cf8af308338fdd9c6867d Mon Sep 17 00:00:00 2001 From: Liquidzk <492928615@qq.com> Date: Tue, 21 Jul 2026 14:29:04 +0800 Subject: [PATCH 3/6] Enable UPF working set mode Signed-off-by: Liquidzk <492928615@qq.com> --- ctriface/iface_test.go | 27 ++++++++++++++++++++++++++- ctriface/orch_options.go | 6 +++--- vhive.go | 5 ----- vhive_test.go | 4 ++-- 4 files changed, 31 insertions(+), 11 deletions(-) diff --git a/ctriface/iface_test.go b/ctriface/iface_test.go index 41be55bec..23305602d 100644 --- a/ctriface/iface_test.go +++ b/ctriface/iface_test.go @@ -37,7 +37,6 @@ import ( "github.com/vhive-serverless/vhive/snapshotting" ) -// TODO: Make it impossible to use lazy mode without UPF var ( isUPFEnabled = flag.Bool("upf", false, "Set UPF enabled") isLazyMode = flag.Bool("lazy", false, "Set lazy serving on or off") @@ -54,6 +53,32 @@ func TestMain(m *testing.M) { os.Exit(m.Run()) } +func TestValidateUPFMode(t *testing.T) { + tests := []struct { + name string + upf bool + lazy bool + wantErr bool + }{ + {name: "disabled", upf: false, lazy: false}, + {name: "working set", upf: true, lazy: false}, + {name: "lazy", upf: true, lazy: true}, + {name: "lazy without UPF", upf: false, lazy: true, wantErr: true}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + orch := &Orchestrator{isUPFEnabled: tt.upf, isLazyMode: tt.lazy} + err := orch.validateUPFMode() + if tt.wantErr { + require.ErrorIs(t, err, errLazyModeRequiresUPF) + return + } + require.NoError(t, err) + }) + } +} + func TestStartSnapStopLoad(t *testing.T) { log.SetFormatter(&log.TextFormatter{ TimestampFormat: ctrdlog.RFC3339NanoFixed, diff --git a/ctriface/orch_options.go b/ctriface/orch_options.go index 385d43d04..3343da2f9 100644 --- a/ctriface/orch_options.go +++ b/ctriface/orch_options.go @@ -28,7 +28,7 @@ import ( "fmt" ) -var errUPFRequiresLazyMode = errors.New("UPF currently requires lazy mode") +var errLazyModeRequiresUPF = errors.New("lazy mode requires UPF") // OrchestratorOption Options to pass to Orchestrator type OrchestratorOption func(*Orchestrator) @@ -73,8 +73,8 @@ func WithLazyMode(isLazyMode bool) OrchestratorOption { } func (o *Orchestrator) validateUPFMode() error { - if o.isUPFEnabled && !o.isLazyMode { - return errUPFRequiresLazyMode + if o.isLazyMode && !o.isUPFEnabled { + return errLazyModeRequiresUPF } return nil } diff --git a/vhive.go b/vhive.go index be784825e..913d687de 100644 --- a/vhive.go +++ b/vhive.go @@ -97,11 +97,6 @@ func main() { log.Error("User-level page faults are not supported without snapshots") return } - if *isUPFEnabled && !*isLazyMode { - log.Error("User-level page faults currently require lazy serving mode") - return - } - if !*isUPFEnabled && *isLazyMode { log.Error("Lazy page fault serving mode is not supported without user-level page faults") return diff --git a/vhive_test.go b/vhive_test.go index d663ee294..c16770250 100644 --- a/vhive_test.go +++ b/vhive_test.go @@ -64,8 +64,8 @@ func TestMain(m *testing.M) { log.SetLevel(log.InfoLevel) flag.Parse() - if *isUPFEnabledTest && !*isLazyModeTest { - log.Error("User-level page faults currently require lazy serving mode") + if !*isUPFEnabledTest && *isLazyModeTest { + log.Error("Lazy page fault serving mode is not supported without user-level page faults") os.Exit(-1) } From 01797990c157996f2ccd646da735053b39cd3e8a Mon Sep 17 00:00:00 2001 From: Liquidzk <492928615@qq.com> Date: Tue, 21 Jul 2026 14:43:48 +0800 Subject: [PATCH 4/6] Test UPF working set record and replay Signed-off-by: Liquidzk <492928615@qq.com> --- ctriface/Makefile | 4 +++ ctriface/iface_test.go | 79 ++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 83 insertions(+) diff --git a/ctriface/Makefile b/ctriface/Makefile index 182b3a8a2..e169e8272 100644 --- a/ctriface/Makefile +++ b/ctriface/Makefile @@ -24,6 +24,7 @@ EXTRAGOARGS:=-v -race -cover EXTRATESTFILES:=iface_test.go iface.go orch_options.go orch.go BENCHFILES:=bench_test.go iface.go orch_options.go orch.go UPFARGS:=-upf -lazy +REAPARGS:=-upf STARGZ:=-ss 'proxy' -img 'ghcr.io/vhive-serverless/helloworld:var_workload-esgz' DOCKER_CREDENTIALS:=-dockerCredentials '{"docker-credentials":{"ghcr.io":{"username":"","password":""}}}' GOBENCH:=-v -timeout 1500s @@ -37,6 +38,9 @@ test: sudo mkdir -m777 -p $(CTRDLOGDIR) && sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/ctriface_log.out 2>$(CTRDLOGDIR)/ctriface_log.err & sudo env "PATH=$(PATH)" go test $(EXTRATESTFILES) $(EXTRAGOARGS) -args $(UPFARGS) ./../scripts/clean_fcctr.sh + sudo mkdir -m777 -p $(CTRDLOGDIR) && sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/ctriface_log.out 2>$(CTRDLOGDIR)/ctriface_log.err & + sudo env "PATH=$(PATH)" go test $(EXTRATESTFILES) $(EXTRAGOARGS) -run '^TestUPFWorkingSetRecordReplay$$' -args $(REAPARGS) + ./../scripts/clean_fcctr.sh test-man: ./../scripts/clean_fcctr.sh diff --git a/ctriface/iface_test.go b/ctriface/iface_test.go index 23305602d..03c5da380 100644 --- a/ctriface/iface_test.go +++ b/ctriface/iface_test.go @@ -22,6 +22,7 @@ package ctriface import ( + "bytes" "context" "flag" "fmt" @@ -135,6 +136,84 @@ func TestStartSnapStopLoad(t *testing.T) { orch.Cleanup() } +func TestUPFWorkingSetRecordReplay(t *testing.T) { + if !*isUPFEnabled || *isLazyMode { + t.Skip("requires UPF working set mode") + } + + log.SetFormatter(&log.TextFormatter{ + TimestampFormat: ctrdlog.RFC3339NanoFixed, + FullTimestamp: true, + }) + log.SetOutput(os.Stdout) + log.SetLevel(log.InfoLevel) + + ctx, cancel := context.WithTimeout(namespaces.WithNamespace(context.Background(), namespaceName), 5*time.Minute) + defer cancel() + + orch := NewOrchestrator( + *snapshotter, + "", + WithTestModeOn(true), + WithUPF(true), + WithLazyMode(false), + WithMetricsMode(true), + WithDockerCredentials(*dockerCredentials), + ) + t.Cleanup(func() { + _ = orch.StopActiveVMs() + orch.Cleanup() + }) + + sourceVMID := "ws-source" + recordVMID := "ws-record" + replayVMID := "ws-replay" + snap := snapshotting.NewSnapshot("myrev-working-set", "/fccd/snapshots", *testImage) + + _, _, err := orch.StartVM(ctx, sourceVMID, *testImage) + require.NoError(t, err, "Failed to start source VM") + require.NoError(t, orch.PauseVM(ctx, sourceVMID), "Failed to pause source VM") + require.NoError(t, snap.CreateSnapDir(), "Failed to create snapshot directory") + require.NoError(t, orch.CreateSnapshot(ctx, sourceVMID, snap), "Failed to create snapshot") + require.NoError(t, orch.StopSingleVM(ctx, sourceVMID), "Failed to stop source VM") + + loadResumeStop := func(vmID string) { + _, _, err := orch.LoadSnapshot(ctx, vmID, snap) + require.NoError(t, err, "Failed to load snapshot into %s", vmID) + _, err = orch.ResumeVM(ctx, vmID) + require.NoError(t, err, "Failed to resume %s", vmID) + time.Sleep(500 * time.Millisecond) + require.NoError(t, orch.StopSingleVM(ctx, vmID), "Failed to stop %s", vmID) + } + + loadResumeStop(recordVMID) + workingSetBefore, err := os.ReadFile(snap.GetWorkingSetFilePath()) + require.NoError(t, err, "Failed to read recorded working set") + require.NotEmpty(t, workingSetBefore, "Recorded working set is empty") + traceBefore, err := os.ReadFile(snap.GetWorkingSetTraceFilePath()) + require.NoError(t, err, "Failed to read recorded working set trace") + require.NotEmpty(t, traceBefore, "Recorded working set trace is empty") + + loadResumeStop(replayVMID) + latencyMetrics, err := orch.GetUPFLatencyStats(replayVMID) + require.NoError(t, err, "Failed to get replay metrics") + workingSetInstalled := false + for _, metric := range latencyMetrics { + if _, ok := metric.MetricMap["InstallWS"]; ok { + workingSetInstalled = true + break + } + } + require.True(t, workingSetInstalled, "Working set was not installed during replay") + + workingSetAfter, err := os.ReadFile(snap.GetWorkingSetFilePath()) + require.NoError(t, err, "Failed to reread working set") + traceAfter, err := os.ReadFile(snap.GetWorkingSetTraceFilePath()) + require.NoError(t, err, "Failed to reread working set trace") + require.True(t, bytes.Equal(workingSetBefore, workingSetAfter), "Replay rewrote working set pages") + require.True(t, bytes.Equal(traceBefore, traceAfter), "Replay rewrote working set trace") +} + func TestPauseSnapResume(t *testing.T) { log.SetFormatter(&log.TextFormatter{ TimestampFormat: ctrdlog.RFC3339NanoFixed, From 49f439a80d49042a6ae5304926f2a3fee3cb07d1 Mon Sep 17 00:00:00 2001 From: Liquidzk <492928615@qq.com> Date: Wed, 22 Jul 2026 21:00:39 +0800 Subject: [PATCH 5/6] Refine UPF working set tests Signed-off-by: Liquidzk <492928615@qq.com> --- ctriface/Makefile | 10 ++++--- ctriface/iface_test.go | 41 ++++++++++++++++----------- memory/manager/manager_test.go | 12 -------- memory/manager/snapshot_state_test.go | 2 +- snapshotting/manager_test.go | 12 -------- 5 files changed, 31 insertions(+), 46 deletions(-) diff --git a/ctriface/Makefile b/ctriface/Makefile index e169e8272..4a3dc1d69 100644 --- a/ctriface/Makefile +++ b/ctriface/Makefile @@ -24,7 +24,7 @@ EXTRAGOARGS:=-v -race -cover EXTRATESTFILES:=iface_test.go iface.go orch_options.go orch.go BENCHFILES:=bench_test.go iface.go orch_options.go orch.go UPFARGS:=-upf -lazy -REAPARGS:=-upf +UPF_WORKING_SET_ARGS:=-upf STARGZ:=-ss 'proxy' -img 'ghcr.io/vhive-serverless/helloworld:var_workload-esgz' DOCKER_CREDENTIALS:=-dockerCredentials '{"docker-credentials":{"ghcr.io":{"username":"","password":""}}}' GOBENCH:=-v -timeout 1500s @@ -38,9 +38,11 @@ test: sudo mkdir -m777 -p $(CTRDLOGDIR) && sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/ctriface_log.out 2>$(CTRDLOGDIR)/ctriface_log.err & sudo env "PATH=$(PATH)" go test $(EXTRATESTFILES) $(EXTRAGOARGS) -args $(UPFARGS) ./../scripts/clean_fcctr.sh - sudo mkdir -m777 -p $(CTRDLOGDIR) && sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/ctriface_log.out 2>$(CTRDLOGDIR)/ctriface_log.err & - sudo env "PATH=$(PATH)" go test $(EXTRATESTFILES) $(EXTRAGOARGS) -run '^TestUPFWorkingSetRecordReplay$$' -args $(REAPARGS) - ./../scripts/clean_fcctr.sh + @set -e; \ + trap './../scripts/clean_fcctr.sh' EXIT; \ + sudo mkdir -m777 -p $(CTRDLOGDIR); \ + sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/ctriface_log_upf_working_set.out 2>$(CTRDLOGDIR)/ctriface_log_upf_working_set.err & \ + sudo env "PATH=$(PATH)" go test $(EXTRATESTFILES) $(EXTRAGOARGS) -run '^TestUPFWorkingSetRecordReplay$$' -args $(UPF_WORKING_SET_ARGS) test-man: ./../scripts/clean_fcctr.sh diff --git a/ctriface/iface_test.go b/ctriface/iface_test.go index 03c5da380..20c67ac50 100644 --- a/ctriface/iface_test.go +++ b/ctriface/iface_test.go @@ -22,10 +22,10 @@ package ctriface import ( - "bytes" "context" "flag" "fmt" + "net" "os" "sync" "testing" @@ -141,13 +141,6 @@ func TestUPFWorkingSetRecordReplay(t *testing.T) { t.Skip("requires UPF working set mode") } - log.SetFormatter(&log.TextFormatter{ - TimestampFormat: ctrdlog.RFC3339NanoFixed, - FullTimestamp: true, - }) - log.SetOutput(os.Stdout) - log.SetLevel(log.InfoLevel) - ctx, cancel := context.WithTimeout(namespaces.WithNamespace(context.Background(), namespaceName), 5*time.Minute) defer cancel() @@ -168,6 +161,7 @@ func TestUPFWorkingSetRecordReplay(t *testing.T) { sourceVMID := "ws-source" recordVMID := "ws-record" replayVMID := "ws-replay" + const installWorkingSetMetric = "InstallWS" snap := snapshotting.NewSnapshot("myrev-working-set", "/fccd/snapshots", *testImage) _, _, err := orch.StartVM(ctx, sourceVMID, *testImage) @@ -178,11 +172,18 @@ func TestUPFWorkingSetRecordReplay(t *testing.T) { require.NoError(t, orch.StopSingleVM(ctx, sourceVMID), "Failed to stop source VM") loadResumeStop := func(vmID string) { - _, _, err := orch.LoadSnapshot(ctx, vmID, snap) + response, _, err := orch.LoadSnapshot(ctx, vmID, snap) require.NoError(t, err, "Failed to load snapshot into %s", vmID) _, err = orch.ResumeVM(ctx, vmID) require.NoError(t, err, "Failed to resume %s", vmID) - time.Sleep(500 * time.Millisecond) + require.Eventually(t, func() bool { + conn, err := net.DialTimeout("tcp", net.JoinHostPort(response.GuestIP, "50051"), 200*time.Millisecond) + if err != nil { + return false + } + _ = conn.Close() + return true + }, 30*time.Second, 100*time.Millisecond, "%s workload did not become ready", vmID) require.NoError(t, orch.StopSingleVM(ctx, vmID), "Failed to stop %s", vmID) } @@ -193,25 +194,31 @@ func TestUPFWorkingSetRecordReplay(t *testing.T) { traceBefore, err := os.ReadFile(snap.GetWorkingSetTraceFilePath()) require.NoError(t, err, "Failed to read recorded working set trace") require.NotEmpty(t, traceBefore, "Recorded working set trace is empty") + workingSetInfoBefore, err := os.Stat(snap.GetWorkingSetFilePath()) + require.NoError(t, err, "Failed to stat recorded working set") + traceInfoBefore, err := os.Stat(snap.GetWorkingSetTraceFilePath()) + require.NoError(t, err, "Failed to stat recorded working set trace") loadResumeStop(replayVMID) latencyMetrics, err := orch.GetUPFLatencyStats(replayVMID) require.NoError(t, err, "Failed to get replay metrics") workingSetInstalled := false for _, metric := range latencyMetrics { - if _, ok := metric.MetricMap["InstallWS"]; ok { + if _, ok := metric.MetricMap[installWorkingSetMetric]; ok { workingSetInstalled = true break } } require.True(t, workingSetInstalled, "Working set was not installed during replay") - workingSetAfter, err := os.ReadFile(snap.GetWorkingSetFilePath()) - require.NoError(t, err, "Failed to reread working set") - traceAfter, err := os.ReadFile(snap.GetWorkingSetTraceFilePath()) - require.NoError(t, err, "Failed to reread working set trace") - require.True(t, bytes.Equal(workingSetBefore, workingSetAfter), "Replay rewrote working set pages") - require.True(t, bytes.Equal(traceBefore, traceAfter), "Replay rewrote working set trace") + workingSetInfoAfter, err := os.Stat(snap.GetWorkingSetFilePath()) + require.NoError(t, err, "Failed to restat working set") + traceInfoAfter, err := os.Stat(snap.GetWorkingSetTraceFilePath()) + require.NoError(t, err, "Failed to restat working set trace") + require.True(t, os.SameFile(workingSetInfoBefore, workingSetInfoAfter), "Replay replaced working set pages") + require.True(t, os.SameFile(traceInfoBefore, traceInfoAfter), "Replay replaced working set trace") + require.Equal(t, workingSetInfoBefore.ModTime(), workingSetInfoAfter.ModTime(), "Replay modified working set pages") + require.Equal(t, traceInfoBefore.ModTime(), traceInfoAfter.ModTime(), "Replay modified working set trace") } func TestPauseSnapResume(t *testing.T) { diff --git a/memory/manager/manager_test.go b/memory/manager/manager_test.go index b1f4efcbc..08138a969 100644 --- a/memory/manager/manager_test.go +++ b/memory/manager/manager_test.go @@ -217,18 +217,6 @@ func TestFetchStateLoadsWorkingSetAcrossVMIDs(t *testing.T) { if !replayState.isRecordReady { t.Fatal("FetchState did not mark the persisted working set ready") } - if got, want := replayState.trace.pageSize, pageSize; got != want { - t.Fatalf("trace page size = %#x, want %#x", got, want) - } - if got, want := len(replayState.trace.trace), 2; got != want { - t.Fatalf("trace length = %d, want %d", got, want) - } - if got, want := replayState.trace.trace[0].offset, pageSize; got != want { - t.Fatalf("first trace offset = %#x, want %#x", got, want) - } - if got, want := replayState.trace.trace[1].offset, 3*pageSize; got != want { - t.Fatalf("second trace offset = %#x, want %#x", got, want) - } guestMem, err := os.ReadFile(guestMemPath) if err != nil { diff --git a/memory/manager/snapshot_state_test.go b/memory/manager/snapshot_state_test.go index e02bf172f..cfdcf295e 100644 --- a/memory/manager/snapshot_state_test.go +++ b/memory/manager/snapshot_state_test.go @@ -209,7 +209,7 @@ func TestPageFaultCopyArgsForGuestOffsetOutsideAllRegions(t *testing.T) { } } -func TestTraceProcessRecordWritesWorkingSet(t *testing.T) { +func TestTraceProcessRecordPersistsWorkingSetAndTrace(t *testing.T) { baseDir := t.TempDir() guestMemPath := filepath.Join(baseDir, "guest_mem") workingSetPath := filepath.Join(baseDir, "working_set_pages") diff --git a/snapshotting/manager_test.go b/snapshotting/manager_test.go index 26e163e2e..81f43e221 100644 --- a/snapshotting/manager_test.go +++ b/snapshotting/manager_test.go @@ -24,8 +24,6 @@ package snapshotting_test import ( "fmt" - "path/filepath" - ctrdlog "github.com/containerd/log" log "github.com/sirupsen/logrus" "github.com/stretchr/testify/require" @@ -81,16 +79,6 @@ func TestSnapshotManagerSingle(t *testing.T) { testSnapshotManager(t, mgr, revision, imageName) } -func TestSnapshotWorkingSetPaths(t *testing.T) { - baseDir := t.TempDir() - revision := "myrevision-working-set" - snap := snapshotting.NewSnapshot(revision, baseDir, "testImage") - snapshotDir := filepath.Join(baseDir, revision) - - require.Equal(t, filepath.Join(snapshotDir, "working_set_pages"), snap.GetWorkingSetFilePath()) - require.Equal(t, filepath.Join(snapshotDir, "working_set_trace"), snap.GetWorkingSetTraceFilePath()) -} - func TestSnapshotManagerConcurrent(t *testing.T) { // Create snapshot manager mgr := snapshotting.NewSnapshotManager(snapshotsDir) From f82cee6807e554ab6c190aba0a4ba689a243a98d Mon Sep 17 00:00:00 2001 From: Liquidzk <492928615@qq.com> Date: Sun, 26 Jul 2026 22:21:42 +0800 Subject: [PATCH 6/6] Address UPF working set review comments Signed-off-by: Liquidzk <492928615@qq.com> --- .github/workflows/nightly_tests.yml | 4 +-- Makefile | 11 +++------ bench_test.go | 18 +++++++++----- ctriface/Makefile | 6 ++--- ctriface/iface.go | 4 --- ctriface/iface_test.go | 6 +++++ ctriface/orch.go | 7 ++++-- memory/manager/manager.go | 12 +++++++-- memory/manager/manager_test.go | 21 ++++++++++++++-- memory/manager/snapshot_state.go | 1 + memory/manager/snapshot_state_test.go | 22 +++++++++++++++++ memory/manager/trace.go | 35 ++++++++++++++++++--------- 12 files changed, 106 insertions(+), 41 deletions(-) diff --git a/.github/workflows/nightly_tests.yml b/.github/workflows/nightly_tests.yml index 50ae70376..ff2bcb399 100644 --- a/.github/workflows/nightly_tests.yml +++ b/.github/workflows/nightly_tests.yml @@ -56,9 +56,7 @@ jobs: strategy: fail-fast: false matrix: - # User-level page faults are temporarily disabled (gh-807) - # vhive_args: ["-dbg", "-dbg -snapshots", "-dbg -snapshots -upf"] - vhive_args: [ "-dbg", "-dbg -snapshots" ] + vhive_args: ["-dbg", "-dbg -snapshots", "-dbg -snapshots -upf"] env: GITHUB_RUN_ID: ${{ github.run_id }} GITHUB_JOB: ${{ github.job }} diff --git a/Makefile b/Makefile index 8e4f6f655..2ba929216 100644 --- a/Makefile +++ b/Makefile @@ -24,11 +24,8 @@ SUBDIRS:=ctriface taps misc profile EXTRAGOARGS:=-v -race -cover EXTRAGOARGS_NORACE:=-v EXTRATESTFILES:=vhive_test.go stats.go vhive.go functions.go -# User-level page faults are temporarily disabled (gh-807) -# WITHUPF:=-upfTest -# WITHLAZY:=-lazyTest -WITHUPF:= -WITHLAZY:= +WITHUPF:=-upfTest +WITHLAZY:=-lazyTest WITHSNAPSHOTS:=-snapshotsTest CTRDLOGDIR:=/tmp/ctrd-logs @@ -99,7 +96,7 @@ bench: sudo env "PATH=$(PATH)" go test $(EXTRAGOARGS) -run TestBenchServe -args -iter 1 $(WITHSNAPSHOTS) $(WITHUPF) -benchDirTest configREAP -metricsTest -funcName helloworld && sudo rm -rf configREAP ./scripts/clean_fcctr.sh sudo mkdir -m777 -p $(CTRDLOGDIR) && sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/fccd_orch_noupf_log_bench.out 2>$(CTRDLOGDIR)/fccd_orch_noupf_log_bench.err & - sudo env "PATH=$(PATH)" go test $(EXTRAGOARGS) -run TestBenchServe -args -iter 1 $(WITHSNAPSHOTS) $(WITHLAZY) -benchDirTest configLazy -metricsTest -funcName helloworld && sudo rm -rf configLazy + sudo env "PATH=$(PATH)" go test $(EXTRAGOARGS) -run TestBenchServe -args -iter 1 $(WITHSNAPSHOTS) $(WITHUPF) $(WITHLAZY) -benchDirTest configLazy -metricsTest -funcName helloworld && sudo rm -rf configLazy ./scripts/clean_fcctr.sh sudo mkdir -m777 -p $(CTRDLOGDIR) && sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/fccd_orch_noupf_log_bench.out 2>$(CTRDLOGDIR)/fccd_orch_noupf_log_bench.err & @@ -109,7 +106,7 @@ bench: sudo env "PATH=$(PATH)" go test $(EXTRAGOARGS) -run TestBenchParallelServe -args $(WITHSNAPSHOTS) $(WITHUPF) -benchDirTest configREAP -metricsTest -funcName helloworld && sudo rm -rf configREAP ./scripts/clean_fcctr.sh sudo mkdir -m777 -p $(CTRDLOGDIR) && sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/fccd_orch_noupf_log_bench.out 2>$(CTRDLOGDIR)/fccd_orch_noupf_log_bench.err & - sudo env "PATH=$(PATH)" go test $(EXTRAGOARGS) -run TestBenchParallelServe -args $(WITHSNAPSHOTS) $(WITHLAZY) -benchDirTest configLazy -metricsTest -funcName helloworld && sudo rm -rf configLazy + sudo env "PATH=$(PATH)" go test $(EXTRAGOARGS) -run TestBenchParallelServe -args $(WITHSNAPSHOTS) $(WITHUPF) $(WITHLAZY) -benchDirTest configLazy -metricsTest -funcName helloworld && sudo rm -rf configLazy ./scripts/clean_fcctr.sh test-man-bench: diff --git a/bench_test.go b/bench_test.go index 7e579c423..1e16a0844 100644 --- a/bench_test.go +++ b/bench_test.go @@ -112,9 +112,10 @@ func TestBenchParallelServe(t *testing.T) { require.NoError(t, err, "Function returned error, "+message) if *isUPFEnabledTest { - memManagerMetrics, err := orch.GetUPFLatencyStats(vmIDString + "-0") - require.NoError(t, err, "Failed to ge tupf metrics") - require.Equal(t, len(memManagerMetrics), 1, "wrong length") + f := funcPool.getFunction(vmIDString, imageName) + memManagerMetrics, err := orch.GetUPFLatencyStats(f.vmID) + require.NoError(t, err, "Failed to get UPF metrics") + require.Len(t, memManagerMetrics, 1, "wrong metrics length") upfMetrics[i] = memManagerMetrics[0] } } @@ -235,6 +236,14 @@ func TestBenchServe(t *testing.T) { message, err := funcPool.RemoveInstance(vmIDString, imageName, isSyncOffload) require.NoError(t, err, "Function returned error, "+message) + if orch.GetUPFEnabled() { + f := funcPool.getFunction(vmIDString, imageName) + instanceMetrics, err := orch.GetUPFLatencyStats(f.vmID) + require.NoError(t, err, "Failed to get UPF metrics for "+f.vmID) + require.Len(t, instanceMetrics, 1, "wrong metrics length for "+f.vmID) + memManagerMetrics = append(memManagerMetrics, instanceMetrics[0]) + } + time.Sleep(3 * time.Second) // this helps kworker hanging } @@ -243,9 +252,6 @@ func TestBenchServe(t *testing.T) { // Page stats err = funcPool.DumpUPFPageStats(vmIDString, imageName, *funcName, getOutFile("pageStats.csv")) require.NoError(t, err, "Failed to dump page stats for"+*funcName) - - memManagerMetrics, err = orch.GetUPFLatencyStats(vmIDString + "-0") - require.NoError(t, err, "Failed to dump get stats for "+*funcName) require.Equal(t, len(serveMetrics), len(memManagerMetrics), "different metrics lengths") } diff --git a/ctriface/Makefile b/ctriface/Makefile index 4a3dc1d69..e0b48db66 100644 --- a/ctriface/Makefile +++ b/ctriface/Makefile @@ -38,11 +38,9 @@ test: sudo mkdir -m777 -p $(CTRDLOGDIR) && sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/ctriface_log.out 2>$(CTRDLOGDIR)/ctriface_log.err & sudo env "PATH=$(PATH)" go test $(EXTRATESTFILES) $(EXTRAGOARGS) -args $(UPFARGS) ./../scripts/clean_fcctr.sh - @set -e; \ - trap './../scripts/clean_fcctr.sh' EXIT; \ - sudo mkdir -m777 -p $(CTRDLOGDIR); \ - sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/ctriface_log_upf_working_set.out 2>$(CTRDLOGDIR)/ctriface_log_upf_working_set.err & \ + sudo env "PATH=$(PATH)" /usr/local/bin/firecracker-containerd --config /etc/firecracker-containerd/config.toml 1>$(CTRDLOGDIR)/ctriface_log_upf_working_set.out 2>$(CTRDLOGDIR)/ctriface_log_upf_working_set.err & sudo env "PATH=$(PATH)" go test $(EXTRATESTFILES) $(EXTRAGOARGS) -run '^TestUPFWorkingSetRecordReplay$$' -args $(UPF_WORKING_SET_ARGS) + ./../scripts/clean_fcctr.sh test-man: ./../scripts/clean_fcctr.sh diff --git a/ctriface/iface.go b/ctriface/iface.go index e2f386088..b08926446 100644 --- a/ctriface/iface.go +++ b/ctriface/iface.go @@ -497,10 +497,6 @@ func configureSnapshotMemoryBackend(conf *proto.CreateVMRequest, backendType, ba // LoadSnapshot Loads a snapshot of a VM func (o *Orchestrator) LoadSnapshot(ctx context.Context, vmID string, snap *snapshotting.Snapshot) (_ *StartVMResponse, _ *metrics.Metric, retErr error) { - if err := o.validateUPFMode(); err != nil { - return nil, nil, err - } - var ( loadSnapshotMetric = metrics.NewMetric() tStart time.Time diff --git a/ctriface/iface_test.go b/ctriface/iface_test.go index 20c67ac50..a07da85b9 100644 --- a/ctriface/iface_test.go +++ b/ctriface/iface_test.go @@ -80,6 +80,12 @@ func TestValidateUPFMode(t *testing.T) { } } +func TestNewOrchestratorRejectsInvalidUPFMode(t *testing.T) { + require.PanicsWithValue(t, errLazyModeRequiresUPF, func() { + NewOrchestrator("", "", WithLazyMode(true)) + }) +} + func TestStartSnapStopLoad(t *testing.T) { log.SetFormatter(&log.TextFormatter{ TimestampFormat: ctrdlog.RFC3339NanoFixed, diff --git a/ctriface/orch.go b/ctriface/orch.go index 8c08cade7..e97468048 100644 --- a/ctriface/orch.go +++ b/ctriface/orch.go @@ -127,12 +127,15 @@ func NewOrchestrator(snapshotter, hostIface string, opts ...OrchestratorOption) o.vethPrefix = "172.17" o.clonePrefix = "172.18" - o.dns = getK8sDNS() - for _, opt := range opts { opt(o) } + if err := o.validateUPFMode(); err != nil { + panic(err) + } + + o.dns = getK8sDNS() o.vmPool = misc.NewVMPool(hostIface, o.netPoolSize, o.vethPrefix, o.clonePrefix, o.setExpIface) if _, err := os.Stat(o.snapshotsDir); err != nil { diff --git a/memory/manager/manager.go b/memory/manager/manager.go index 7994a6f1f..fbaa6267e 100644 --- a/memory/manager/manager.go +++ b/memory/manager/manager.go @@ -264,6 +264,9 @@ func (m *MemoryManager) Deactivate(vmID string) error { m.Unlock() + state.deactivateMu.Lock() + defer state.deactivateMu.Unlock() + if !state.isEverActivated { return nil } @@ -288,12 +291,17 @@ func (m *MemoryManager) Deactivate(vmID string) error { defer func() { _ = state.userFaultFD.Close() }() } - if !state.isRecordReady && !state.IsLazyMode { + if !state.isRecordReady { pageSize, err := guestMappingPageSize(state.guestRegionMappings) if err != nil { return err } - if err := state.trace.ProcessRecord(state.GuestMemPath, state.WorkingSetPath, pageSize); err != nil { + if state.IsLazyMode { + err = state.trace.persistTrace(pageSize) + } else { + err = state.trace.ProcessRecord(state.GuestMemPath, state.WorkingSetPath, pageSize) + } + if err != nil { return err } } diff --git a/memory/manager/manager_test.go b/memory/manager/manager_test.go index 08138a969..0bfc8dfdb 100644 --- a/memory/manager/manager_test.go +++ b/memory/manager/manager_test.go @@ -331,8 +331,25 @@ func TestMemoryManagerActivateReceivesFirecrackerMappings(t *testing.T) { t.Fatalf("validateGuestMemory mapped memory returned error: %v", err) } - if err := manager.Deactivate(vmID); err != nil { - t.Fatalf("Deactivate returned error: %v", err) + const deactivatorCount = 8 + start := make(chan struct{}) + deactivateErrCh := make(chan error, deactivatorCount) + for range deactivatorCount { + go func() { + <-start + deactivateErrCh <- manager.Deactivate(vmID) + }() + } + close(start) + + successes := 0 + for range deactivatorCount { + if err := <-deactivateErrCh; err == nil { + successes++ + } + } + if successes != 1 { + t.Fatalf("successful Deactivate calls = %d, want 1", successes) } if err := manager.DeregisterVM(vmID); err != nil { t.Fatalf("DeregisterVM returned error: %v", err) diff --git a/memory/manager/snapshot_state.go b/memory/manager/snapshot_state.go index c69982030..5a95f4041 100644 --- a/memory/manager/snapshot_state.go +++ b/memory/manager/snapshot_state.go @@ -53,6 +53,7 @@ type SnapshotStateCfg struct { // of the VM. type SnapshotState struct { SnapshotStateCfg + deactivateMu sync.Mutex firstPageFaultOnce *sync.Once userFaultFD *os.File guestRegionMappings []GuestRegionUffdMapping diff --git a/memory/manager/snapshot_state_test.go b/memory/manager/snapshot_state_test.go index cfdcf295e..80142b356 100644 --- a/memory/manager/snapshot_state_test.go +++ b/memory/manager/snapshot_state_test.go @@ -263,6 +263,28 @@ func TestTraceProcessRecordPersistsWorkingSetAndTrace(t *testing.T) { } } +func TestTracePersistTraceForLazyReplay(t *testing.T) { + baseDir := t.TempDir() + tracePath := filepath.Join(baseDir, "working_set_trace") + pageSize := uint64(os.Getpagesize()) + + trace := initTrace(tracePath) + trace.AppendRecord(Record{offset: 2 * pageSize}) + trace.AppendRecord(Record{offset: pageSize}) + + if err := trace.persistTrace(pageSize); err != nil { + t.Fatalf("persistTrace returned error: %v", err) + } + + loadedTrace := initTrace(tracePath) + if err := loadedTrace.readTrace(); err != nil { + t.Fatalf("readTrace returned error: %v", err) + } + if !reflect.DeepEqual(loadedTrace.trace, trace.trace) { + t.Fatal("loaded trace records do not match persisted records") + } +} + func TestReceiveUffdMappingsAndFD(t *testing.T) { mappings := []GuestRegionUffdMapping{{ BaseHostVirtAddr: 0x100000, diff --git a/memory/manager/trace.go b/memory/manager/trace.go index 86b239c60..ee36470ff 100644 --- a/memory/manager/trace.go +++ b/memory/manager/trace.go @@ -50,7 +50,7 @@ type Trace struct { regions map[uint64]int } -type workingSetTraceMetadata struct { +type serializedWorkingSet struct { Version int `json:"version"` PageSize uint64 `json:"page_size"` Offsets []uint64 `json:"offsets"` @@ -82,20 +82,20 @@ func (t *Trace) readTrace() error { return err } - var metadata workingSetTraceMetadata - if err := json.Unmarshal(data, &metadata); err != nil { + var workingSet serializedWorkingSet + if err := json.Unmarshal(data, &workingSet); err != nil { return err } - if metadata.Version != workingSetTraceVersion { - return fmt.Errorf("unsupported working set trace version: %d", metadata.Version) + if workingSet.Version != workingSetTraceVersion { + return fmt.Errorf("unsupported working set trace version: %d", workingSet.Version) } - if metadata.PageSize == 0 { + if workingSet.PageSize == 0 { return errInvalidGuestRegionPageSize } - containedOffsets := make(map[uint64]struct{}, len(metadata.Offsets)) - records := make([]Record, 0, len(metadata.Offsets)) - for _, offset := range metadata.Offsets { + containedOffsets := make(map[uint64]struct{}, len(workingSet.Offsets)) + records := make([]Record, 0, len(workingSet.Offsets)) + for _, offset := range workingSet.Offsets { if _, ok := containedOffsets[offset]; ok { return fmt.Errorf("duplicate working set trace offset: %#x", offset) } @@ -106,7 +106,7 @@ func (t *Trace) readTrace() error { t.Lock() defer t.Unlock() - t.pageSize = metadata.PageSize + t.pageSize = workingSet.PageSize t.containedOffsets = containedOffsets t.trace = records t.buildRegionsLocked() @@ -146,6 +146,19 @@ func (t *Trace) ProcessRecord(guestMemPath, workingSetPath string, pageSize uint return t.writeTraceLocked() } +func (t *Trace) persistTrace(pageSize uint64) error { + if pageSize == 0 { + return errInvalidGuestRegionPageSize + } + + t.Lock() + defer t.Unlock() + + t.pageSize = pageSize + t.buildRegionsLocked() + return t.writeTraceLocked() +} + func (t *Trace) buildRegionsLocked() { sort.Slice(t.trace, func(i, j int) bool { return t.trace[i].offset < t.trace[j].offset @@ -173,7 +186,7 @@ func (t *Trace) writeTraceLocked() error { offsets[i] = rec.offset } - data, err := json.Marshal(workingSetTraceMetadata{ + data, err := json.Marshal(serializedWorkingSet{ Version: workingSetTraceVersion, PageSize: t.pageSize, Offsets: offsets,