add filter

Signed-off-by: yaroslavborbat <[email protected]>
deckhouse · Oct 30, 2024 · 79d25b5 · 79d25b5
1 parent 2f705b1
commit 79d25b5
Show file tree

Hide file tree

Showing 4 changed files with 43 additions and 14 deletions.
diff --git a/images/virtualization-artifact/cmd/virtualization-controller/main.go b/images/virtualization-artifact/cmd/virtualization-controller/main.go
@@ -253,7 +253,7 @@ func main() {
 		os.Exit(1)
 	}
 
-	if _, err = vmclass.NewController(ctx, mgr, log); err != nil {
+	if _, err = vmclass.NewController(ctx, mgr, controllerNamespace, log); err != nil {
 		log.Error(err.Error())
 		os.Exit(1)
 	}

diff --git a/images/virtualization-artifact/pkg/controller/vmclass/internal/state/state.go b/images/virtualization-artifact/pkg/controller/vmclass/internal/state/state.go
@@ -40,12 +40,13 @@ type VirtualMachineClassState interface {
 }
 
 type state struct {
-	client  client.Client
-	vmClass *service.Resource[*virtv2.VirtualMachineClass, virtv2.VirtualMachineClassStatus]
+	controllerNamespace string
+	client              client.Client
+	vmClass             *service.Resource[*virtv2.VirtualMachineClass, virtv2.VirtualMachineClassStatus]
 }
 
-func New(c client.Client, vmClass *service.Resource[*virtv2.VirtualMachineClass, virtv2.VirtualMachineClassStatus]) VirtualMachineClassState {
-	return &state{client: c, vmClass: vmClass}
+func New(c client.Client, controllerNamespace string, vmClass *service.Resource[*virtv2.VirtualMachineClass, virtv2.VirtualMachineClassStatus]) VirtualMachineClassState {
+	return &state{client: c, controllerNamespace: controllerNamespace, vmClass: vmClass}
 }
 
 func (s *state) VirtualMachineClass() *service.Resource[*virtv2.VirtualMachineClass, virtv2.VirtualMachineClassStatus] {
@@ -79,8 +80,17 @@ func (s *state) Nodes(ctx context.Context) ([]corev1.Node, error) {
 	var (
 		curr        = s.vmClass.Current()
 		matchLabels map[string]string
-		filters     []common.FilterFunc[corev1.Node]
 	)
+	virtHandlerNodes, err := s.getVirtHandlerNodeNames(ctx)
+	if err != nil {
+		return nil, err
+	}
+	filters := []common.FilterFunc[corev1.Node]{
+		func(node *corev1.Node) (keep bool) {
+			_, found := virtHandlerNodes[node.GetName()]
+			return found
+		},
+	}
 
 	switch curr.Spec.CPU.Type {
 	case virtv2.CPUTypeHost, virtv2.CPUTypeHostPassthrough:
@@ -102,7 +112,7 @@ func (s *state) Nodes(ctx context.Context) ([]corev1.Node, error) {
 		return nil, fmt.Errorf("unexpected cpu type %s", curr.Spec.CPU.Type)
 	}
 	nodes := &corev1.NodeList{}
-	err := s.client.List(
+	err = s.client.List(
 		ctx,
 		nodes,
 		client.MatchingLabelsSelector{Selector: labels.SelectorFromSet(matchLabels)})
@@ -113,6 +123,22 @@ func (s *state) Nodes(ctx context.Context) ([]corev1.Node, error) {
 	return nodeFilter(nodes.Items, filters...), nil
 }
 
+func (s *state) getVirtHandlerNodeNames(ctx context.Context) (map[string]struct{}, error) {
+	pods := &corev1.PodList{}
+	err := s.client.List(ctx, pods, client.InNamespace(s.controllerNamespace),
+		client.MatchingLabelsSelector{
+			Selector: labels.SelectorFromSet(map[string]string{
+				virtv1.AppLabel: "virt-handler"})})
+	if err != nil {
+		return nil, fmt.Errorf("failed to list pods: %w", err)
+	}
+	nodes := make(map[string]struct{}, len(pods.Items))
+	for _, pod := range pods.Items {
+		nodes[pod.Spec.NodeName] = struct{}{}
+	}
+	return nodes, nil
+}
+
 func (s *state) AvailableNodes(nodes []corev1.Node) ([]corev1.Node, error) {
 	if s.vmClass == nil || s.vmClass.IsEmpty() {
 		return nil, nil

diff --git a/images/virtualization-artifact/pkg/controller/vmclass/vmclass_controller.go b/images/virtualization-artifact/pkg/controller/vmclass/vmclass_controller.go
@@ -38,6 +38,7 @@ const (
 func NewController(
 	ctx context.Context,
 	mgr manager.Manager,
+	controllerNamespace string,
 	log *slog.Logger,
 ) (controller.Controller, error) {
 	log = log.With(logger.SlogController(controllerName))
@@ -49,7 +50,7 @@ func NewController(
 		internal.NewDiscoveryHandler(),
 		internal.NewLifeCycleHandler(client),
 	}
-	r := NewReconciler(client, handlers...)
+	r := NewReconciler(controllerNamespace, client, handlers...)
 
 	c, err := controller.New(controllerName, mgr, controller.Options{
 		Reconciler:       r,

diff --git a/images/virtualization-artifact/pkg/controller/vmclass/vmclass_reconciler.go b/images/virtualization-artifact/pkg/controller/vmclass/vmclass_reconciler.go
@@ -45,16 +45,18 @@ type Handler interface {
 	Name() string
 }
 
-func NewReconciler(client client.Client, handlers ...Handler) *Reconciler {
+func NewReconciler(controllerNamespace string, client client.Client, handlers ...Handler) *Reconciler {
 	return &Reconciler{
-		client:   client,
-		handlers: handlers,
+		controllerNamespace: controllerNamespace,
+		client:              client,
+		handlers:            handlers,
 	}
 }
 
 type Reconciler struct {
-	client   client.Client
-	handlers []Handler
+	controllerNamespace string
+	client              client.Client
+	handlers            []Handler
 }
 
 func (r *Reconciler) SetupController(_ context.Context, mgr manager.Manager, ctr controller.Controller) error {
@@ -138,7 +140,7 @@ func (r *Reconciler) Reconcile(ctx context.Context, req reconcile.Request) (reco
 		log.Info("Reconcile observe an absent VirtualMachineClass: it may be deleted")
 		return reconcile.Result{}, nil
 	}
-	s := state.New(r.client, class)
+	s := state.New(r.client, r.controllerNamespace, class)
 
 	log.Debug("Start reconcile VMClass")