Fix upgrade for same versions (#7635)

Author: michalpristas

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

@@ -92,16 +92,28 @@ func (h *Upgrade) Handle(ctx context.Context, a fleetapi.Action, ack acker.Acker
 // User is responsible for obtaining and releasing bkgMutex lock
 func (h *Upgrade) ackActions(ctx context.Context, ack acker.Acker) {
 	for _, a := range h.bkgActions {
-		if err := ack.Ack(ctx, a); err != nil {
-			h.log.Errorf("ack of failed upgrade failed: %v", err)
-		}
+		h.ackAction(ctx, ack, a, false)
 	}
 	h.bkgActions = nil
 	if err := ack.Commit(ctx); err != nil {
 		h.log.Errorf("commit of ack for failed upgrade failed: %v", err)
 	}
 }
 
+// ackActions Acks all the actions in bkgActions, and deletes entries from bkgActions.
+// User is responsible for obtaining and releasing bkgMutex lock
+func (h *Upgrade) ackAction(ctx context.Context, ack acker.Acker, action fleetapi.Action, commit bool) {
+	if err := ack.Ack(ctx, action); err != nil {
+		h.log.Errorf("ack of failed upgrade failed: %v", err)
+	}
+
+	if commit {
+		if err := ack.Commit(ctx); err != nil {
+			h.log.Errorf("commit of ack for failed upgrade failed: %v", err)
+		}
+	}
+}
+
 // getAsyncContext returns a cancelContext and whether or not to run the upgrade
 func (h *Upgrade) getAsyncContext(ctx context.Context, action fleetapi.Action, ack acker.Acker) (context.Context, bool) {
 	h.bkgMutex.Lock()
@@ -125,11 +137,22 @@ func (h *Upgrade) getAsyncContext(ctx context.Context, action fleetapi.Action, a
 		h.log.Errorf("invalid type, expected ActionUpgrade and received %T", action)
 		return nil, false
 	}
-	if (upgradeAction.Data.Version == bkgAction.Data.Version) &&
-		(upgradeAction.Data.SourceURI == bkgAction.Data.SourceURI) {
+	if upgradeAction.ActionID == bkgAction.ActionID {
 		h.log.Infof("Duplicate upgrade to version %s received",
 			bkgAction.Data.Version)
-		h.bkgActions = append(h.bkgActions, action)
+		return nil, false
+	}
+
+	if upgradeAction.Data.Version == bkgAction.Data.Version &&
+		upgradeAction.Data.SourceURI == bkgAction.Data.SourceURI {
+		// not the same action this one needs to be acked
+		go func() {
+			// kick it off and don't block, lock to prevent race with ackActions from finished upgrade
+			h.bkgMutex.Lock()
+			defer h.bkgMutex.Unlock()
+
+			h.ackAction(ctx, ack, upgradeAction, true)
+		}()
 		return nil, false
 	}
 

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

@@ -11,6 +11,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/stretchr/testify/mock"
 	"github.com/stretchr/testify/require"
 
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/coordinator"
@@ -68,7 +69,11 @@ func (u *mockUpgradeManager) Upgrade(
 		pgpBytes...)
 }
 
-func (u *mockUpgradeManager) Ack(ctx context.Context, acker acker.Acker) error {
+func (u *mockUpgradeManager) Ack(_ context.Context, _ acker.Acker) error {
+	return nil
+}
+
+func (u *mockUpgradeManager) AckAction(_ context.Context, _ acker.Acker, _ fleetapi.Action) error {
 	return nil
 }
 
@@ -110,7 +115,7 @@ func TestUpgradeHandler(t *testing.T) {
 				return nil, nil
 			},
 		},
-		nil, nil, nil, nil, nil, false, nil)
+		nil, nil, nil, nil, nil, false, nil, nil)
 	//nolint:errcheck // We don't need the termination state of the Coordinator
 	go c.Run(ctx)
 
@@ -169,7 +174,7 @@ func TestUpgradeHandlerSameVersion(t *testing.T) {
 				return nil, err
 			},
 		},
-		nil, nil, nil, nil, nil, false, nil)
+		nil, nil, nil, nil, nil, false, nil, nil)
 	//nolint:errcheck // We don't need the termination state of the Coordinator
 	go c.Run(ctx)
 
