From 2157e9f037a795930d8dded8985bcc80a08dc8ed Mon Sep 17 00:00:00 2001
From: Michal Pristas <michal.pristas@gmail.com>
Date: Fri, 3 Sep 2021 16:09:13 +0200
Subject: [PATCH] [Elastic Agent] Await execution finish on windows (#27730)

* Await execution finish on windows

* don't wait if not running as a service
---
 libbeat/service/service_unix.go           |  3 +++
 libbeat/service/service_windows.go        | 26 +++++++++++++++++++----
 x-pack/elastic-agent/pkg/agent/cmd/run.go |  8 +++++--
 3 files changed, 31 insertions(+), 6 deletions(-)

diff --git a/libbeat/service/service_unix.go b/libbeat/service/service_unix.go
index 7d20b04620e6..68851d7e54f3 100644
--- a/libbeat/service/service_unix.go
+++ b/libbeat/service/service_unix.go
@@ -25,3 +25,6 @@ func ProcessWindowsControlEvents(stopCallback func()) {
 
 func notifyWindowsServiceStopped() {
 }
+
+// WaitExecutionDone is not used on non-windows platforms.
+func WaitExecutionDone() {}
diff --git a/libbeat/service/service_windows.go b/libbeat/service/service_windows.go
index a81f4fb5a0ff..027011911c20 100644
--- a/libbeat/service/service_windows.go
+++ b/libbeat/service/service_windows.go
@@ -29,13 +29,15 @@ import (
 )
 
 type beatService struct {
-	stopCallback func()
-	done         chan struct{}
+	stopCallback    func()
+	done            chan struct{}
+	executeFinished chan struct{}
 }
 
 var serviceInstance = &beatService{
-	stopCallback: nil,
-	done:         make(chan struct{}, 0),
+	stopCallback:    nil,
+	done:            make(chan struct{}, 0),
+	executeFinished: make(chan struct{}, 0),
 }
 
 // Execute runs the beat service with the arguments and manages changes that
@@ -85,6 +87,8 @@ const couldNotConnect syscall.Errno = 1063
 // stopCallback function is called when the Stop/Shutdown
 // request is received.
 func ProcessWindowsControlEvents(stopCallback func()) {
+	defer close(serviceInstance.executeFinished)
+
 	isInteractive, err := svc.IsAnInteractiveSession()
 	if err != nil {
 		logp.Err("IsAnInteractiveSession: %v", err)
@@ -125,3 +129,17 @@ func ProcessWindowsControlEvents(stopCallback func()) {
 
 	logp.Err("Windows service setup failed: %+v", err)
 }
+
+// WaitExecutionDone returns only after stop was reported to service manager.
+// If response is not retrieved within 500 millisecond wait is aborted.
+func WaitExecutionDone() {
+	if isWinService, err := svc.IsWindowsService(); err != nil || !isWinService {
+		// not a service, don't wait
+		return
+	}
+
+	select {
+	case <-serviceInstance.executeFinished:
+	case <-time.After(500 * time.Millisecond):
+	}
+}
diff --git a/x-pack/elastic-agent/pkg/agent/cmd/run.go b/x-pack/elastic-agent/pkg/agent/cmd/run.go
index c9f02de59fbb..cf24a932a919 100644
--- a/x-pack/elastic-agent/pkg/agent/cmd/run.go
+++ b/x-pack/elastic-agent/pkg/agent/cmd/run.go
@@ -59,10 +59,14 @@ func newRunCommandWithArgs(_ []string, streams *cli.IOStreams) *cobra.Command {
 	}
 }
 
-func run(streams *cli.IOStreams, override cfgOverrider) error { // Windows: Mark service as stopped.
+func run(streams *cli.IOStreams, override cfgOverrider) error {
+	// Windows: Mark service as stopped.
 	// After this is run, the service is considered by the OS to be stopped.
 	// This must be the first deferred cleanup task (last to execute).
-	defer service.NotifyTermination()
+	defer func() {
+		service.NotifyTermination()
+		service.WaitExecutionDone()
+	}()
 
 	locker := filelock.NewAppLocker(paths.Data(), paths.AgentLockFileName)
 	if err := locker.TryLock(); err != nil {