Likun/upf ws restore

.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 }}

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:

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")
 	}
 

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
+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
@@ -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 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

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
@@ -595,14 +591,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
 		}

ctriface/iface_test.go

@@ -25,6 +25,7 @@ import (
 	"context"
 	"flag"
 	"fmt"
+	"net"
 	"os"
 	"sync"
 	"testing"
@@ -37,7 +38,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 +54,38 @@ 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 TestNewOrchestratorRejectsInvalidUPFMode(t *testing.T) {
+	require.PanicsWithValue(t, errLazyModeRequiresUPF, func() {
+		NewOrchestrator("", "", WithLazyMode(true))
+	})
+}
+
 func TestStartSnapStopLoad(t *testing.T) {
 	log.SetFormatter(&log.TextFormatter{
 		TimestampFormat: ctrdlog.RFC3339NanoFixed,
@@ -110,6 +142,91 @@ func TestStartSnapStopLoad(t *testing.T) {
 	orch.Cleanup()
 }
 
+func TestUPFWorkingSetRecordReplay(t *testing.T) {
+	if !*isUPFEnabled || *isLazyMode {
+		t.Skip("requires UPF working set mode")
+	}
+
+	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"
+	const installWorkingSetMetric = "InstallWS"
+	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) {
+		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)
+		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)
+	}
+
+	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")
+	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[installWorkingSetMetric]; ok {
+			workingSetInstalled = true
+			break
+		}
+	}
+	require.True(t, workingSetInstalled, "Working set was not installed during replay")
+
+	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) {
 	log.SetFormatter(&log.TextFormatter{
 		TimestampFormat: ctrdlog.RFC3339NanoFixed,

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 {

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
 }

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
@@ -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
 		}
 	}