Test

cri/firecracker/coordinator.go

@@ -181,6 +181,7 @@ func (c *coordinator) orchLoadInstance(ctx context.Context, snap *snapshotting.S
 	ctxTimeout, cancel := context.WithTimeout(ctx, time.Second*30)
 	defer cancel()
 
+	logger.Debug("FIXME: temp pass same lastVmID")
 	resp, _, err := c.orch.LoadSnapshot(ctxTimeout, vmID, vmID, snap)
 	if err != nil {
 		logger.WithError(err).Error("failed to load VM")

ctriface/failing_test.go

@@ -69,7 +69,7 @@ func TestStartSnapStop(t *testing.T) {
 	err = orch.StopSingleVM(ctx, vmID)
 	require.NoError(t, err, "Failed to stop VM")
 
-	_, _, err = orch.LoadSnapshot(ctx, "", vmID, snap)
+	_, _, err = orch.LoadSnapshot(ctx, "1", vmID, snap)
 	require.NoError(t, err, "Failed to load snapshot of VM")
 
 	_, err = orch.ResumeVM(ctx, vmID)

ctriface/iface.go

@@ -220,7 +220,7 @@ func (o *Orchestrator) StartVMWithEnvironment(ctx context.Context, vmID, imageNa
 		}
 		logger.Debugf("TEST: show to-reg snapStat: %+v", stateCfg)
 		logger.Debugf("TEST: show socket path: %s", resp.GetSocketPath())
-		if err := o.memoryManager.RegisterVM(stateCfg); err != nil {
+		if err := o.memoryManager.RegisterVM(stateCfg, false, ""); err != nil {
 			return nil, nil, errors.Wrap(err, "failed to register VM with memory manager")
 			// NOTE (Plamen): Potentially need a defer(DeregisteVM) here if RegisterVM is not last to execute
 		}
@@ -451,7 +451,7 @@ func (o *Orchestrator) CreateSnapshot(ctx context.Context, vmID string, snap *sn
 }
 
 // LoadSnapshot Loads a snapshot of a VM
-func (o *Orchestrator) LoadSnapshot(ctx context.Context, snapVmID string, vmID string, snap *snapshotting.Snapshot) (_ *StartVMResponse, _ *metrics.Metric, retErr error) {
+func (o *Orchestrator) LoadSnapshot(ctx context.Context, lastVmID string, vmID string, snap *snapshotting.Snapshot) (_ *StartVMResponse, _ *metrics.Metric, retErr error) {
 	var (
 		loadSnapshotMetric   *metrics.Metric = metrics.NewMetric()
 		tStart               time.Time
@@ -507,13 +507,13 @@ func (o *Orchestrator) LoadSnapshot(ctx context.Context, snapVmID string, vmID s
 	if o.GetUPFEnabled() {
 		logger.Debug("TEST: UPF is enabled")
 		conf.MemBackend.BackendType = uffdBackend
-		conf.MemBackend.BackendPath, err = o.memoryManager.GetUPFSockPath(snapVmID)
+		conf.MemBackend.BackendPath, err = o.memoryManager.GetUPFSockPath(lastVmID, true)
 		logger.Debugf("TEST: the upf socket: %s", conf.MemBackend.BackendPath)
 		if err != nil {
 			return nil, nil, errors.Wrapf(err, "failed to get UPF socket path for uffd backend")
 		}
 
-		if err := o.memoryManager.FetchState(snapVmID); err != nil {
+		if err := o.memoryManager.FetchState(lastVmID); err != nil {
 			return nil, nil, err
 		}
 	}
@@ -569,12 +569,11 @@ func (o *Orchestrator) LoadSnapshot(ctx context.Context, snapVmID string, vmID s
 			WorkingSetPath:   o.getWorkingSetFile(vmID),
 			InstanceSockAddr: newUPFSockPath,
 		}
-		if err := o.memoryManager.RegisterVM(stateCfg); err != nil {
+		if err := o.memoryManager.RegisterVM(stateCfg, true, vmID); err != nil {
 			logger.Error(err, "failed to register new VM with memory manager")
 		}
 
 
-
 		if activateErr = o.memoryManager.Activate(vmID); activateErr != nil {
 			logger.Warn("Failed to activate VM in the memory manager", activateErr)
 		}

ctriface/manual_cleanup_test.go

@@ -91,7 +91,7 @@ func TestSnapLoad(t *testing.T) {
 
 	vmID = "2"
 
-	_, _, err = orch.LoadSnapshot(ctx, vmID, vmID, snap)
+	_, _, err = orch.LoadSnapshot(ctx, "1", vmID, snap)
 	require.NoError(t, err, "Failed to load snapshot of VM")
 
 	_, err = orch.ResumeVM(ctx, vmID)
@@ -152,7 +152,7 @@ func TestSnapLoadMultiple(t *testing.T) {
 
 	vmID = "4"
 
-	_, _, err = orch.LoadSnapshot(ctx, vmID, vmID, snap)
+	_, _, err = orch.LoadSnapshot(ctx, "3", vmID, snap)
 	require.NoError(t, err, "Failed to load snapshot of VM")
 
 	_, err = orch.ResumeVM(ctx, vmID)
@@ -163,7 +163,7 @@ func TestSnapLoadMultiple(t *testing.T) {
 
 	vmID = "5"
 
-	_, _, err = orch.LoadSnapshot(ctx, vmID, vmID, snap)
+	_, _, err = orch.LoadSnapshot(ctx, "4", vmID, snap)
 	require.NoError(t, err, "Failed to load snapshot of VM")
 
 	_, err = orch.ResumeVM(ctx, vmID)
@@ -234,10 +234,11 @@ func TestParallelSnapLoad(t *testing.T) {
 			err = orch.StopSingleVM(ctx, vmID)
 			require.NoError(t, err, "Failed to offload VM, "+vmID)
 
+			lastVmID := vmID
 			vmIDInt, _ := strconv.Atoi(vmID)
 			vmID = strconv.Itoa(vmIDInt + 1)
 
-			_, _, err = orch.LoadSnapshot(ctx, vmID, vmID, snap)
+			_, _, err = orch.LoadSnapshot(ctx, lastVmID, vmID, snap)
 			require.NoError(t, err, "Failed to load snapshot of VM, "+vmID)
 
 			_, err = orch.ResumeVM(ctx, vmID)
@@ -357,9 +358,10 @@ func TestParallelPhasedSnapLoad(t *testing.T) {
 				defer vmGroup.Done()
 				vmID := fmt.Sprintf("%d", i+vmIDBase)
 				snap := snapshotting.NewSnapshot(vmID, "/fccd/snapshots", testImageName)
+				lastVmID := vmID
 				vmIDInt, _ := strconv.Atoi(vmID)
 				vmID = strconv.Itoa(vmIDInt + 1)
-				_, _, err := orch.LoadSnapshot(ctx, vmID, vmID, snap)
+				_, _, err := orch.LoadSnapshot(ctx, lastVmID, vmID, snap)
 				require.NoError(t, err, "Failed to load snapshot of VM, "+vmID)
 			}(i)
 		}
@@ -482,7 +484,7 @@ func TestRemoteSnapLoad(t *testing.T) {
 
 	snap := snapshotting.NewSnapshot(revision, remoteSnapshotsDir, testImageName)
 
-	_, _, err = orch.LoadSnapshot(ctx, vmID, vmID, snap)
+	_, _, err = orch.LoadSnapshot(ctx, "36", vmID, snap)
 	require.NoError(t, err, "Failed to load remote snapshot of VM")
 
 	_, err = orch.ResumeVM(ctx, vmID)

functions.go

@@ -361,11 +361,12 @@ func (f *Function) AddInstance() *metrics.Metric {
 
 	if f.isSnapshotReady {
 		var resp *ctriface.StartVMResponse
-
 		// resp, metr = f.LoadInstance(f.getVMID())
-		snapVmID := fmt.Sprintf("%s-%d", f.fID, f.lastInstanceID - 1)
+
+		lastVmID := fmt.Sprintf("%s-%d", f.fID, f.lastInstanceID - 1)
 		currVmID := f.getVMID()
-		resp, metr = f.LoadInstance(snapVmID, currVmID)	
+		resp, metr = f.LoadInstance(lastVmID, currVmID)	
+
 		f.guestIP = resp.GuestIP
 		f.vmID = f.getVMID()
 		f.lastInstanceID++
@@ -482,7 +483,7 @@ func (f *Function) CreateInstanceSnapshot() {
 
 // LoadInstance Loads a new instance of the function from its snapshot and resumes it
 // The tap, the shim and the vmID remain the same
-func (f *Function) LoadInstance(snapVmID string, vmID string) (*ctriface.StartVMResponse, *metrics.Metric) {
+func (f *Function) LoadInstance(lastVmID string, vmID string) (*ctriface.StartVMResponse, *metrics.Metric) {
 	logger := log.WithFields(log.Fields{"fID": f.fID})
 
 	logger.Debug("Loading instance")
@@ -495,7 +496,7 @@ func (f *Function) LoadInstance(snapVmID string, vmID string) (*ctriface.StartVM
 		log.Panic(err)
 	}
 
-	resp, loadMetr, err := orch.LoadSnapshot(ctx, snapVmID, vmID, snap)
+	resp, loadMetr, err := orch.LoadSnapshot(ctx, lastVmID, vmID, snap)
 	if err != nil {
 		log.Panic(err)
 	}

memory/manager/manager.go

@@ -53,6 +53,7 @@ type MemoryManager struct {
 	sync.Mutex
 	MemoryManagerCfg
 	instances map[string]*SnapshotState // Indexed by vmID
+	origins map[string]string
 }
 
 // NewMemoryManager Initializes a new memory manager
@@ -62,12 +63,13 @@ func NewMemoryManager(cfg MemoryManagerCfg) *MemoryManager {
 	m := new(MemoryManager)
 	m.instances = make(map[string]*SnapshotState)
 	m.MemoryManagerCfg = cfg
+	m.origins = make(map[string]string)
 
 	return m
 }
 
 // RegisterVM Registers a VM within the memory manager
-func (m *MemoryManager) RegisterVM(cfg SnapshotStateCfg) error {
+func (m *MemoryManager) RegisterVM(cfg SnapshotStateCfg, isSnapshotReady bool, originID string) error {
 	m.Lock()
 	defer m.Unlock()
 
@@ -86,6 +88,10 @@ func (m *MemoryManager) RegisterVM(cfg SnapshotStateCfg) error {
 	state := NewSnapshotState(cfg)
 
 	m.instances[vmID] = state
+	// if isSnapshotReady {
+	// 	logger.Debugf("TEST: register current vmID %s with originID %s", vmID, originID)
+	// 	m.origins[vmID] = originID
+	// }
 
 	return nil
 }
@@ -130,7 +136,17 @@ func (m *MemoryManager) Activate(vmID string) error {
 	m.Lock()
 
 	logger.Debug("TEST: Activate: fetch snapstate by vmID for UFFD")
+
+	// originID, ok := m.origins[vmID]
+
+	// if !ok {
+	// 	logger.Debug("TEST: not loaded from snapshot")
+	// }
+
+	// state, ok = m.instances[originID]
+
 	state, ok = m.instances[vmID]
+
 	if !ok {
 		m.Unlock()
 		logger.Error("VM not registered with the memory manager")
@@ -178,6 +194,12 @@ func (m *MemoryManager) FetchState(vmID string) error {
 
 	m.Lock()
 
+	// originID, ok := m.origins[vmID]
+	// if !ok {
+	// 	logger.Debug("TEST: not loaded from snapshot")
+	// }
+	// state, ok = m.instances[originID]
+
 	state, ok = m.instances[vmID]
 	if !ok {
 		m.Unlock()
@@ -352,13 +374,24 @@ func (m *MemoryManager) GetUPFLatencyStats(vmID string) ([]*metrics.Metric, erro
 	return state.latencyMetrics, nil
 }
 
-func (m *MemoryManager) GetUPFSockPath(vmID string) (string, error) {
+func (m *MemoryManager) GetUPFSockPath(vmID string, isSnapshotReady bool) (string, error) {
 	logger := log.WithFields(log.Fields{"vmID": vmID})
 
 	logger.Debug("Get the path of firecracker unix domain socket")
 
 	m.Lock()
 
+	// id := ""
+	// if isSnapshotReady {
+	// 	logger.Debugf("TEST: to find originID by vmID %s", vmID)
+	// 	originID, ok := m.origins[vmID]
+	// 	if !ok {
+	// 		logger.Debug("TEST: not loaded from snapshot")
+	// 	}
+	// 	id = originID
+	// }
+	// state, ok := m.instances[id]
+
 	state, ok := m.instances[vmID]
 	if !ok {
 		m.Unlock()