feat(agent): lazy reattach on restart — restore procs + IP allocation + VTEP

Author: syscod3

cmd/agent/main.go

@@ -97,6 +97,11 @@ func main() {
 		log.Info("Using FirecrackerDriver")
 	}
 
+	var alloc *network.Allocator
+	if fcDrv, ok := driver.(*agent.FirecrackerDriver); ok {
+		alloc = fcDrv.Alloc
+	}
+
 	if err := (&agent.ImpVMReconciler{
 		Client:   mgr.GetClient(),
 		Scheme:   mgr.GetScheme(),
@@ -105,6 +110,7 @@ func main() {
 		Driver:   driver,
 		Metrics:  mc,
 		Net:      prodNet,
+		Alloc:    alloc,
 	}).SetupWithManager(mgr); err != nil {
 		log.Error(err, "Unable to set up ImpVMReconciler")
 		os.Exit(1)

internal/agent/firecracker_driver.go

@@ -642,12 +642,32 @@ func (d *FirecrackerDriver) applySnapshotBoot(ctx context.Context, vm *impdevv1a
 	return nil
 }
 
-// IsAlive is a placeholder — real implementation in Task 2.
-func (d *FirecrackerDriver) IsAlive(_ int64) bool {
-	return false
+// IsAlive reports whether the process with the given PID is still running.
+// Uses kill(pid, 0) — succeeds if the process exists (even if zombie).
+func (d *FirecrackerDriver) IsAlive(pid int64) bool {
+	p, err := os.FindProcess(int(pid))
+	if err != nil {
+		return false
+	}
+	return p.Signal(syscall.Signal(0)) == nil
 }
 
-// Reattach is a placeholder — real implementation in Task 2.
-func (d *FirecrackerDriver) Reattach(_ context.Context, _ *impdevv1alpha1.ImpVM) error {
-	return fmt.Errorf("not implemented")
+// Reattach re-registers an already-running Firecracker VM into the driver's
+// procs map without launching a new process. Called after an agent restart when
+// the Firecracker process survived but the in-memory procs map was lost.
+// Returns an error if the VM's API socket is not present on disk (which would
+// mean the PID has been reused by a different process).
+func (d *FirecrackerDriver) Reattach(_ context.Context, vm *impdevv1alpha1.ImpVM) error {
+	sock := d.socketPath(vm)
+	if _, err := os.Stat(sock); err != nil {
+		return fmt.Errorf("reattach %s: socket %s not found — PID may be reused: %w",
+			vmKey(vm), sock, err)
+	}
+	d.mu.Lock()
+	d.procs[vmKey(vm)] = &fcProc{
+		pid:    vm.Status.RuntimePID,
+		socket: sock,
+	}
+	d.mu.Unlock()
+	return nil
 }

internal/agent/reconciler.go

@@ -31,6 +31,9 @@ type ImpVMReconciler struct {
 	Metrics *VMMetricsCollector
 	// Net is optional. When non-nil, used for VXLAN/FDB operations after VTEP sync.
 	Net network.NetManager
+	// Alloc is the in-memory IP allocator. When non-nil, Reserve is called during
+	// lazy reattach to restore IP state after agent restart.
+	Alloc *network.Allocator
 }
 
 // +kubebuilder:rbac:groups=imp.dev,resources=impvms,verbs=get;list;watch;update;patch
@@ -139,6 +142,32 @@ func (r *ImpVMReconciler) handleRunning(ctx context.Context, vm *impdevv1alpha1.
 		return ctrl.Result{}, nil // watch-driven steady state
 	}
 
+	// Inspect returned Running=false. Before declaring the VM dead, check whether
+	// the Firecracker process is still alive (procs map may be empty after an
+	// agent pod restart). If PID is alive, reattach and restore allocator state.
+	if pid := vm.Status.RuntimePID; pid > 0 && r.Driver.IsAlive(pid) {
+		if err := r.Driver.Reattach(ctx, vm); err != nil {
+			log.Error(err, "Reattach failed — treating VM as dead")
+		} else {
+			// Restore in-memory IP allocation so Release works correctly later.
+			if r.Alloc != nil && vm.Spec.NetworkRef != nil && vm.Status.IP != "" {
+				netKey := vm.Namespace + "/" + vm.Spec.NetworkRef.Name
+				r.Alloc.Reserve(netKey, vm.Status.IP)
+			}
+			// Re-publish VTEP entry and sync FDB in case they were lost.
+			if vm.Spec.NetworkRef != nil && vm.Status.IP != "" && r.NodeIP != "" {
+				macAddr := network.MACAddr(vm.Namespace + "/" + vm.Name)
+				if err := r.registerVTEP(ctx, vm, vm.Status.IP, macAddr); err != nil {
+					log.Error(err, "registerVTEP after reattach failed")
+				} else if err := r.syncFDB(ctx, vm); err != nil {
+					log.Error(err, "syncFDB after reattach failed")
+				}
+			}
+			log.Info("VM reattached after agent restart", "pid", pid)
+			return ctrl.Result{}, nil
+		}
+	}
+
 	log.Info("VM process exited", "lifecycle", vm.Spec.Lifecycle)
 	if vm.Spec.Lifecycle == impdevv1alpha1.VMLifecycleEphemeral {
 		return r.finishSucceeded(ctx, vm)

internal/agent/reconciler_test.go

@@ -401,6 +401,104 @@ var _ = Describe("ImpVM Agent: Reconcile non-existent VM (NotFound)", func() {
 	})
 })
 
+var _ = Describe("ImpVM Agent: Running — lazy reattach on agent restart", func() {
+	ctx := context.Background()
+
+	It("reattaches and stays Running when PID is alive", func() {
+		driver := NewStubDriver()
+		driver.IsAliveResult = true
+
+		vm := &impdevv1alpha1.ImpVM{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "tc-reattach-alive", Namespace: "default",
+				Finalizers: []string{"imp/finalizer"},
+			},
+			Spec: impdevv1alpha1.ImpVMSpec{NodeName: testNode},
+		}
+		Expect(k8sClient.Create(ctx, vm)).To(Succeed())
+		DeferCleanup(func() { k8sClient.Delete(ctx, vm) }) //nolint:errcheck
+
+		base := vm.DeepCopy()
+		vm.Status.Phase = impdevv1alpha1.VMPhaseRunning
+		vm.Status.RuntimePID = 99999
+		Expect(k8sClient.Status().Patch(ctx, vm, client.MergeFrom(base))).To(Succeed())
+
+		// Inspect returns Running=false (no entry in states map — simulates restart)
+		// IsAliveResult=true means the PID check passes
+
+		_, err := newReconciler(driver).Reconcile(ctx, reconcile.Request{
+			NamespacedName: types.NamespacedName{Name: "tc-reattach-alive", Namespace: "default"},
+		})
+		Expect(err).NotTo(HaveOccurred())
+
+		updated := &impdevv1alpha1.ImpVM{}
+		Expect(k8sClient.Get(ctx, types.NamespacedName{Name: "tc-reattach-alive", Namespace: "default"}, updated)).To(Succeed())
+		Expect(updated.Status.Phase).To(Equal(impdevv1alpha1.VMPhaseRunning))
+		Expect(driver.ReattachCalls).To(HaveLen(1))
+		Expect(driver.ReattachCalls[0]).To(Equal("default/tc-reattach-alive"))
+	})
+
+	It("transitions to Failed when PID is dead", func() {
+		driver := NewStubDriver()
+		driver.IsAliveResult = false
+
+		vm := &impdevv1alpha1.ImpVM{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "tc-reattach-dead", Namespace: "default",
+				Finalizers: []string{"imp/finalizer"},
+			},
+			Spec: impdevv1alpha1.ImpVMSpec{NodeName: testNode},
+		}
+		Expect(k8sClient.Create(ctx, vm)).To(Succeed())
+		DeferCleanup(func() { k8sClient.Delete(ctx, vm) }) //nolint:errcheck
+
+		base := vm.DeepCopy()
+		vm.Status.Phase = impdevv1alpha1.VMPhaseRunning
+		vm.Status.RuntimePID = 99999
+		Expect(k8sClient.Status().Patch(ctx, vm, client.MergeFrom(base))).To(Succeed())
+
+		_, err := newReconciler(driver).Reconcile(ctx, reconcile.Request{
+			NamespacedName: types.NamespacedName{Name: "tc-reattach-dead", Namespace: "default"},
+		})
+		Expect(err).NotTo(HaveOccurred())
+
+		updated := &impdevv1alpha1.ImpVM{}
+		Expect(k8sClient.Get(ctx, types.NamespacedName{Name: "tc-reattach-dead", Namespace: "default"}, updated)).To(Succeed())
+		Expect(updated.Status.Phase).To(Equal(impdevv1alpha1.VMPhaseFailed))
+		Expect(driver.ReattachCalls).To(BeEmpty())
+	})
+
+	It("transitions to Failed when RuntimePID is zero", func() {
+		driver := NewStubDriver()
+		driver.IsAliveResult = true // shouldn't matter — PID is 0
+
+		vm := &impdevv1alpha1.ImpVM{
+			ObjectMeta: metav1.ObjectMeta{
+				Name: "tc-reattach-nopid", Namespace: "default",
+				Finalizers: []string{"imp/finalizer"},
+			},
+			Spec: impdevv1alpha1.ImpVMSpec{NodeName: testNode},
+		}
+		Expect(k8sClient.Create(ctx, vm)).To(Succeed())
+		DeferCleanup(func() { k8sClient.Delete(ctx, vm) }) //nolint:errcheck
+
+		base := vm.DeepCopy()
+		vm.Status.Phase = impdevv1alpha1.VMPhaseRunning
+		vm.Status.RuntimePID = 0
+		Expect(k8sClient.Status().Patch(ctx, vm, client.MergeFrom(base))).To(Succeed())
+
+		_, err := newReconciler(driver).Reconcile(ctx, reconcile.Request{
+			NamespacedName: types.NamespacedName{Name: "tc-reattach-nopid", Namespace: "default"},
+		})
+		Expect(err).NotTo(HaveOccurred())
+
+		updated := &impdevv1alpha1.ImpVM{}
+		Expect(k8sClient.Get(ctx, types.NamespacedName{Name: "tc-reattach-nopid", Namespace: "default"}, updated)).To(Succeed())
+		Expect(updated.Status.Phase).To(Equal(impdevv1alpha1.VMPhaseFailed))
+		Expect(driver.ReattachCalls).To(BeEmpty())
+	})
+})
+
 var _ = Describe("ImpVM Agent: handleTerminating driver.Stop error", func() {
 	ctx := context.Background()