@@ -190,6 +195,98 @@ func TestUpgradeHandlerSameVersion(t *testing.T) {
 	}
 }
 
+func TestDuplicateActionsHandled(t *testing.T) {
+	// Create a cancellable context that will shut down the coordinator after
+	// the test.
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	log, _ := logger.New("", false)
+	upgradeCalledChan := make(chan string)
+
+	agentInfo := &info.AgentInfo{}
+	acker := &fakeAcker{}
+
+	// Create and start the Coordinator
+	c := coordinator.New(
+		log,
+		configuration.DefaultConfiguration(),
+		logger.DefaultLogLevel,
+		agentInfo,
+		component.RuntimeSpecs{},
+		nil,
+		&mockUpgradeManager{
+			UpgradeFn: func(
+				ctx context.Context,
+				version string,
+				sourceURI string,
+				action *fleetapi.ActionUpgrade,
+				details *details.Details,
+				skipVerifyOverride bool,
+				skipDefaultPgp bool,
+				pgpBytes ...string) (reexec.ShutdownCallbackFn, error) {
+
+				defer func() {
+					upgradeCalledChan <- action.ActionID
+				}()
+
+				return nil, nil
+			},
+		},
+		nil, nil, nil, nil, nil, false, nil, acker)
+	//nolint:errcheck // We don't need the termination state of the Coordinator
+	go c.Run(ctx)
+
+	u := NewUpgrade(log, c)
+	a1 := fleetapi.ActionUpgrade{
+		ActionID: "action-8.5-1",
+		Data: fleetapi.ActionUpgradeData{
+			Version: "8.5.0", SourceURI: "http://localhost",
+		},
+	}
+	a2 := fleetapi.ActionUpgrade{
+		ActionID: "action-8.5-2",
+		Data: fleetapi.ActionUpgradeData{
+			Version: "8.5.0", SourceURI: "http://localhost",
+		},
+	}
+
+	checkMsg := func(c <-chan string, expected, errMsg string) error {
+		t.Helper()
+		// Make sure this test does not dead lock or wait for too long
+		// For some reason < 1s sometimes makes the test fail.
+		select {
+		case <-time.Tick(1500 * time.Millisecond):
+			return errors.New("timed out waiting for Upgrade to return")
+		case msg := <-c:
+			require.Equal(t, expected, msg, errMsg)
+		}
+
+		return nil
+	}
+
+	acker.On("Ack", mock.Anything, mock.Anything).Return(nil)
+	acker.On("Commit", mock.Anything).Return(nil)
+
+	t.Log("First upgrade action should be processed")
+	require.NoError(t, u.Handle(ctx, &a1, acker))
+	require.Nil(t, checkMsg(upgradeCalledChan, a1.ActionID, "action was not processed"))
+	c.ClearOverrideState() // it's upgrading, normally we would restart
+
+	t.Log("Action with different ID but same version should not be propagated to upgrader but acked")
+	require.NoError(t, u.Handle(ctx, &a2, acker))
+	require.NotNil(t, checkMsg(upgradeCalledChan, a2.ActionID, "action was not processed"))
+	acker.AssertCalled(t, "Ack", ctx, &a2)
+	acker.AssertCalled(t, "Commit", ctx)
+
+	c.ClearOverrideState() // it's upgrading, normally we would restart
+
+	t.Log("Resending action with same ID should be skipped")
+	require.NoError(t, u.Handle(ctx, &a1, acker))
+	require.NotNil(t, checkMsg(upgradeCalledChan, a1.ActionID, "action was not processed"))
+	acker.AssertNotCalled(t, "Ack", ctx, &a1)
+}
+
 func TestUpgradeHandlerNewVersion(t *testing.T) {
 	// Create a cancellable context that will shut down the coordinator after
 	// the test.
@@ -230,15 +327,23 @@ func TestUpgradeHandlerNewVersion(t *testing.T) {
 				return nil, nil
 			},
 		},
-		nil, nil, nil, nil, nil, false, nil)
+		nil, nil, nil, nil, nil, false, nil, nil)
 	//nolint:errcheck // We don't need the termination state of the Coordinator
 	go c.Run(ctx)
 
 	u := NewUpgrade(log, c)
