fix: correctly handle input finalizers in InfraMachineController

Improve the handling of extra mapped input finalizers in `InfraMachineController`, also add a finalizer to `MachineExtensions` resources. Signed-off-by: Utku Ozdemir <[email protected]>
siderolabs · Dec 18, 2024 · 3332684 · 3332684
1 parent b7c3c50
commit 3332684
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 33 deletions.
diff --git a/internal/backend/runtime/omni/controllers/omni/infra_machine.go b/internal/backend/runtime/omni/controllers/omni/infra_machine.go
@@ -84,6 +84,16 @@ func NewInfraMachineController() *InfraMachineController {
 type infraMachineControllerHelper struct{}
 
 func (h *infraMachineControllerHelper) transformExtraOutput(ctx context.Context, r controller.ReaderWriter, _ *zap.Logger, link *siderolink.Link, infraMachine *infra.Machine) error {
+	config, err := helpers.HandleInput[*omni.InfraMachineConfig](ctx, r, InfraMachineControllerName, link)
+	if err != nil {
+		return err
+	}
+
+	machineExts, err := helpers.HandleInput[*omni.MachineExtensions](ctx, r, InfraMachineControllerName, link)
+	if err != nil {
+		return err
+	}
+
 	providerID, ok := link.Metadata().Annotations().Get(omni.LabelInfraProviderID)
 	if !ok {
 		return xerrors.NewTaggedf[qtransform.SkipReconcileTag]("the link is not created by an infra provider")
@@ -98,7 +108,7 @@ func (h *infraMachineControllerHelper) transformExtraOutput(ctx context.Context,
 		return xerrors.NewTaggedf[qtransform.SkipReconcileTag]("the link is not created by a static infra provider")
 	}
 
-	if err = h.applyInfraMachineConfig(ctx, r, link, infraMachine); err != nil {
+	if err = h.applyInfraMachineConfig(infraMachine, config); err != nil {
 		return err
 	}
 
@@ -127,18 +137,16 @@ func (h *infraMachineControllerHelper) transformExtraOutput(ctx context.Context,
 		infraMachine.TypedSpec().Value.ClusterTalosVersion = ""
 		infraMachine.TypedSpec().Value.Extensions = nil
 
-		if err = r.RemoveFinalizer(ctx, clusterMachine.Metadata(), InfraMachineControllerName); err != nil {
-			return err
-		}
+		_, err = helpers.HandleInput[*omni.ClusterMachine](ctx, r, InfraMachineControllerName, link)
 
-		return nil
+		return err
 	}
 
 	if err = r.AddFinalizer(ctx, clusterMachine.Metadata(), InfraMachineControllerName); err != nil {
 		return err
 	}
 
-	talosVersion, extensions, err := h.getClusterInfo(ctx, r, link.Metadata().ID())
+	talosVersion, extensions, err := h.getClusterInfo(ctx, r, link.Metadata().ID(), machineExts)
 	if err != nil {
 		return err
 	}
@@ -152,34 +160,21 @@ func (h *infraMachineControllerHelper) transformExtraOutput(ctx context.Context,
 }
 
 func (h *infraMachineControllerHelper) finalizerRemovalExtraOutput(ctx context.Context, r controller.ReaderWriter, _ *zap.Logger, link *siderolink.Link) error {
-	clusterMachine, err := safe.ReaderGetByID[*omni.ClusterMachine](ctx, r, link.Metadata().ID())
-	if err != nil {
-		if state.IsNotFoundError(err) {
-			return nil
-		}
-
+	if _, err := helpers.HandleInput[*omni.InfraMachineConfig](ctx, r, InfraMachineControllerName, link); err != nil {
 		return err
 	}
 
-	if err = r.RemoveFinalizer(ctx, clusterMachine.Metadata(), InfraMachineControllerName); err != nil {
+	if _, err := helpers.HandleInput[*omni.MachineExtensions](ctx, r, InfraMachineControllerName, link); err != nil {
 		return err
 	}
 
-	_, err = helpers.HandleInput[*omni.InfraMachineConfig](ctx, r, InfraMachineControllerName, link)
-	if err != nil {
-		return err
-	}
+	_, err := helpers.HandleInput[*omni.ClusterMachine](ctx, r, InfraMachineControllerName, link)
 
-	return nil
+	return err
 }
 
 // applyInfraMachineConfig applies the user-managed configuration from the omni.InfraMachineConfig resource into the infra.Machine.
-func (h *infraMachineControllerHelper) applyInfraMachineConfig(ctx context.Context, r controller.ReaderWriter, link *siderolink.Link, infraMachine *infra.Machine) error {
-	config, err := helpers.HandleInput[*omni.InfraMachineConfig](ctx, r, InfraMachineControllerName, link)
-	if err != nil {
-		return err
-	}
-
+func (h *infraMachineControllerHelper) applyInfraMachineConfig(infraMachine *infra.Machine, config *omni.InfraMachineConfig) error {
 	const defaultPreferredPowerState = specs.InfraMachineSpec_POWER_STATE_OFF // todo: introduce a resource to configure this globally or per-provider level
 
 	// reset the user-override fields except the "Accepted" field
@@ -219,7 +214,7 @@ func (h *infraMachineControllerHelper) applyInfraMachineConfig(ctx context.Conte
 // getClusterInfo returns the Talos version and extensions for the given machine.
 //
 // At this point, the machine is known to be associated with a cluster.
-func (h *infraMachineControllerHelper) getClusterInfo(ctx context.Context, r controller.Reader, id resource.ID) (string, []string, error) {
+func (h *infraMachineControllerHelper) getClusterInfo(ctx context.Context, r controller.Reader, id resource.ID, machineExtensions *omni.MachineExtensions) (string, []string, error) {
 	schematicConfig, err := safe.ReaderGetByID[*omni.SchematicConfiguration](ctx, r, id)
 	if err != nil {
 		if state.IsNotFoundError(err) {
@@ -229,14 +224,11 @@ func (h *infraMachineControllerHelper) getClusterInfo(ctx context.Context, r con
 		return "", nil, err
 	}
 
-	machineExts, err := safe.ReaderGetByID[*omni.MachineExtensions](ctx, r, schematicConfig.Metadata().ID())
-	if err != nil {
-		if state.IsNotFoundError(err) { // no extensions
-			return schematicConfig.TypedSpec().Value.TalosVersion, nil, nil
-		}
+	var extensions []string
 
-		return "", nil, err
+	if machineExtensions != nil {
+		extensions = machineExtensions.TypedSpec().Value.Extensions
 	}
 
-	return schematicConfig.TypedSpec().Value.TalosVersion, machineExts.TypedSpec().Value.Extensions, nil
+	return schematicConfig.TypedSpec().Value.TalosVersion, extensions, nil
 }
diff --git a/internal/backend/runtime/omni/controllers/omni/infra_machine_test.go b/internal/backend/runtime/omni/controllers/omni/infra_machine_test.go
@@ -62,6 +62,10 @@ func (suite *InfraMachineControllerSuite) TestReconcile() {
 
 	suite.Require().NoError(suite.state.Create(suite.ctx, config))
 
+	assertResource[*omni.InfraMachineConfig](&suite.OmniSuite, config.Metadata(), func(r *omni.InfraMachineConfig, assertion *assert.Assertions) {
+		assertion.True(r.Metadata().Finalizers().Has(omnictrl.InfraMachineControllerName))
+	})
+
 	assertResource[*infra.Machine](&suite.OmniSuite, infraMachineMD, func(r *infra.Machine, assertion *assert.Assertions) {
 		assertion.Equal(specs.InfraMachineConfigSpec_ACCEPTED, r.TypedSpec().Value.AcceptanceStatus)
 		assertion.Equal(specs.InfraMachineSpec_POWER_STATE_ON, r.TypedSpec().Value.PreferredPowerState)
@@ -97,6 +101,10 @@ func (suite *InfraMachineControllerSuite) TestReconcile() {
 
 	suite.Require().NoError(suite.state.Create(suite.ctx, extensions))
 
+	assertResource[*omni.MachineExtensions](&suite.OmniSuite, extensions.Metadata(), func(r *omni.MachineExtensions, assertion *assert.Assertions) {
+		assertion.True(r.Metadata().Finalizers().Has(omnictrl.InfraMachineControllerName))
+	})
+
 	// assert that the cluster machine has the correct extensions
 	assertResource[*infra.Machine](&suite.OmniSuite, clusterMachine.Metadata(), func(r *infra.Machine, assertion *assert.Assertions) {
 		assertion.ElementsMatch([]string{"foo", "bar"}, r.TypedSpec().Value.Extensions)
@@ -127,11 +135,19 @@ func (suite *InfraMachineControllerSuite) TestReconcile() {
 	// destroy the link
 	rtestutils.Destroy[*siderolink.Link](suite.ctx, suite.T(), suite.state, []string{link.Metadata().ID()})
 
-	// assert that the finalizer is removed
+	// assert that the finalizers are removed
 	assertResource[*omni.ClusterMachine](&suite.OmniSuite, infraMachineMD, func(r *omni.ClusterMachine, assertion *assert.Assertions) {
 		assertion.False(r.Metadata().Finalizers().Has(omnictrl.InfraMachineControllerName))
 	})
 
+	assertResource[*omni.InfraMachineConfig](&suite.OmniSuite, infraMachineMD, func(r *omni.InfraMachineConfig, assertion *assert.Assertions) {
+		assertion.False(r.Metadata().Finalizers().Has(omnictrl.InfraMachineControllerName))
+	})
+
+	assertResource[*omni.MachineExtensions](&suite.OmniSuite, infraMachineMD, func(r *omni.MachineExtensions, assertion *assert.Assertions) {
+		assertion.False(r.Metadata().Finalizers().Has(omnictrl.InfraMachineControllerName))
+	})
+
 	// assert that infra.Machine is removed
 	assertNoResource[*infra.Machine](&suite.OmniSuite, infraMachine)
 }