feat(agent): scale-to-zero wake-on-traffic datapath (Phase 3)

cmd/agent/main.go

@@ -6,6 +6,7 @@ import (
 	"context"
 	"flag"
 	"os"
+	"time"
 
 	_ "k8s.io/client-go/plugin/pkg/client/auth"
 
@@ -111,6 +112,14 @@ func main() {
 		alloc = fcDrv.Alloc
 	}
 
+	// Scale-to-zero (wake-on-traffic) is experimental and its capture hook is not
+	// yet hardware-validated — opt in explicitly via IMP_SCALE_TO_ZERO=true.
+	var sz *agent.ScaleToZero
+	if os.Getenv("IMP_SCALE_TO_ZERO") == "true" {
+		sz = agent.NewLinuxScaleToZero(1024, 15*time.Second)
+		log.Info("scale-to-zero enabled (experimental; wake path not hardware-validated)")
+	}
+
 	if err := (&agent.ImpVMReconciler{
 		Client:   mgr.GetClient(),
 		Scheme:   mgr.GetScheme(),
@@ -120,6 +129,7 @@ func main() {
 		Metrics:  mc,
 		Net:      prodNet,
 		Alloc:    alloc,
+		SZ:       sz,
 		Recorder: mgr.GetEventRecorderFor("imp-agent"), //nolint:staticcheck // controller-runtime returns legacy recorder type expected by reconciler
 	}).SetupWithManager(mgr); err != nil {
 		log.Error(err, "Unable to set up ImpVMReconciler")

internal/agent/reconciler.go

@@ -22,7 +22,9 @@ import (
 	"k8s.io/client-go/util/retry"
 	ctrl "sigs.k8s.io/controller-runtime"
 	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/handler"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
+	"sigs.k8s.io/controller-runtime/pkg/source"
 
 	impdevv1alpha1 "github.com/syscode-labs/imp/api/v1alpha1"
 	"github.com/syscode-labs/imp/internal/agent/network"
@@ -56,6 +58,9 @@ type ImpVMReconciler struct {
 	// SuspendDir is the node-local base directory under which suspend snapshots
 	// are written (one subdir per VM). Defaults to /var/lib/imp/suspend when empty.
 	SuspendDir string
+	// SZ is optional. When non-nil it enables scale-to-zero: ScaleToZero VMs
+	// auto-suspend on traffic idle and auto-resume on the first inbound packet.
+	SZ *ScaleToZero
 }
 
 // suspendBaseDir returns the configured suspend snapshot base directory,
@@ -256,14 +261,27 @@ func (r *ImpVMReconciler) handleRunning(ctx context.Context, vm *impdevv1alpha1.
 		return ctrl.Result{}, nil
 	}
 
+	// ScaleToZero: poll TAP traffic and auto-suspend after idleTimeout of silence.
+	if vm.Spec.DesiredState == impdevv1alpha1.VMDesiredStateScaleToZero && r.SZ != nil {
+		suspended, err := r.maybeSuspendIdle(ctx, vm)
+		if err != nil {
+			log.Error(err, "ScaleToZero idle probe failed; will retry")
+			return ctrl.Result{RequeueAfter: r.SZ.interval}, nil
+		}
+		if suspended {
+			return ctrl.Result{}, nil
+		}
+		// Not idle — fall through to the liveness check; runningResult requeues.
+	}
+
 	state, err := r.Driver.Inspect(ctx, vm)
 	if err != nil {
 		log.Error(err, "Driver Inspect failed")
 		return ctrl.Result{}, err
 	}
 
 	if state.Running {
-		return ctrl.Result{}, nil // watch-driven steady state
+		return r.runningResult(vm), nil // watch-driven steady state (ScaleToZero: polled)
 	}
 
 	// Inspect returned Running=false. Before declaring the VM dead, check whether
@@ -333,6 +351,46 @@ func (r *ImpVMReconciler) handleRunning(ctx context.Context, vm *impdevv1alpha1.
 	return r.finishFailed(ctx, vm)
 }
 
+// runningResult is the steady-state reconcile result for a Running VM: empty
+// (watch-driven) normally, or a poll interval for ScaleToZero VMs whose idle
+// detector needs periodic wake-ups.
+func (r *ImpVMReconciler) runningResult(vm *impdevv1alpha1.ImpVM) ctrl.Result {
+	if r.SZ != nil && vm.Spec.DesiredState == impdevv1alpha1.VMDesiredStateScaleToZero {
+		return ctrl.Result{RequeueAfter: r.SZ.interval}
+	}
+	return ctrl.Result{}
+}
+
+// maybeSuspendIdle probes vm's TAP traffic. If the VM has been idle for its
+// idleTimeout it stamps lastActivityTime and transitions to Suspending, returning
+// true; otherwise it returns false without touching status. We deliberately do NOT
+// write lastActivityTime on every busy probe — a genuinely active VM changes bytes
+// each interval, so that would be a status write every interval per VM (churn at the
+// density this feature targets). The field records when the VM last had traffic
+// before suspending; live activity is observable via /metrics.
+func (r *ImpVMReconciler) maybeSuspendIdle(ctx context.Context, vm *impdevv1alpha1.ImpVM) (bool, error) {
+	key := client.ObjectKeyFromObject(vm)
+	tap := network.TAPName(vmKey(vm))
+	idle, lastActivity, err := r.SZ.observe(key, tap, idleTimeoutOrDefault(vm), time.Now())
+	if err != nil {
+		return false, err
+	}
+	if !idle {
+		return false, nil
+	}
+	base := vm.DeepCopy()
+	t := metav1.NewTime(lastActivity)
+	vm.Status.LastActivityTime = &t
+	vm.Status.Phase = impdevv1alpha1.VMPhaseSuspending
+	if err := r.Status().Patch(ctx, vm, client.MergeFrom(base)); err != nil {
+		if apierrors.IsConflict(err) {
+			return false, nil // requeue re-evaluates
+		}
+		return false, err
+	}
+	return true, nil
+}
+
 func (r *ImpVMReconciler) handleTerminating(ctx context.Context, vm *impdevv1alpha1.ImpVM) (result ctrl.Result, err error) {
 	log := logf.FromContext(ctx)
 
@@ -401,8 +459,11 @@ func (r *ImpVMReconciler) handleSuspending(ctx context.Context, vm *impdevv1alph
 		return ctrl.Result{RequeueAfter: r.retryInterval()}, err
 	}
 
-	// Deregister VTEP so other nodes stop routing to the VM about to be stopped.
-	if vm.Spec.NetworkRef != nil && vm.Status.IP != "" {
+	// Deregister VTEP so other nodes stop routing to the VM about to be stopped —
+	// EXCEPT for ScaleToZero, which keeps the VTEP so the overlay still delivers
+	// the wake packet to this node while the VM is suspended.
+	if vm.Spec.NetworkRef != nil && vm.Status.IP != "" &&
+		vm.Spec.DesiredState != impdevv1alpha1.VMDesiredStateScaleToZero {
 		if derr := r.deregisterVTEP(ctx, vm); derr != nil {
 			log.Error(derr, "Failed to deregister VTEP during suspend")
 		}
@@ -431,16 +492,32 @@ func (r *ImpVMReconciler) handleSuspending(ctx context.Context, vm *impdevv1alph
 	if r.Metrics != nil {
 		r.Metrics.SetVMState(vm.Namespace+"/"+vm.Name, "Suspended", r.NodeName)
 	}
+
+	// ScaleToZero: register the VM's IP so the activator wakes it on the first
+	// inbound packet. Reset any idle sample so a fresh window applies on resume.
+	if vm.Spec.DesiredState == impdevv1alpha1.VMDesiredStateScaleToZero && r.SZ != nil && vm.Status.IP != "" {
+		r.SZ.reg.register(vm.Status.IP, vm)
+		r.SZ.resetIdle(client.ObjectKeyFromObject(vm))
+	}
+
 	log.Info("VM suspended", "snapshotPath", destDir)
 	return ctrl.Result{}, nil
 }
 
 // handleSuspended holds the VM in the Suspended state until desiredState flips
 // back to Running, at which point it advances to Resuming.
 func (r *ImpVMReconciler) handleSuspended(ctx context.Context, vm *impdevv1alpha1.ImpVM) (ctrl.Result, error) {
-	if vm.Spec.DesiredState == impdevv1alpha1.VMDesiredStateSuspended {
+	switch vm.Spec.DesiredState {
+	case impdevv1alpha1.VMDesiredStateSuspended:
 		return ctrl.Result{}, nil // steady state
+	case impdevv1alpha1.VMDesiredStateScaleToZero:
+		// Stay suspended until the activator observes an inbound packet.
+		if r.SZ == nil || !r.SZ.reg.pending(client.ObjectKeyFromObject(vm)) {
+			return ctrl.Result{}, nil
+		}
+		// Wake packet observed — fall through to resume.
 	}
+	// desiredState=Running (explicit resume) or ScaleToZero with a pending wake.
 	base := vm.DeepCopy()
 	vm.Status.Phase = impdevv1alpha1.VMPhaseResuming
 	if err := r.Status().Patch(ctx, vm, client.MergeFrom(base)); err != nil {
@@ -501,6 +578,13 @@ func (r *ImpVMReconciler) handleResuming(ctx context.Context, vm *impdevv1alpha1
 		r.Metrics.SetVMState(vm.Namespace+"/"+vm.Name, "Running", r.NodeName)
 	}
 
+	// Clear scale-to-zero wake state now that the VM is running again.
+	if r.SZ != nil {
+		key := client.ObjectKeyFromObject(vm)
+		r.SZ.reg.clear(key)
+		r.SZ.resetIdle(key)
+	}
+
 	// Re-register VTEP + sync FDB so other nodes route to the resumed VM.
 	r.ensureVTEPAndFDB(ctx, vm, state.IP)
 
@@ -597,10 +681,22 @@ func (r *ImpVMReconciler) SetupWithManager(mgr ctrl.Manager) error {
 	// Detect and patch CPU model onto ClusterImpNodeProfile at startup (best-effort).
 	go detectAndPatchCPUModel(context.Background(), r.Client, r.NodeName)
 
-	return ctrl.NewControllerManagedBy(mgr).
+	b := ctrl.NewControllerManagedBy(mgr).
 		For(&impdevv1alpha1.ImpVM{}).
-		Named("agent-impvm").
-		Complete(r)
+		Named("agent-impvm")
+
+	// Scale-to-zero: let the activator fire reconciles via a channel, and run the
+	// packet source as a per-node runnable.
+	if r.SZ != nil {
+		b = b.WatchesRawSource(source.Channel(r.SZ.reg.events, &handler.EnqueueRequestForObject{}))
+		if r.SZ.src != nil {
+			if err := mgr.Add(&activator{src: r.SZ.src, reg: r.SZ.reg}); err != nil {
+				return err
+			}
+		}
+	}
+
+	return b.Complete(r)
 }
 
 // registerVTEP adds or updates the VTEPEntry for vm in ImpNetwork.status.vtepTable.

internal/agent/scaletozero.go

@@ -0,0 +1,180 @@
+package agent
+
+// Scale-to-zero wake-on-traffic support (Phase 3).
+//
+// This file holds the platform-neutral, fully unit-tested core: the wake
+// registry (which suspended VMs are awaiting a packet), the traffic-idle
+// detector, and the reconcile-triggering plumbing. The two pieces that need a
+// real host — reading NIC byte counters and capturing packets — are injected as
+// a linkStatsFunc and a PacketSource, faked in tests and implemented for real in
+// scaletozero_linux.go.
+//
+// UNVALIDATED: the real PacketSource (AF_PACKET on the overlay) has NOT been
+// confirmed to observe the first frame destined to a TAP-less (suspended) VM.
+// The first cluster spike (Phase 1 of the wake-on-traffic plan) must validate
+// hook placement. The PacketSource interface exists precisely so the hook can be
+// swapped to tc-BPF later without touching any of the logic below.
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"k8s.io/apimachinery/pkg/types"
+	"sigs.k8s.io/controller-runtime/pkg/client"
+	"sigs.k8s.io/controller-runtime/pkg/event"
+
+	impdevv1alpha1 "github.com/syscode-labs/imp/api/v1alpha1"
+)
+
+// defaultIdleTimeout is used when a ScaleToZero VM leaves spec.idleTimeout unset.
+const defaultIdleTimeout = 5 * time.Minute
+
+// PacketSource delivers the destination IP of each inbound overlay frame to
+// onDstIP until ctx is cancelled. Implementations are host-specific.
+type PacketSource interface {
+	Run(ctx context.Context, onDstIP func(ip string)) error
+}
+
+// linkStatsFunc returns cumulative rx+tx bytes for a host interface (a VM TAP).
+type linkStatsFunc func(iface string) (bytes uint64, err error)
+
+// wakeRegistry tracks suspended ScaleToZero VMs by IP and fires a reconcile when
+// a packet arrives for one. Safe for concurrent use: the activator goroutine
+// calls onDstIP while the reconcile loop calls register/clear/pending.
+type wakeRegistry struct {
+	mu        sync.Mutex
+	keyByIP   map[string]types.NamespacedName
+	ipByKey   map[types.NamespacedName]string
+	objByKey  map[types.NamespacedName]client.Object
+	signalled map[types.NamespacedName]bool
+	events    chan event.GenericEvent
+}
+
+func newWakeRegistry(bufSize int) *wakeRegistry {
+	if bufSize <= 0 {
+		bufSize = 1024
+	}
+	return &wakeRegistry{
+		keyByIP:   map[string]types.NamespacedName{},
+		ipByKey:   map[types.NamespacedName]string{},
+		objByKey:  map[types.NamespacedName]client.Object{},
+		signalled: map[types.NamespacedName]bool{},
+		events:    make(chan event.GenericEvent, bufSize),
+	}
+}
+
+// register marks vm as suspended and awaiting a wake packet at ip.
+func (w *wakeRegistry) register(ip string, vm client.Object) {
+	key := client.ObjectKeyFromObject(vm)
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	// Drop any stale IP mapping for this VM before recording the new one.
+	if old, ok := w.ipByKey[key]; ok {
+		delete(w.keyByIP, old)
+	}
+	w.keyByIP[ip] = key
+	w.ipByKey[key] = ip
+	w.objByKey[key] = vm
+}
+
+// onDstIP is the PacketSource callback: a frame arrived for ip. If ip belongs to
+// a registered VM not already signalled, enqueue a reconcile for it. The
+// signalled flag is set only when the event is actually enqueued, so a full
+// channel never silently loses a wake — the next packet retries.
+func (w *wakeRegistry) onDstIP(ip string) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	key, ok := w.keyByIP[ip]
+	if !ok || w.signalled[key] {
+		return
+	}
+	obj := w.objByKey[key]
+	select {
+	case w.events <- event.GenericEvent{Object: obj}:
+		w.signalled[key] = true
+	default:
+		// Channel full; leave unsignalled so a later packet retries.
+	}
+}
+
+// pending reports whether a wake packet has been observed for key.
+func (w *wakeRegistry) pending(key types.NamespacedName) bool {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	return w.signalled[key]
+}
+
+// clear drops all state for key (called once the VM has resumed).
+func (w *wakeRegistry) clear(key types.NamespacedName) {
+	w.mu.Lock()
+	defer w.mu.Unlock()
+	if ip, ok := w.ipByKey[key]; ok {
+		delete(w.keyByIP, ip)
+	}
+	delete(w.ipByKey, key)
+	delete(w.objByKey, key)
+	delete(w.signalled, key)
+}
+
+// ScaleToZero bundles the wake registry, the idle detector, and the packet
+// source into the optional feature attached to the reconciler as SZ.
+type ScaleToZero struct {
+	reg      *wakeRegistry
+	stats    linkStatsFunc
+	src      PacketSource
+	interval time.Duration
+
+	mu      sync.Mutex
+	samples map[types.NamespacedName]idleSample
+}
+
+type idleSample struct {
+	bytes uint64
+	since time.Time
+}
+
+func newScaleToZero(stats linkStatsFunc, src PacketSource, interval time.Duration, bufSize int) *ScaleToZero {
+	if interval <= 0 {
+		interval = 15 * time.Second
+	}
+	return &ScaleToZero{
+		reg:      newWakeRegistry(bufSize),
+		stats:    stats,
+		src:      src,
+		interval: interval,
+		samples:  map[types.NamespacedName]idleSample{},
+	}
+}
+
+// observe samples iface's byte counter and reports whether the VM has seen no
+// traffic for at least idleTimeout. Any change resets the idle clock, so a
+// freshly-resumed VM (no prior sample) always gets a full idleTimeout of grace —
+// this is the anti-thrash hysteresis.
+//
+// ASSUMPTION (validate on cluster): byte-idle suspends a VM holding an idle-but-
+// open connection (long-poll, pooled DB conn), which the resume must transparently
+// re-establish or the connection breaks. Combining with the guest CPU-idle signal
+// is an open design question; not built here.
+func (s *ScaleToZero) observe(key types.NamespacedName, iface string, idleTimeout time.Duration, now time.Time) (idle bool, lastActivity time.Time, err error) {
+	b, err := s.stats(iface)
+	if err != nil {
+		return false, time.Time{}, err
+	}
+	s.mu.Lock()
+	defer s.mu.Unlock()
+	prev, ok := s.samples[key]
+	if !ok || b != prev.bytes {
+		s.samples[key] = idleSample{bytes: b, since: now}
+		return false, now, nil
+	}
+	return now.Sub(prev.since) >= idleTimeout, prev.since, nil
+}
+
+// idleTimeoutOrDefault resolves the effective idle window for a VM.
+func idleTimeoutOrDefault(vm *impdevv1alpha1.ImpVM) time.Duration {
+	if vm.Spec.IdleTimeout != nil && vm.Spec.IdleTimeout.Duration > 0 {
+		return vm.Spec.IdleTimeout.Duration
+	}
+	return defaultIdleTimeout
+}