treat component notification as critical during migration process

Author: michalpristas

internal/pkg/agent/application/actions/handlers/handler_action_migrate.go

@@ -27,7 +27,7 @@ const ()
 type migrateCoordinator interface {
 	actionCoordinator
 
-	Migrate(_ context.Context, _ *fleetapi.ActionMigrate, _ func(done <-chan struct{}) backoff.Backoff) error
+	Migrate(_ context.Context, _ *fleetapi.ActionMigrate, _ func(done <-chan struct{}) backoff.Backoff, _ func(context.Context, *fleetapi.ActionMigrate) error) error
 	ReExec(callback reexec.ShutdownCallbackFn, argOverrides ...string)
 	Protection() protection.Config
 }
@@ -85,7 +85,7 @@ func (h *Migrate) Handle(ctx context.Context, a fleetapi.Action, ack acker.Acker
 		}
 	}
 
-	if err := h.coord.Migrate(ctx, action, fleetgateway.RequestBackoff); err != nil {
+	if err := h.coord.Migrate(ctx, action, fleetgateway.RequestBackoff, h.notifyComponents); err != nil {
 		// this should not happen, unmanaged agent should not receive the action
 		// defensive coding to avoid misbehavior
 		if errors.Is(err, coordinator.ErrNotManaged) {
@@ -102,12 +102,6 @@ func (h *Migrate) Handle(ctx context.Context, a fleetapi.Action, ack acker.Acker
 		return fmt.Errorf("migration of agent to a new cluster failed: %w", err)
 	}
 
-	// action is all rigth we can notify endpoint
-	if err := h.notifyComponents(ctx, action); err != nil {
-		// config is cleaned up already we cannot revert
-		h.log.Warnf("failed to notify components, aborting: %v", err)
-	}
-
 	// reexec and load new config
 	h.coord.ReExec(nil)
 	return nil

internal/pkg/agent/application/actions/handlers/handler_action_migrate_test.go

@@ -17,15 +17,12 @@ import (
 	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
-	"github.com/elastic/elastic-agent-client/v7/pkg/client"
-	"github.com/elastic/elastic-agent-client/v7/pkg/proto"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/coordinator"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/reexec"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/protection"
 	"github.com/elastic/elastic-agent/internal/pkg/core/backoff"
 	"github.com/elastic/elastic-agent/internal/pkg/fleetapi"
 	"github.com/elastic/elastic-agent/pkg/component"
-	"github.com/elastic/elastic-agent/pkg/component/runtime"
 	"github.com/elastic/elastic-agent/pkg/core/logger/loggertest"
 	mockinfo "github.com/elastic/elastic-agent/testing/mocks/internal_/pkg/agent/application/info"
 )
@@ -42,7 +39,7 @@ func TestActionMigratelHandler(t *testing.T) {
 		ack.On("Commit", t.Context()).Return(nil)
 
 		coord := &fakeMigrateCoordinator{}
-		coord.On("Migrate", mock.Anything, mock.Anything).Return(nil)
+		coord.On("Migrate", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		coord.On("ReExec", mock.Anything, mock.Anything)
 		coord.On("Protection").Return(protection.Config{SignatureValidationKey: nil})
 
@@ -82,7 +79,7 @@ func TestActionMigratelHandler(t *testing.T) {
 
 				coord := &fakeMigrateCoordinator{}
 				coord.On("State").Return(coordinator.State{})
-				coord.On("Migrate", mock.Anything, mock.Anything).Return(nil)
+				coord.On("Migrate", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 				coord.On("ReExec", mock.Anything, mock.Anything)
 				coord.On("Protection").Return(protection.Config{SignatureValidationKey: nil, Enabled: tc.protectionEnabled})
 
@@ -120,7 +117,7 @@ func TestActionMigratelHandler(t *testing.T) {
 
 		coord := &fakeMigrateCoordinator{}
 		coord.On("State").Return(coordinator.State{})
-		coord.On("Migrate", mock.Anything, mock.Anything).Return(nil)
+		coord.On("Migrate", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		coord.On("ReExec", mock.Anything, mock.Anything)
 		coord.On("Protection").Return(protection.Config{SignatureValidationKey: nil})
 
@@ -170,7 +167,7 @@ func TestActionMigratelHandler(t *testing.T) {
 
 		coord := &fakeMigrateCoordinator{}
 		coord.On("State").Return(coordinator.State{})
-		coord.On("Migrate", mock.Anything, mock.Anything).Return(nil)
+		coord.On("Migrate", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		coord.On("ReExec", mock.Anything, mock.Anything)
 		coord.On("Protection").Return(protection.Config{SignatureValidationKey: signatureValidationKey})
 
@@ -206,7 +203,7 @@ func TestActionMigratelHandler(t *testing.T) {
 		ack.On("Commit", t.Context()).Return(nil)
 
 		coord := &fakeMigrateCoordinator{}
-		coord.On("Migrate", mock.Anything, mock.Anything).Return(nil)
+		coord.On("Migrate", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		coord.On("ReExec", mock.Anything, mock.Anything)
 		coord.On("Protection").Return(protection.Config{SignatureValidationKey: signatureValidationKey})
 
@@ -256,7 +253,7 @@ func TestActionMigratelHandler(t *testing.T) {
 
 		coord := &fakeMigrateCoordinator{}
 		coord.On("State").Return(coordinator.State{})
-		coord.On("Migrate", mock.Anything, mock.Anything).Return(nil)
+		coord.On("Migrate", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		coord.On("ReExec", mock.Anything, mock.Anything)
 		coord.On("Protection").Return(protection.Config{SignatureValidationKey: nil})
 
@@ -308,7 +305,7 @@ func TestActionMigratelHandler(t *testing.T) {
 		ack.On("Commit", t.Context()).Return(nil)
 
 		coord := &fakeMigrateCoordinator{}
-		coord.On("Migrate", mock.Anything, mock.Anything).Return(nil)
+		coord.On("Migrate", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil)
 		coord.On("ReExec", mock.Anything, mock.Anything)
 		coord.On("Protection").Return(protection.Config{SignatureValidationKey: signatureValidationKey})
 
@@ -331,7 +328,7 @@ func TestActionMigratelHandler(t *testing.T) {
 
 		coord := &fakeMigrateCoordinator{}
 		coord.On("State").Return(coordinator.State{})
-		coord.On("Migrate", mock.Anything, mock.Anything).Return(coordinator.ErrFleetServer)
+		coord.On("Migrate", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(coordinator.ErrFleetServer)
 		coord.On("ReExec", mock.Anything, mock.Anything)
 		coord.On("Protection").Return(protection.Config{SignatureValidationKey: nil})
 
@@ -346,65 +343,14 @@ func TestActionMigratelHandler(t *testing.T) {
 		ack.AssertNumberOfCalls(t, "Commit", 1)
 		coord.AssertNotCalled(t, "ReExec", mock.Anything, mock.Anything)
 	})
-
-	t.Run("endpoint notified", func(t *testing.T) {
-		mockAgentInfo := mockinfo.NewAgent(t)
-		mockAgentInfo.On("AgentID").Return("agent-id")
-		action := &fleetapi.ActionMigrate{
-			ActionType: fleetapi.ActionTypeMigrate,
-		}
-
-		ack := &fakeAcker{}
-		ack.On("Ack", t.Context(), action).Return(nil)
-		ack.On("Commit", t.Context()).Return(nil)
-
-		coord := &fakeMigrateCoordinator{}
-		coord.On("State").Return(coordinator.State{
-			Components: []runtime.ComponentComponentState{
-				runtime.ComponentComponentState{
-					Component: component.Component{
-						InputSpec: &component.InputRuntimeSpec{
-							Spec: component.InputSpec{
-								ProxiedActions: []string{fleetapi.ActionTypeMigrate},
-							},
-						},
-						Units: []component.Unit{
-							component.Unit{
-								Type: client.UnitTypeInput,
-								Config: &proto.UnitExpectedConfig{
-									Type: "migrate-sensitive-input",
-								},
-							},
-						},
-						InputType: "migrate-sensitive-input",
-					},
-				},
-			},
-		})
-		coord.On("PerformAction", mock.Anything, mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(map[string]interface{}{}, nil)
-		coord.On("Migrate", mock.Anything, mock.Anything).Return(nil)
-		coord.On("ReExec", mock.Anything, mock.Anything)
-		coord.On("Protection").Return(protection.Config{SignatureValidationKey: nil})
-
-		h := NewMigrate(log, mockAgentInfo, coord)
-		h.tamperProtectionFn = func() bool { return false }
-
-		require.NoError(t, h.Handle(t.Context(), action, ack))
-		coord.AssertNumberOfCalls(t, "Migrate", 1)
-
-		// ack not delegated to migrate coordinator, failure is reported
-		ack.AssertNumberOfCalls(t, "Ack", 0)
-		ack.AssertNumberOfCalls(t, "Commit", 0)
-		coord.AssertNumberOfCalls(t, "ReExec", 1)
-	})
 }
 
 type fakeMigrateCoordinator struct {
 	mock.Mock
 }
 
-func (f *fakeMigrateCoordinator) Migrate(ctx context.Context, a *fleetapi.ActionMigrate, _ func(done <-chan struct{}) backoff.Backoff) error {
-	args := f.Called(ctx, a)
+func (f *fakeMigrateCoordinator) Migrate(ctx context.Context, a *fleetapi.ActionMigrate, b func(done <-chan struct{}) backoff.Backoff, n func(context.Context, *fleetapi.ActionMigrate) error) error {
+	args := f.Called(ctx, a, b, n)
 	return args.Error(0)
 }
 

internal/pkg/agent/application/coordinator/coordinator.go

@@ -10,6 +10,7 @@ import (
 	"fmt"
 	"reflect"
 	"strings"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -364,7 +365,7 @@ type Coordinator struct {
 	// migrationProgressWg is used to block processing of incoming policies after enroll is done
 	// incomming policies are blocked until we reboot so components receiving proxied MIGRATE action
 	// are not confused
-	isMigrationProgress bool
+	migrationProgressWg sync.WaitGroup
 }
 
 // The channels Coordinator reads to receive updates from the various managers.
@@ -600,7 +601,12 @@ func (c *Coordinator) ReExec(callback reexec.ShutdownCallbackFn, argOverrides ..
 
 // Migrate migrates agent to a new cluster and ACKs success to the old one.
 // In case of failure no ack is performed and error is returned.
-func (c *Coordinator) Migrate(ctx context.Context, action *fleetapi.ActionMigrate, backoffFactory func(done <-chan struct{}) backoff.Backoff) error {
+func (c *Coordinator) Migrate(
+	ctx context.Context,
+	action *fleetapi.ActionMigrate,
+	backoffFactory func(done <-chan struct{}) backoff.Backoff,
+	notifyFn func(context.Context, *fleetapi.ActionMigrate) error,
+) error {
 	if !c.isManaged {
 		return ErrNotManaged
 	}
@@ -670,6 +676,24 @@ func (c *Coordinator) Migrate(ctx context.Context, action *fleetapi.ActionMigrat
 		return errors.Join(fmt.Errorf("failed to enroll: %w", err), restoreErr)
 	}
 
+	// lock processing of new config before notifying components
+	// hold lock until notification failure or reexec
+	c.migrationProgressWg.Add(1)
+	if notifyFn != nil {
+		// notify before completing migration
+		// components such endpoint are crucial to work even though it's on stale cluster
+		// error on component side is returned as part of Action response
+		if err := notifyFn(ctx, action); err != nil {
+			restoreErr := RestoreConfig()
+
+			// in case of failure no need to lock processing
+			// safe to forward policy from source cluster
+			c.migrationProgressWg.Done()
+
+			return errors.Join(fmt.Errorf("failed to notify components: %w", err), restoreErr)
+		}
+	}
+
 	// ACK success to source fleet server
 	if err := c.ackMigration(ctx, action, c.fleetAcker); err != nil {
 		c.logger.Warnf("failed to ACK success: %v", err)
@@ -682,7 +706,6 @@ func (c *Coordinator) Migrate(ctx context.Context, action *fleetapi.ActionMigrat
 	}
 
 	c.bestEffortUnenroll(ctx, originalOptions)
-	c.isMigrationProgress = true
 
 	return nil
 }
@@ -1535,9 +1558,7 @@ func (c *Coordinator) runLoopIteration(ctx context.Context) {
 
 // Always called on the main Coordinator goroutine.
 func (c *Coordinator) processConfig(ctx context.Context, cfg *config.Config) (err error) {
-	if c.isMigrationProgress {
-		return nil
-	}
+	c.migrationProgressWg.Wait()
 
 	if c.otelMgr != nil {
 		c.otelCfg = cfg.OTel