From 453e76c275cc42a39decbc43b3b51b38dae576ce Mon Sep 17 00:00:00 2001
From: Andrew Wilkins <axw@elastic.co>
Date: Mon, 15 Feb 2021 13:00:28 +0800
Subject: [PATCH] beater: don't hang if runServer returns error

If runServer returns an error without Stop being
called, signal the "done" channel.
---
 beater/beater.go | 12 +++++++++++-
 1 file changed, 11 insertions(+), 1 deletion(-)

diff --git a/beater/beater.go b/beater/beater.go
index 35de9b57036..b5616384cdf 100644
--- a/beater/beater.go
+++ b/beater/beater.go
@@ -198,9 +198,12 @@ func (bt *beater) start(ctx context.Context, cancelContext context.CancelFunc, b
 		if err != nil {
 			return nil, err
 		}
-		bt.stopServer = func() {
+		go func() {
 			defer close(done)
 			defer closeTracer()
+			s.Wait()
+		}()
+		bt.stopServer = func() {
 			if bt.config.ShutdownTimeout > 0 {
 				time.AfterFunc(bt.config.ShutdownTimeout, cancelContext)
 			}
@@ -316,11 +319,18 @@ func (s *serverRunner) String() string {
 	return "APMServer"
 }
 
+// Stop stops the server.
 func (s *serverRunner) Stop() {
 	s.stopOnce.Do(s.cancelRunServerContext)
+	s.Wait()
+}
+
+// Wait waits for the server to stop.
+func (s *serverRunner) Wait() {
 	s.wg.Wait()
 }
 
+// Start starts the server.
 func (s *serverRunner) Start() {
 	s.wg.Add(1)
 	go func() {