From 11cb50d9fca504649b5817dea08096c0edbdd4ec Mon Sep 17 00:00:00 2001
From: Shivam Verma <113686362+vermaaatul07@users.noreply.github.com>
Date: Thu, 13 Jun 2024 01:28:58 +0530
Subject: [PATCH] [query] Avoid errors when closing shared listener (#5559)

## Description
This PR aims to address the intermittent timeout issue in the
`TestServerSinglePort` test by adding enhanced logging to the `Start`
and `Close` methods of the server. These additional logs provide better
visibility into the server's behavior during the test, which should help
in diagnosing the root cause of the issue.

## Changes
### Enhanced Logging:
- Added error logging in the `Start` method for listener initialization
failures.
- Added info logging in the `Start` method to log the initialized HTTP
and gRPC ports.
- Added info logging in the `Close` method to indicate an attempt to
close the server and confirm when the server has stopped.

These changes provide better diagnostics and insights into server
operations, which should help in identifying and resolving the flaky
test issue if it occurs again.

## Related Issue
Closes #5519

## Checklist
- [x] I have read `CONTRIBUTING_GUIDELINES`
- [x] I have signed all commits
- [ ] I have added unit tests for the new functionality
- [x] I have run lint and test steps successfully
  - for `jaeger`: `make lint test`
  - for `jaeger-ui`: `yarn lint` and `yarn test`

---------

Signed-off-by: Shivam Verma <vermaaatul07@gmail.com>
Signed-off-by: Yuri Shkuro <github@ysh.us>
Co-authored-by: Yuri Shkuro <yurishkuro@users.noreply.github.com>
Co-authored-by: Yuri Shkuro <github@ysh.us>
---
 cmd/query/app/server.go      | 27 ++++++++++++++++++++-------
 cmd/query/app/server_test.go |  2 +-
 2 files changed, 21 insertions(+), 8 deletions(-)

diff --git a/cmd/query/app/server.go b/cmd/query/app/server.go
index 8e1af7662c3..c2debad540a 100644
--- a/cmd/query/app/server.go
+++ b/cmd/query/app/server.go
@@ -274,7 +274,7 @@ func (s *Server) initListener() (cmux.CMux, error) {
 func (s *Server) Start() error {
 	cmuxServer, err := s.initListener()
 	if err != nil {
-		return err
+		return fmt.Errorf("query server failed to initialize listener: %w", err)
 	}
 	s.cmuxServer = cmuxServer
 
@@ -317,7 +317,8 @@ func (s *Server) Start() error {
 	go func() {
 		s.logger.Info("Starting GRPC server", zap.Int("port", grpcPort), zap.String("addr", s.queryOptions.GRPCHostPort))
 
-		if err := s.grpcServer.Serve(s.grpcConn); err != nil {
+		err := s.grpcServer.Serve(s.grpcConn)
+		if err != nil && !errors.Is(err, cmux.ErrListenerClosed) && !errors.Is(err, cmux.ErrServerClosed) {
 			s.logger.Error("Could not start GRPC server", zap.Error(err))
 		}
 		s.logger.Info("GRPC server stopped", zap.Int("port", grpcPort), zap.String("addr", s.queryOptions.GRPCHostPort))
@@ -336,6 +337,7 @@ func (s *Server) Start() error {
 			if err != nil && !strings.Contains(err.Error(), "use of closed network connection") {
 				s.logger.Error("Could not start multiplexed server", zap.Error(err))
 			}
+			s.logger.Info("CMUX server stopped", zap.Int("port", tcpPort), zap.String("addr", s.queryOptions.HTTPHostPort))
 			s.healthCheck.Set(healthcheck.Unavailable)
 			s.bgFinished.Done()
 		}()
@@ -344,16 +346,27 @@ func (s *Server) Start() error {
 	return nil
 }
 
-// Close stops http, GRPC servers and closes the port listener.
+// Close stops HTTP, GRPC servers and closes the port listener.
 func (s *Server) Close() error {
-	var errs []error
-	errs = append(errs, s.queryOptions.TLSGRPC.Close())
-	errs = append(errs, s.queryOptions.TLSHTTP.Close())
+	errs := []error{
+		s.queryOptions.TLSGRPC.Close(),
+		s.queryOptions.TLSHTTP.Close(),
+	}
+
+	s.logger.Info("Closing HTTP server")
+	if err := s.httpServer.Close(); err != nil {
+		errs = append(errs, fmt.Errorf("failed to close HTTP server: %w", err))
+	}
+
+	s.logger.Info("Stopping gRPC server")
 	s.grpcServer.Stop()
-	errs = append(errs, s.httpServer.Close())
+
 	if !s.separatePorts {
+		s.logger.Info("Closing CMux server")
 		s.cmuxServer.Close()
 	}
+
 	s.bgFinished.Wait()
+	s.logger.Info("Server stopped")
 	return errors.Join(errs...)
 }
diff --git a/cmd/query/app/server_test.go b/cmd/query/app/server_test.go
index 83cc081c026..b556aecd87f 100644
--- a/cmd/query/app/server_test.go
+++ b/cmd/query/app/server_test.go
@@ -601,7 +601,7 @@ func TestServerInUseHostPort(t *testing.T) {
 
 func TestServerSinglePort(t *testing.T) {
 	flagsSvc := flags.NewService(ports.QueryAdminHTTP)
-	flagsSvc.Logger = zaptest.NewLogger(t)
+	flagsSvc.Logger = zaptest.NewLogger(t, zaptest.WrapOptions(zap.AddCaller()))
 	hostPort := ports.GetAddressFromCLIOptions(ports.QueryHTTP, "")
 	querySvc := makeQuerySvc()
 	server, err := NewServer(flagsSvc.Logger, flagsSvc.HC(), querySvc.qs, nil,