From cb83cc19cedd5ba216f952220ac596b43a648186 Mon Sep 17 00:00:00 2001
From: kruskall <99559985+kruskall@users.noreply.github.com>
Date: Wed, 19 Jun 2024 19:42:24 +0200
Subject: [PATCH] refactor: replace custom atomic types with stdlib atomic

Go 1.19 added new atomic types. Replace custom atomic implementation with native types.
---
 .../actions/handlers/handler_action_unenroll_test.go     | 9 ++++-----
 .../upgrade/artifact/download/http/progress_reporter.go  | 7 +++----
 pkg/component/runtime/manager.go                         | 4 ++--
 pkg/component/runtime/runtime.go                         | 5 ++---
 4 files changed, 11 insertions(+), 14 deletions(-)

diff --git a/internal/pkg/agent/application/actions/handlers/handler_action_unenroll_test.go b/internal/pkg/agent/application/actions/handlers/handler_action_unenroll_test.go
index e28612e9502..79ab501d829 100644
--- a/internal/pkg/agent/application/actions/handlers/handler_action_unenroll_test.go
+++ b/internal/pkg/agent/application/actions/handlers/handler_action_unenroll_test.go
@@ -6,10 +6,9 @@ package handlers
 
 import (
 	"context"
+	"sync/atomic"
 	"testing"
 
-	"github.com/elastic/elastic-agent-libs/atomic"
-
 	"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"
@@ -43,7 +42,7 @@ func makeComponentState(name string, proxiedActions []string) runtime.ComponentC
 
 type MockActionCoordinator struct {
 	st               coordinator.State
-	performedActions atomic.Int
+	performedActions atomic.Int64
 }
 
 func (c *MockActionCoordinator) State() coordinator.State {
@@ -51,7 +50,7 @@ func (c *MockActionCoordinator) State() coordinator.State {
 }
 
 func (c *MockActionCoordinator) PerformAction(ctx context.Context, comp component.Component, unit component.Unit, name string, params map[string]interface{}) (map[string]interface{}, error) {
-	c.performedActions.Inc()
+	c.performedActions.Add(1)
 	return nil, nil
 }
 
@@ -118,7 +117,7 @@ func TestActionUnenrollHandler(t *testing.T) {
 		name                 string
 		st                   coordinator.State
 		wantErr              error // Handler error
-		wantPerformedActions int
+		wantPerformedActions int64
 		tamperProtectionFn   func() bool
 	}{
 		{
diff --git a/internal/pkg/agent/application/upgrade/artifact/download/http/progress_reporter.go b/internal/pkg/agent/application/upgrade/artifact/download/http/progress_reporter.go
index 491646b3ab5..c94292c0752 100644
--- a/internal/pkg/agent/application/upgrade/artifact/download/http/progress_reporter.go
+++ b/internal/pkg/agent/application/upgrade/artifact/download/http/progress_reporter.go
@@ -6,9 +6,8 @@ package http
 
 import (
 	"context"
+	"sync/atomic"
 	"time"
-
-	"github.com/elastic/elastic-agent-libs/atomic"
 )
 
 type downloadProgressReporter struct {
@@ -17,7 +16,7 @@ type downloadProgressReporter struct {
 	warnTimeout time.Duration
 	length      float64
 
-	downloaded atomic.Int
+	downloaded atomic.Int64
 	started    time.Time
 
 	progressObservers []progressObserver
@@ -42,7 +41,7 @@ func newDownloadProgressReporter(sourceURI string, timeout time.Duration, length
 
 func (dp *downloadProgressReporter) Write(b []byte) (int, error) {
 	n := len(b)
-	dp.downloaded.Add(n)
+	dp.downloaded.Add(int64(n))
 	return n, nil
 }
 
diff --git a/pkg/component/runtime/manager.go b/pkg/component/runtime/manager.go
index 17e86c38d68..b382dfac39a 100644
--- a/pkg/component/runtime/manager.go
+++ b/pkg/component/runtime/manager.go
@@ -16,6 +16,7 @@ import (
 	"path"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	"github.com/gofrs/uuid"
@@ -28,7 +29,6 @@ import (
 
 	"github.com/elastic/elastic-agent-client/v7/pkg/client"
 	"github.com/elastic/elastic-agent-client/v7/pkg/proto"
-	"github.com/elastic/elastic-agent-libs/atomic"
 
 	"github.com/elastic/elastic-agent/internal/pkg/agent/application/info"
 	"github.com/elastic/elastic-agent/internal/pkg/agent/configuration"
@@ -193,7 +193,7 @@ func NewManager(
 		errCh:          make(chan error),
 		monitor:        monitor,
 		grpcConfig:     grpcConfig,
-		serverReady:    atomic.NewBool(false),
+		serverReady:    &atomic.Bool{},
 		doneChan:       make(chan struct{}),
 	}
 	return m, nil
diff --git a/pkg/component/runtime/runtime.go b/pkg/component/runtime/runtime.go
index 68baded754b..48d499960b8 100644
--- a/pkg/component/runtime/runtime.go
+++ b/pkg/component/runtime/runtime.go
@@ -8,10 +8,10 @@ import (
 	"context"
 	"errors"
 	"sync"
+	"sync/atomic"
 
 	"github.com/elastic/elastic-agent-client/v7/pkg/client"
 	"github.com/elastic/elastic-agent-client/v7/pkg/proto"
-	"github.com/elastic/elastic-agent-libs/atomic"
 	"github.com/elastic/elastic-agent/internal/pkg/runner"
 	"github.com/elastic/elastic-agent/pkg/component"
 	"github.com/elastic/elastic-agent/pkg/core/logger"
@@ -188,11 +188,10 @@ func (s *componentRuntimeState) start() error {
 }
 
 func (s *componentRuntimeState) stop(teardown bool, signed *component.Signed) error {
-	if s.shuttingDown.Load() {
+	if !s.shuttingDown.CompareAndSwap(false, true) {
 		// already stopping
 		return nil
 	}
-	s.shuttingDown.Store(true)
 	if teardown {
 		return s.runtime.Teardown(signed)
 	}