fix: read ScaleToZero TAP stats via runtime RPC

cmd/agent/driver_linux.go

@@ -15,9 +15,13 @@ import (
 // Reads FC_BIN, FC_SOCK_DIR, FC_KERNEL, FC_KERNEL_ARGS, and IMP_IMAGE_CACHE.
 // Returns the driver, the shared NetManager, and any error.
 func newProductionDriver(_ ctrlclient.Client, _ *agent.VMMetricsCollector, _ string) (agent.VMDriver, network.NetManager, error) {
-	endpoint := os.Getenv("IMP_RUNTIME_SOCKET")
-	if endpoint == "" {
-		endpoint = "/run/imp/runtime.sock"
-	}
+	endpoint := runtimeSocketEndpoint()
 	return agent.NewRuntimeDriver(endpoint), agent.NewRuntimeNetManager(endpoint), nil
 }
+
+func runtimeSocketEndpoint() string {
+	if endpoint := os.Getenv("IMP_RUNTIME_SOCKET"); endpoint != "" {
+		return endpoint
+	}
+	return "/run/imp/runtime.sock"
+}

cmd/agent/main.go

@@ -28,6 +28,7 @@ import (
 	"github.com/syscode-labs/imp/internal/agent/api"
 	"github.com/syscode-labs/imp/internal/agent/network"
 	"github.com/syscode-labs/imp/internal/capability"
+	"github.com/syscode-labs/imp/internal/runtimeapi"
 	"github.com/syscode-labs/imp/internal/telemetry"
 )
 
@@ -141,7 +142,7 @@ func main() {
 	// 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)
+		sz = agent.NewLinuxScaleToZeroWithRuntime(runtimeapi.NewClient(runtimeSocketEndpoint()), 1024, 15*time.Second)
 		log.Info("scale-to-zero enabled (experimental; wake path not hardware-validated)")
 	}
 

internal/agent/scaletozero_linux.go

@@ -11,6 +11,8 @@ import (
 	"golang.org/x/sys/unix"
 	"k8s.io/apimachinery/pkg/types"
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
+
+	"github.com/syscode-labs/imp/internal/runtimeapi"
 )
 
 // resetIdle forgets any idle sample for key (called on suspend/resume so the VM
@@ -131,3 +133,11 @@ func (afpacketSource) Run(ctx context.Context, onDstIP func(string)) error {
 func NewLinuxScaleToZero(bufSize int, interval time.Duration) *ScaleToZero {
 	return newScaleToZero(netlinkLinkStats, afpacketSource{}, interval, bufSize)
 }
+
+// NewLinuxScaleToZeroWithRuntime wires TAP statistics through the node runtime,
+// which owns the TAP network namespace, while retaining the packet capture hook.
+func NewLinuxScaleToZeroWithRuntime(client *runtimeapi.Client, bufSize int, interval time.Duration) *ScaleToZero {
+	return newScaleToZero(func(iface string) (uint64, error) {
+		return client.LinkStats(context.Background(), iface)
+	}, afpacketSource{}, interval, bufSize)
+}

internal/agent/scaletozero_runtime_linux_test.go

@@ -0,0 +1,58 @@
+//go:build linux
+
+package agent
+
+import (
+	"context"
+	"os"
+	"path/filepath"
+	"testing"
+	"time"
+
+	"k8s.io/apimachinery/pkg/types"
+
+	"github.com/syscode-labs/imp/internal/runtimeapi"
+)
+
+func TestScaleToZeroRuntimeLinkStatsUsesRuntimeUnixSocket(t *testing.T) {
+	t.Parallel()
+
+	const wantTAP = "imptap-scale-to-zero"
+	const wantBytes = uint64(1234)
+	server := runtimeapi.NewServer(runtimeapi.BackendFuncs{
+		LinkStatsFunc: func(_ context.Context, tapName string) (uint64, error) {
+			if tapName != wantTAP {
+				t.Errorf("tapName = %q, want %q", tapName, wantTAP)
+			}
+			return wantBytes, nil
+		},
+	})
+	dir, err := os.MkdirTemp("/tmp", "imp-runtime-")
+	if err != nil {
+		t.Fatalf("MkdirTemp() error = %v", err)
+	}
+	t.Cleanup(func() { _ = os.RemoveAll(dir) })
+	endpoint := filepath.Join(dir, "runtime.sock")
+	if err := server.Start(endpoint); err != nil {
+		t.Fatalf("Start() error = %v", err)
+	}
+	t.Cleanup(func() { _ = server.Close() })
+
+	sz := NewLinuxScaleToZeroWithRuntime(runtimeapi.NewClient(endpoint), 1, time.Second)
+	key := types.NamespacedName{Namespace: "default", Name: "worker"}
+	t0 := time.Now()
+	idle, _, err := sz.observe(key, wantTAP, time.Minute, t0)
+	if err != nil {
+		t.Fatalf("first observe() error = %v", err)
+	}
+	if idle {
+		t.Fatal("first observe() reported idle")
+	}
+	idle, _, err = sz.observe(key, wantTAP, time.Minute, t0.Add(2*time.Minute))
+	if err != nil {
+		t.Fatalf("second observe() error = %v", err)
+	}
+	if !idle {
+		t.Fatal("second observe() did not use the unchanged runtime byte count")
+	}
+}

internal/noderuntime/backend.go

@@ -27,6 +27,8 @@ type driver interface {
 type Backend struct {
 	Driver driver
 	Net    network.NetManager
+	// LinkStatsFunc is an optional test seam; production uses runtime-owned netlink.
+	LinkStatsFunc func(context.Context, string) (uint64, error)
 	// StatePath stores VM inventory that survives a runtime process restart.
 	StatePath string
 }
@@ -114,3 +116,10 @@ func (b *Backend) EnsureVXLAN(ctx context.Context, vni uint32, ifaceName, nodeIP
 func (b *Backend) SyncFDB(ctx context.Context, ifaceName string, entries []network.FDBEntry) error {
 	return b.Net.SyncFDB(ctx, ifaceName, entries)
 }
+
+func (b *Backend) LinkStats(ctx context.Context, tapName string) (uint64, error) {
+	if b.LinkStatsFunc != nil {
+		return b.LinkStatsFunc(ctx, tapName)
+	}
+	return runtimeLinkStats(ctx, tapName)
+}

internal/noderuntime/backend_test.go

@@ -152,3 +152,20 @@ func (d *fakeDriver) Snapshot(context.Context, *impdevv1alpha1.ImpVM, string) (a
 	return agent.SnapshotResult{}, nil
 }
 func (d *fakeDriver) GetVSockPath(string) (string, bool) { return "", false }
+
+func TestBackendLinkStatsUsesRuntimeOwnedImplementation(t *testing.T) {
+	backend := &Backend{LinkStatsFunc: func(_ context.Context, tapName string) (uint64, error) {
+		if tapName != "imptap-test" {
+			t.Fatalf("tapName = %q, want imptap-test", tapName)
+		}
+		return 5678, nil
+	}}
+
+	got, err := backend.LinkStats(context.Background(), "imptap-test")
+	if err != nil {
+		t.Fatalf("LinkStats() error = %v", err)
+	}
+	if got != 5678 {
+		t.Fatalf("LinkStats() = %d, want 5678", got)
+	}
+}

internal/noderuntime/linkstats_linux.go

@@ -0,0 +1,23 @@
+//go:build linux
+
+package noderuntime
+
+import (
+	"context"
+
+	"github.com/vishvananda/netlink"
+)
+
+// runtimeLinkStats reads the TAP from the runtime's network namespace. The
+// agent must use the RPC instead of attempting this lookup in its own namespace.
+func runtimeLinkStats(_ context.Context, tapName string) (uint64, error) {
+	link, err := netlink.LinkByName(tapName)
+	if err != nil {
+		return 0, err
+	}
+	stats := link.Attrs().Statistics
+	if stats == nil {
+		return 0, nil
+	}
+	return stats.RxBytes + stats.TxBytes, nil
+}

internal/noderuntime/linkstats_other.go

@@ -0,0 +1,12 @@
+//go:build !linux
+
+package noderuntime
+
+import (
+	"context"
+	"errors"
+)
+
+func runtimeLinkStats(_ context.Context, _ string) (uint64, error) {
+	return 0, errors.New("runtime link stats is not supported on this platform")
+}

internal/runtimeapi/runtime.go

@@ -48,6 +48,7 @@ type Backend interface {
 	RemoveEgressDeny(context.Context, string) error
 	EnsureVXLAN(context.Context, uint32, string, string, string) error
 	SyncFDB(context.Context, string, []network.FDBEntry) error
+	LinkStats(context.Context, string) (uint64, error)
 }
 
 // BackendFuncs makes protocol behavior testable without Firecracker.
@@ -68,6 +69,7 @@ type BackendFuncs struct {
 	RemoveEgressDenyFunc func(context.Context, string) error
 	EnsureVXLANFunc      func(context.Context, uint32, string, string, string) error
 	SyncFDBFunc          func(context.Context, string, []network.FDBEntry) error
+	LinkStatsFunc        func(context.Context, string) (uint64, error)
 }
 
 func (f BackendFuncs) Start(ctx context.Context, vm *impdevv1alpha1.ImpVM) (int64, error) {
@@ -170,6 +172,12 @@ func (f BackendFuncs) SyncFDB(ctx context.Context, ifaceName string, entries []n
 	}
 	return f.SyncFDBFunc(ctx, ifaceName, entries)
 }
+func (f BackendFuncs) LinkStats(ctx context.Context, tapName string) (uint64, error) {
+	if f.LinkStatsFunc == nil {
+		return 0, errors.New("link stats is not supported")
+	}
+	return f.LinkStatsFunc(ctx, tapName)
+}
 
 // Server serves one runtime backend over a node-local Unix socket.
 type Server struct {
@@ -243,6 +251,7 @@ type StartReply struct{ PID int64 }
 type StateReply struct{ State VMState }
 type SnapshotReply struct{ Result SnapshotResult }
 type AliveReply struct{ Alive bool }
+type LinkStatsReply struct{ Bytes uint64 }
 type VSockReply struct {
 	Path  string
 	Found bool
@@ -343,6 +352,11 @@ func (s *rpcService) EnsureVXLAN(args VXLANArgs, _ *Empty) error {
 func (s *rpcService) SyncFDB(args FDBArgs, _ *Empty) error {
 	return s.backend.SyncFDB(context.Background(), args.InterfaceName, args.Entries)
 }
+func (s *rpcService) LinkStats(args TAPArgs, reply *LinkStatsReply) error {
+	stats, err := s.backend.LinkStats(context.Background(), args.TAPName)
+	reply.Bytes = stats
+	return err
+}
 
 // Client invokes a node runtime through its Unix socket.
 type Client struct{ endpoint string }
@@ -444,3 +458,12 @@ func (c *Client) EnsureVXLAN(ctx context.Context, vni uint32, ifaceName, nodeIP,
 func (c *Client) SyncFDB(ctx context.Context, ifaceName string, entries []network.FDBEntry) error {
 	return c.call(ctx, "SyncFDB", FDBArgs{InterfaceName: ifaceName, Entries: entries}, &Empty{})
 }
+
+// LinkStats returns cumulative receive and transmit bytes for a runtime-owned TAP.
+func (c *Client) LinkStats(ctx context.Context, tapName string) (uint64, error) {
+	var reply LinkStatsReply
+	if err := c.call(ctx, "LinkStats", TAPArgs{TAPName: tapName}, &reply); err != nil {
+		return 0, err
+	}
+	return reply.Bytes, nil
+}

internal/runtimeapi/runtime_test.go

@@ -2,6 +2,7 @@ package runtimeapi_test
 
 import (
 	"context"
+	"errors"
 	"os"
 	"path/filepath"
 	"strconv"
@@ -77,3 +78,56 @@ func TestClientEnsureNetworkForwardsHostNetworkRequest(t *testing.T) {
 		t.Fatal("EnsureNetwork() did not reach runtime")
 	}
 }
+
+func TestClientLinkStatsForwardsRuntimeOwnedTAPStats(t *testing.T) {
+	t.Parallel()
+	server := runtimeapi.NewServer(runtimeapi.BackendFuncs{
+		LinkStatsFunc: func(_ context.Context, tapName string) (uint64, error) {
+			if tapName != "imptap-test" {
+				t.Fatalf("tapName = %q, want imptap-test", tapName)
+			}
+			return 1234, nil
+		},
+	})
+	dir, err := os.MkdirTemp("/tmp", "imp-runtime-")
+	if err != nil {
+		t.Fatalf("MkdirTemp() error = %v", err)
+	}
+	t.Cleanup(func() { _ = os.RemoveAll(dir) })
+	endpoint := filepath.Join(dir, "runtime.sock")
+	if err := server.Start(endpoint); err != nil {
+		t.Fatalf("Start() error = %v", err)
+	}
+	t.Cleanup(func() { _ = server.Close() })
+
+	got, err := runtimeapi.NewClient(endpoint).LinkStats(context.Background(), "imptap-test")
+	if err != nil {
+		t.Fatalf("LinkStats() error = %v", err)
+	}
+	if got != 1234 {
+		t.Fatalf("LinkStats() = %d, want 1234", got)
+	}
+}
+
+func TestClientLinkStatsReturnsRuntimeError(t *testing.T) {
+	t.Parallel()
+	server := runtimeapi.NewServer(runtimeapi.BackendFuncs{
+		LinkStatsFunc: func(context.Context, string) (uint64, error) {
+			return 0, errors.New("link missing")
+		},
+	})
+	dir, err := os.MkdirTemp("/tmp", "imp-runtime-")
+	if err != nil {
+		t.Fatalf("MkdirTemp() error = %v", err)
+	}
+	t.Cleanup(func() { _ = os.RemoveAll(dir) })
+	endpoint := filepath.Join(dir, "runtime.sock")
+	if err := server.Start(endpoint); err != nil {
+		t.Fatalf("Start() error = %v", err)
+	}
+	t.Cleanup(func() { _ = server.Close() })
+
+	if _, err := runtimeapi.NewClient(endpoint).LinkStats(context.Background(), "missing"); err == nil || err.Error() != "link missing" {
+		t.Fatalf("LinkStats() error = %v, want link missing", err)
+	}
+}