-	a1 := fleetapi.ActionUpgrade{Data: fleetapi.ActionUpgradeData{
-		Version: "8.2.0", SourceURI: "http://localhost"}}
-	a2 := fleetapi.ActionUpgrade{Data: fleetapi.ActionUpgradeData{
-		Version: "8.5.0", SourceURI: "http://localhost"}}
+	a1 := fleetapi.ActionUpgrade{
+		ActionID: "action-8.2",
+		Data: fleetapi.ActionUpgradeData{
+			Version: "8.2.0", SourceURI: "http://localhost",
+		},
+	}
+	a2 := fleetapi.ActionUpgrade{
+		ActionID: "action-8.5",
+		Data: fleetapi.ActionUpgradeData{
+			Version: "8.5.0", SourceURI: "http://localhost",
+		},
+	}
 	ack := noopacker.New()
 
 	checkMsg := func(c <-chan string, expected, errMsg string) {
@@ -262,3 +367,17 @@ func TestUpgradeHandlerNewVersion(t *testing.T) {
 	require.NoError(t, err2)
 	checkMsg(upgradeCalledChan, "8.5.0", "second call to Upgrade must be with version 8.5.0")
 }
+
+type fakeAcker struct {
+	mock.Mock
+}
+
+func (f *fakeAcker) Ack(ctx context.Context, action fleetapi.Action) error {
+	args := f.Called(ctx, action)
+	return args.Error(0)
+}
+
+func (f *fakeAcker) Commit(ctx context.Context) error {
+	args := f.Called(ctx)
+	return args.Error(0)
+}

internal/pkg/agent/application/application.go

@@ -21,10 +21,16 @@ import (
 	"github.com/elastic/elastic-agent/internal/pkg/agent/configuration"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/errors"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/storage"
+	stateStore "github.com/elastic/elastic-agent/internal/pkg/agent/storage/store"
 	"github.com/elastic/elastic-agent/internal/pkg/capabilities"
 	"github.com/elastic/elastic-agent/internal/pkg/composable"
 	"github.com/elastic/elastic-agent/internal/pkg/composable/providers/kubernetes"
 	"github.com/elastic/elastic-agent/internal/pkg/config"
+	"github.com/elastic/elastic-agent/internal/pkg/fleetapi/acker"
+	"github.com/elastic/elastic-agent/internal/pkg/fleetapi/acker/fleet"
+	"github.com/elastic/elastic-agent/internal/pkg/fleetapi/acker/lazy"
+	"github.com/elastic/elastic-agent/internal/pkg/fleetapi/acker/retrier"
+	fleetclient "github.com/elastic/elastic-agent/internal/pkg/fleetapi/client"
 	otelmanager "github.com/elastic/elastic-agent/internal/pkg/otel/manager"
 	"github.com/elastic/elastic-agent/internal/pkg/release"
 	"github.com/elastic/elastic-agent/pkg/component"
@@ -136,6 +142,7 @@ func New(
 	var composableManaged bool
 	var isManaged bool
 
+	var actionAcker acker.Acker
 	if testingMode {
 		log.Info("Elastic Agent has been started in testing mode and is managed through the control protocol")
 
@@ -182,8 +189,29 @@ func New(
 				InjectProxyEndpointModifier(),
 			)
 
+			client, err := fleetclient.NewAuthWithConfig(log, cfg.Fleet.AccessAPIKey, cfg.Fleet.Client)
+			if err != nil {
+				return nil, nil, nil, errors.New(err,
+					"fail to create API client",
+					errors.TypeNetwork,
+					errors.M(errors.MetaKeyURI, cfg.Fleet.Client.Host))
+			}
+			stateStorage, err := stateStore.NewStateStoreWithMigration(ctx, log, paths.AgentActionStoreFile(), paths.AgentStateStoreFile())
+			if err != nil {
+				return nil, nil, nil, errors.New(err, fmt.Sprintf("fail to read state store '%s'", paths.AgentStateStoreFile()))
+			}
+
+			fleetAcker, err := fleet.NewAcker(log, agentInfo, client)
+			if err != nil {
+				return nil, nil, nil, fmt.Errorf("failed to create acker: %w", err)
+			}
+
+			retrier := retrier.New(fleetAcker, log)
+			batchedAcker := lazy.NewAcker(fleetAcker, log, lazy.WithRetrier(retrier))
+			actionAcker = stateStore.NewStateStoreActionAcker(batchedAcker, stateStorage)
+
 			// TODO: stop using global state
-			managed, err = newManagedConfigManager(ctx, log, agentInfo, cfg, store, runtime, fleetInitTimeout, paths.Top(), upgrader)
+			managed, err = newManagedConfigManager(ctx, log, agentInfo, cfg, store, runtime, fleetInitTimeout, paths.Top(), client, fleetAcker, actionAcker, retrier, stateStorage, upgrader)
 			if err != nil {
 				return nil, nil, nil, err
 			}
@@ -197,7 +225,7 @@ func New(
 	}
 
 	otelManager := otelmanager.NewOTelManager(log.Named("otel_manager"))
-	coord := coordinator.New(log, cfg, logLevel, agentInfo, specs, reexec, upgrader, runtime, configMgr, varsManager, caps, monitor, isManaged, otelManager, compModifiers...)
+	coord := coordinator.New(log, cfg, logLevel, agentInfo, specs, reexec, upgrader, runtime, configMgr, varsManager, caps, monitor, isManaged, otelManager, actionAcker, compModifiers...)
 	if managed != nil {
 		// the coordinator requires the config manager as well as in managed-mode the config manager requires the
 		// coordinator, so it must be set here once the coordinator is created

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

@@ -75,6 +75,9 @@ type UpgradeManager interface {
 	// Ack is used on startup to check if the agent has upgraded and needs to send an ack for the action
 	Ack(ctx context.Context, acker acker.Acker) error
 
+	// AckAction is used to ack not persisted action.
+	AckAction(ctx context.Context, acker acker.Acker, action fleetapi.Action) error
+
 	// MarkerWatcher returns a watcher for the upgrade marker.
 	MarkerWatcher() upgrade.MarkerWatcher
 }
@@ -204,8 +207,9 @@ type Coordinator struct {
 	agentInfo info.Agent
 	isManaged bool
 
-	cfg   *configuration.Configuration
-	specs component.RuntimeSpecs
+	cfg        *configuration.Configuration
+	specs      component.RuntimeSpecs
+	fleetAcker acker.Acker
 
 	reexecMgr  ReExecManager
 	upgradeMgr UpgradeManager
@@ -372,7 +376,7 @@ type UpdateComponentChange struct {
 }
 
 // New creates a new coordinator.
-func New(logger *logger.Logger, cfg *configuration.Configuration, logLevel logp.Level, agentInfo info.Agent, specs component.RuntimeSpecs, reexecMgr ReExecManager, upgradeMgr UpgradeManager, runtimeMgr RuntimeManager, configMgr ConfigManager, varsMgr VarsManager, caps capabilities.Capabilities, monitorMgr MonitorManager, isManaged bool, otelMgr OTelManager, modifiers ...ComponentsModifier) *Coordinator {
+func New(logger *logger.Logger, cfg *configuration.Configuration, logLevel logp.Level, agentInfo info.Agent, specs component.RuntimeSpecs, reexecMgr ReExecManager, upgradeMgr UpgradeManager, runtimeMgr RuntimeManager, configMgr ConfigManager, varsMgr VarsManager, caps capabilities.Capabilities, monitorMgr MonitorManager, isManaged bool, otelMgr OTelManager, fleetAcker acker.Acker, modifiers ...ComponentsModifier) *Coordinator {
 	var fleetState cproto.State
 	var fleetMessage string
 	if !isManaged {
@@ -426,6 +430,8 @@ func New(logger *logger.Logger, cfg *configuration.Configuration, logLevel logp.
 		heartbeatChan:              make(chan struct{}),
 		componentPIDTicker:         time.NewTicker(time.Second * 30),
 		componentPidRequiresUpdate: &atomic.Bool{},
+
+		fleetAcker: fleetAcker,
 	}
 	// Setup communication channels for any non-nil components. This pattern
 	// lets us transparently accept nil managers / simulated events during
@@ -585,7 +591,7 @@ func (c *Coordinator) Upgrade(ctx context.Context, version string, sourceURI str
 		if errors.Is(err, upgrade.ErrUpgradeSameVersion) {
 			// Set upgrade state to completed so update no longer shows in-progress.
 			det.SetState(details.StateCompleted)
-			return nil
+			return c.upgradeMgr.AckAction(ctx, c.fleetAcker, action)
 		}
 		det.Fail(err)
 		return err