fix: wait server server function termination

server/serverimpl.go

@@ -19,17 +19,17 @@ import (
 	serverTypes "github.com/open-telemetry/opamp-go/server/types"
 )
 
-var (
-	errAlreadyStarted = errors.New("already started")
+var errAlreadyStarted = errors.New("already started")
+
+const (
+	defaultOpAMPPath      = "/v1/opamp"
+	headerContentType     = "Content-Type"
+	headerContentEncoding = "Content-Encoding"
+	headerAcceptEncoding  = "Accept-Encoding"
+	contentEncodingGzip   = "gzip"
+	contentTypeProtobuf   = "application/x-protobuf"
 )
 
-const defaultOpAMPPath = "/v1/opamp"
-const headerContentType = "Content-Type"
-const headerContentEncoding = "Content-Encoding"
-const headerAcceptEncoding = "Accept-Encoding"
-const contentEncodingGzip = "gzip"
-const contentTypeProtobuf = "application/x-protobuf"
-
 type server struct {
 	logger   types.Logger
 	settings Settings
@@ -39,7 +39,8 @@ type server struct {
 
 	// The listening HTTP Server after successful Start() call. Nil if Start()
 	// is not called or was not successful.
-	httpServer *http.Server
+	httpServer        *http.Server
+	httpServerServeWg *sync.WaitGroup
 
 	// The network address Server is listening on. Nil if not started.
 	addr net.Addr
@@ -108,6 +109,9 @@ func (s *server) Start(settings StartSettings) error {
 		ConnContext: contextWithConn,
 	}
 	s.httpServer = hs
+	httpServerServeWg := sync.WaitGroup{}
+	httpServerServeWg.Add(1)
+	s.httpServerServeWg = &httpServerServeWg
 
 	listenAddr := s.httpServer.Addr
 
@@ -118,15 +122,21 @@ func (s *server) Start(settings StartSettings) error {
 		}
 		err = s.startHttpServer(
 			listenAddr,
-			func(l net.Listener) error { return hs.ServeTLS(l, "", "") },
+			func(l net.Listener) error {
+				defer httpServerServeWg.Done()
+				return hs.ServeTLS(l, "", "")
+			},
 		)
 	} else {
 		if listenAddr == "" {
 			listenAddr = ":http"
 		}
 		err = s.startHttpServer(
 			listenAddr,
-			func(l net.Listener) error { return hs.Serve(l) },
+			func(l net.Listener) error {
+				defer httpServerServeWg.Done()
+				return hs.Serve(l)
+			},
 		)
 	}
 	return err
@@ -159,7 +169,11 @@ func (s *server) Stop(ctx context.Context) error {
 		defer func() { s.httpServer = nil }()
 		// This stops accepting new connections. TODO: close existing
 		// connections and wait them to be terminated.
-		return s.httpServer.Shutdown(ctx)
+		err := s.httpServer.Shutdown(ctx)
+		if err != nil {
+			return err
+		}
+		s.httpServerServeWg.Wait()
 	}
 	return nil
 }
@@ -366,7 +380,6 @@ func (s *server) handlePlainHTTPRequest(req *http.Request, w http.ResponseWriter
 		w.Header().Set(headerContentEncoding, contentEncodingGzip)
 	}
 	_, err = w.Write(bodyBytes)
-
 	if err != nil {
 		s.logger.Debugf(req.Context(), "Cannot send HTTP response: %v", err)
 	}