fix(controller,agent): restart flow, runner demand scaling, and composite layer wiring

Author: syscod3

api/v1alpha1/impvmrunnerpool_types.go

@@ -49,7 +49,7 @@ type RunnerPlatformSpec struct {
 
 // RunnerScopeSpec selects org-level or repo-level runner registration.
 // Exactly one of Org or Repo must be set.
-// +kubebuilder:validation:XValidation:rule="!(size(self.org) > 0 && size(self.repo) > 0)",message="org and repo are mutually exclusive; set exactly one"
+// +kubebuilder:validation:XValidation:rule="(size(self.org) > 0) != (size(self.repo) > 0)",message="set exactly one of org or repo"
 type RunnerScopeSpec struct {
 	// Org registers a runner for the entire organisation.
 	// +optional

api/v1alpha1/impvmrunnerpool_types_test.go

@@ -2,6 +2,8 @@ package v1alpha1
 
 import (
 	"encoding/json"
+	"os"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -36,3 +38,14 @@ func TestImpVMRunnerPool_roundTrip(t *testing.T) {
 	assert.Equal(t, "my-org", out.Spec.Platform.Scope.Org)
 	assert.Equal(t, int32(10), out.Spec.Scaling.MaxConcurrent)
 }
+
+func TestImpVMRunnerPoolCRD_scopeValidationRequiresExactlyOne(t *testing.T) {
+	b, err := os.ReadFile("../../config/crd/bases/imp.dev_impvmrunnerpools.yaml")
+	require.NoError(t, err)
+
+	yaml := string(b)
+	assert.Contains(t, yaml, "set exactly one of org or repo")
+	assert.Contains(t, yaml, "(size(self.org) > 0) != (size(self.repo) > 0)")
+	assert.False(t, strings.Contains(yaml, "!(size(self.org) > 0 && size(self.repo) > 0)"),
+		"scope validation should require at least one of org or repo, not only mutual exclusion")
+}

config/crd/bases/imp.dev_impvmrunnerpools.yaml

@@ -117,8 +117,8 @@ spec:
                         type: string
                     type: object
                     x-kubernetes-validations:
-                    - message: org and repo are mutually exclusive; set exactly one
-                      rule: '!(size(self.org) > 0 && size(self.repo) > 0)'
+                    - message: set exactly one of org or repo
+                      rule: (size(self.org) > 0) != (size(self.repo) > 0)
                   serverURL:
                     description: ServerURL is required for GitLab and Forgejo. Leave
                       empty for github.com.

config/rbac/role.yaml

@@ -20,6 +20,12 @@ rules:
   verbs:
   - create
   - patch
+- apiGroups:
+  - ""
+  resources:
+  - secrets
+  verbs:
+  - get
 - apiGroups:
   - cilium.io
   resources:

internal/agent/firecracker_driver.go

@@ -50,6 +50,11 @@ type fcProc struct {
 	probeCancel  context.CancelFunc   // non-nil when probe goroutine is running
 }
 
+type rootfsBuilder interface {
+	Build(ctx context.Context, imageRef string, opts ...rootfs.BuildOption) (string, error)
+	BuildComposite(ctx context.Context, baseImage string, extraLayers []string, opts ...rootfs.BuildOption) (string, error)
+}
+
 // FirecrackerDriver is a VMDriver that launches real Firecracker microVMs.
 // It is safe for concurrent use.
 type FirecrackerDriver struct {
@@ -62,7 +67,7 @@ type FirecrackerDriver struct {
 	// KernelArgs is the kernel command-line string.
 	KernelArgs string
 	// Cache builds and caches ext4 rootfs images from OCI images.
-	Cache *rootfs.Builder
+	Cache rootfsBuilder
 	// Client is the controller-runtime Kubernetes client.
 	Client ctrlclient.Client
 	// Net manages host-level networking (bridge, TAP, NAT).
@@ -153,7 +158,7 @@ func (d *FirecrackerDriver) Start(ctx context.Context, vm *impdevv1alpha1.ImpVM)
 	if gaEnabled {
 		buildOpts = append(buildOpts, rootfs.WithGuestAgent(d.guestAgentPath()))
 	}
-	rootfsPath, err := d.Cache.Build(ctx, vm.Spec.Image, buildOpts...)
+	rootfsPath, err := d.buildRootfs(ctx, vm, buildOpts...)
 	if err != nil {
 		return 0, fmt.Errorf("build rootfs for %s: %w", vm.Spec.Image, err)
 	}
@@ -235,6 +240,22 @@ func (d *FirecrackerDriver) Start(ctx context.Context, vm *impdevv1alpha1.ImpVM)
 	return int64(pid), nil
 }
 
+func (d *FirecrackerDriver) buildRootfs(
+	ctx context.Context, vm *impdevv1alpha1.ImpVM, opts ...rootfs.BuildOption,
+) (string, error) {
+	extraLayers := make([]string, 0, 2)
+	if vm.Spec.RunnerLayer != "" {
+		extraLayers = append(extraLayers, vm.Spec.RunnerLayer)
+	}
+	if vm.Spec.CiliumLayer != "" {
+		extraLayers = append(extraLayers, vm.Spec.CiliumLayer)
+	}
+	if len(extraLayers) == 0 {
+		return d.Cache.Build(ctx, vm.Spec.Image, opts...)
+	}
+	return d.Cache.BuildComposite(ctx, vm.Spec.Image, extraLayers, opts...)
+}
+
 // Stop implements VMDriver. Sends a graceful ACPI shutdown signal, waits up to
 // shuttingDownTimeout, then force-terminates the Firecracker process and cleans up
 // the Unix socket. Safe to call on a VM that was never started or already stopped.

internal/agent/firecracker_driver_test.go

@@ -4,6 +4,7 @@ package agent
 
 import (
 	"context"
+	"fmt"
 	"os"
 	"os/exec"
 	"testing"
@@ -20,6 +21,36 @@ import (
 	pb "github.com/syscode-labs/imp/internal/proto/guest"
 )
 
+type fakeRootfsBuilder struct {
+	buildCalled          bool
+	buildCompositeCalled bool
+	lastBaseImage        string
+	lastExtraLayers      []string
+	returnPath           string
+	returnErr            error
+}
+
+func (f *fakeRootfsBuilder) Build(_ context.Context, imageRef string, _ ...rootfs.BuildOption) (string, error) {
+	f.buildCalled = true
+	f.lastBaseImage = imageRef
+	if f.returnErr != nil {
+		return "", f.returnErr
+	}
+	return f.returnPath, nil
+}
+
+func (f *fakeRootfsBuilder) BuildComposite(
+	_ context.Context, baseImage string, extraLayers []string, _ ...rootfs.BuildOption,
+) (string, error) {
+	f.buildCompositeCalled = true
+	f.lastBaseImage = baseImage
+	f.lastExtraLayers = append([]string(nil), extraLayers...)
+	if f.returnErr != nil {
+		return "", f.returnErr
+	}
+	return f.returnPath, nil
+}
+
 // hasFirecrackerBin returns true if the firecracker binary is available.
 func hasFirecrackerBin() bool {
 	_, err := exec.LookPath("firecracker")
@@ -36,6 +67,56 @@ func TestFirecrackerDriverPlaceholder(t *testing.T) {
 	t.Log("FirecrackerDriver test file compiles correctly")
 }
 
+func TestFirecrackerDriver_buildRootfs_UsesBuildWhenNoExtraLayers(t *testing.T) {
+	b := &fakeRootfsBuilder{returnPath: "/tmp/base.ext4"}
+	d := &FirecrackerDriver{Cache: b}
+
+	vm := &impdevv1alpha1.ImpVM{}
+	vm.Spec.Image = "ghcr.io/example/base:latest"
+
+	path, err := d.buildRootfs(context.Background(), vm)
+	if err != nil {
+		t.Fatalf("buildRootfs: %v", err)
+	}
+	if !b.buildCalled {
+		t.Fatal("expected Build to be called")
+	}
+	if b.buildCompositeCalled {
+		t.Fatal("did not expect BuildComposite to be called")
+	}
+	if path != "/tmp/base.ext4" {
+		t.Fatalf("path = %q, want %q", path, "/tmp/base.ext4")
+	}
+}
+
+func TestFirecrackerDriver_buildRootfs_UsesBuildCompositeWithLayerOrder(t *testing.T) {
+	b := &fakeRootfsBuilder{returnPath: "/tmp/composite.ext4"}
+	d := &FirecrackerDriver{Cache: b}
+
+	vm := &impdevv1alpha1.ImpVM{}
+	vm.Spec.Image = "ghcr.io/example/base:latest"
+	vm.Spec.RunnerLayer = "ghcr.io/example/runner:v1"
+	vm.Spec.CiliumLayer = "ghcr.io/example/cilium:v1"
+
+	path, err := d.buildRootfs(context.Background(), vm)
+	if err != nil {
+		t.Fatalf("buildRootfs: %v", err)
+	}
+	if !b.buildCompositeCalled {
+		t.Fatal("expected BuildComposite to be called")
+	}
+	if b.buildCalled {
+		t.Fatal("did not expect Build to be called")
+	}
+	wantLayers := []string{"ghcr.io/example/runner:v1", "ghcr.io/example/cilium:v1"}
+	if fmt.Sprint(b.lastExtraLayers) != fmt.Sprint(wantLayers) {
+		t.Fatalf("extraLayers = %v, want %v", b.lastExtraLayers, wantLayers)
+	}
+	if path != "/tmp/composite.ext4" {
+		t.Fatalf("path = %q, want %q", path, "/tmp/composite.ext4")
+	}
+}
+
 func TestFirecrackerDriver_SocketPath(t *testing.T) {
 	d := &FirecrackerDriver{SocketDir: "/run/imp/sockets"}
 

internal/controller/impvmrunnerpool_controller.go

@@ -18,8 +18,10 @@ package controller
 
 import (
 	"context"
+	"fmt"
 	"time"
 
+	corev1 "k8s.io/api/core/v1"
 	apierrors "k8s.io/apimachinery/pkg/api/errors"
 	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
 	"k8s.io/apimachinery/pkg/runtime"
@@ -28,18 +30,31 @@ import (
 	logf "sigs.k8s.io/controller-runtime/pkg/log"
 
 	impv1alpha1 "github.com/syscode-labs/imp/api/v1alpha1"
+	"github.com/syscode-labs/imp/internal/runner"
 )
 
 // ImpVMRunnerPoolReconciler reconciles ImpVMRunnerPool objects.
 type ImpVMRunnerPoolReconciler struct {
 	client.Client
-	Scheme *runtime.Scheme
+	Scheme        *runtime.Scheme
+	DriverFactory RunnerDriverFactory
 }
 
+type runnerQueueDepthReader interface {
+	QueueDepth(ctx context.Context) (int, error)
+}
+
+type RunnerDriverFactory func(
+	ctx context.Context,
+	c client.Client,
+	pool *impv1alpha1.ImpVMRunnerPool,
+) (runnerQueueDepthReader, error)
+
 // +kubebuilder:rbac:groups=imp.dev,resources=impvmrunnerpools,verbs=get;list;watch;create;update;patch;delete
 // +kubebuilder:rbac:groups=imp.dev,resources=impvmrunnerpools/status,verbs=get;update;patch
 // +kubebuilder:rbac:groups=imp.dev,resources=impvmtemplates,verbs=get;list;watch
 // +kubebuilder:rbac:groups=imp.dev,resources=impvms,verbs=get;list;watch;create;delete
+// +kubebuilder:rbac:groups=core,resources=secrets,verbs=get
 
 func (r *ImpVMRunnerPoolReconciler) Reconcile(ctx context.Context, req ctrl.Request) (ctrl.Result, error) {
 	log := logf.FromContext(ctx)
@@ -113,7 +128,15 @@ func (r *ImpVMRunnerPoolReconciler) Reconcile(ctx context.Context, req ctrl.Requ
 		maxConcurrent = pool.Spec.Scaling.MaxConcurrent
 	}
 
-	toCreate := minIdle - activeCount
+	desiredCount := minIdle
+	queueDepth, err := r.queueDepth(ctx, pool)
+	if err != nil {
+		log.Info("could not fetch runner queue depth; falling back to minIdle", "pool", pool.Name, "err", err)
+	} else if int32(queueDepth) > desiredCount {
+		desiredCount = int32(queueDepth)
+	}
+
+	toCreate := desiredCount - activeCount
 	if toCreate < 0 {
 		toCreate = 0
 	}
@@ -153,6 +176,24 @@ func (r *ImpVMRunnerPoolReconciler) Reconcile(ctx context.Context, req ctrl.Requ
 	return ctrl.Result{RequeueAfter: requeueAfter}, nil
 }
 
+func (r *ImpVMRunnerPoolReconciler) queueDepth(ctx context.Context, pool *impv1alpha1.ImpVMRunnerPool) (int, error) {
+	if pool.Spec.JobDetection == nil ||
+		pool.Spec.JobDetection.Polling == nil ||
+		!pool.Spec.JobDetection.Polling.Enabled {
+		return 0, nil
+	}
+
+	factory := r.DriverFactory
+	if factory == nil {
+		factory = defaultRunnerDriverFactory
+	}
+	d, err := factory(ctx, r.Client, pool)
+	if err != nil {
+		return 0, err
+	}
+	return d.QueueDepth(ctx)
+}
+
 func (r *ImpVMRunnerPoolReconciler) createRunnerVM(ctx context.Context, pool *impv1alpha1.ImpVMRunnerPool, tpl *impv1alpha1.ImpVMTemplate) error {
 	classRef := tpl.Spec.ClassRef
 	vm := &impv1alpha1.ImpVM{
@@ -179,6 +220,14 @@ func (r *ImpVMRunnerPoolReconciler) createRunnerVM(ctx context.Context, pool *im
 	if tpl.Spec.NetworkGroup != "" {
 		vm.Spec.NetworkGroup = tpl.Spec.NetworkGroup
 	}
+	if pool.Spec.RunnerLayer != "" {
+		vm.Spec.RunnerLayer = pool.Spec.RunnerLayer
+	} else if tpl.Spec.RunnerLayer != "" {
+		vm.Spec.RunnerLayer = tpl.Spec.RunnerLayer
+	}
+	if tpl.Spec.CiliumLayer != "" {
+		vm.Spec.CiliumLayer = tpl.Spec.CiliumLayer
+	}
 	if err := ctrl.SetControllerReference(pool, vm, r.Scheme); err != nil {
 		return err
 	}
@@ -191,3 +240,76 @@ func (r *ImpVMRunnerPoolReconciler) SetupWithManager(mgr ctrl.Manager) error {
 		Owns(&impv1alpha1.ImpVM{}).
 		Complete(r)
 }
+
+func defaultRunnerDriverFactory(
+	ctx context.Context,
+	c client.Client,
+	pool *impv1alpha1.ImpVMRunnerPool,
+) (runnerQueueDepthReader, error) {
+	var creds corev1.Secret
+	if err := c.Get(ctx, client.ObjectKey{
+		Namespace: pool.Namespace,
+		Name:      pool.Spec.Platform.CredentialsSecret,
+	}, &creds); err != nil {
+		return nil, err
+	}
+	token := pickSecretValue(creds.Data, "token")
+	if token == "" {
+		return nil, fmt.Errorf("credentials secret %s/%s has no token value", pool.Namespace, creds.Name)
+	}
+
+	scope, err := platformScope(pool)
+	if err != nil {
+		return nil, err
+	}
+
+	switch pool.Spec.Platform.Type {
+	case "github-actions":
+		return runner.NewGitHubDriver(token, scope, nil)
+	case "forgejo":
+		return runner.NewForgejoDriver(token, pool.Spec.Platform.ServerURL, scope, nil)
+	case "gitlab":
+		return runner.NewGitLabDriver(token, pool.Spec.Platform.ServerURL, scope, nil)
+	default:
+		return nil, fmt.Errorf("unsupported platform type %q", pool.Spec.Platform.Type)
+	}
+}
+
+func platformScope(pool *impv1alpha1.ImpVMRunnerPool) (string, error) {
+	if pool.Spec.Platform.Scope == nil {
+		return "", fmt.Errorf("platform.scope is required")
+	}
+	scope := pool.Spec.Platform.Scope
+	switch pool.Spec.Platform.Type {
+	case "github-actions", "forgejo":
+		if scope.Org != "" {
+			return "org:" + scope.Org, nil
+		}
+		if scope.Repo != "" {
+			return "repo:" + scope.Repo, nil
+		}
+	case "gitlab":
+		if scope.Org != "" {
+			return "group:" + scope.Org, nil
+		}
+		if scope.Repo != "" {
+			return "project:" + scope.Repo, nil
+		}
+	}
+	return "", fmt.Errorf("invalid platform.scope for type %q", pool.Spec.Platform.Type)
+}
+
+func pickSecretValue(m map[string][]byte, preferredKey string) string {
+	if len(m) == 0 {
+		return ""
+	}
+	if v, ok := m[preferredKey]; ok && len(v) > 0 {
+		return string(v)
+	}
+	for _, v := range m {
+		if len(v) > 0 {
+			return string(v)
+		}
+	}
+	return ""
+}